fixing benchmark tests (#1244)

Author: parmeet

benchmark/benchmark_basic_english_normalize.py

@@ -18,17 +18,17 @@ def _run_benchmark_lookup(train, tokenizer):
     experimental_jit_basic_english_normalize = torch.jit.script(experimental_basic_english_normalize)
 
     # existing eager lookup
-    train, _ = AG_NEWS()
+    train = AG_NEWS(split='train')
     print("BasicEnglishNormalize - Eager Mode")
     _run_benchmark_lookup(train, existing_basic_english_tokenizer)
 
     # experimental eager lookup
-    train, _ = AG_NEWS()
+    train = AG_NEWS(split='train')
     print("BasicEnglishNormalize Experimental - Eager Mode")
     _run_benchmark_lookup(train, experimental_basic_english_normalize)
 
     # experimental jit lookup
-    train, _ = AG_NEWS()
+    train = AG_NEWS(split='train')
     print("BasicEnglishNormalize Experimental - Jit Mode")
     _run_benchmark_lookup(train, experimental_jit_basic_english_normalize)
 

benchmark/benchmark_experimental_vectors.py

@@ -13,7 +13,7 @@ def _run_benchmark_lookup(tokens, vector):
             vector[token]
         print("Lookup time:", time.monotonic() - t0)
 
-    train, = AG_NEWS(data_select='train')
+    train = AG_NEWS(split='train')
     vocab = train.get_vocab()
     tokens = []
     for (label, text) in train:

benchmark/benchmark_experimental_vocab.py

@@ -94,7 +94,7 @@ def _run_benchmark_lookup(tokens, vocab):
     tokens = []
     tokens_lists = []
 
-    train, = AG_NEWS(data_select='train')
+    train = AG_NEWS(split='train')
     vocab = train.get_vocab()
     for (_, text) in train:
         cur_tokens = []

benchmark/benchmark_sentencepiece.py

@@ -3,7 +3,7 @@
 from torchtext.experimental.transforms import load_sp_model as load_pybind_sp_model
 from torchtext.data.functional import load_sp_model as load_torchbind_sp_model
 from torchtext.utils import download_from_url
-from torchtext.datasets import text_classification as raw
+from torchtext.datasets import DATASETS
 
 
 def benchmark_sentencepiece(args):
@@ -17,13 +17,13 @@ def _run_benchmark(train, spm_processor):
     sp_model_path = download_from_url('https://pytorch.s3.amazonaws.com/models/text/pretrained_spm/text_unigram_15000.model')
 
     # existing sentencepiece model with torchbind
-    train, _ = raw.DATASETS[args.dataset]()
+    train = DATASETS[args.dataset](split='train')
     sp_model = load_torchbind_sp_model(sp_model_path)
     print("SentencePiece EncodeAsIds - torchbind")
     _run_benchmark(train, sp_model.EncodeAsIds)
 
     # experimental sentencepiece model with pybind
-    train, _ = raw.DATASETS[args.dataset]()
+    train = DATASETS[args.dataset](split='train')
     sp_model = load_pybind_sp_model(sp_model_path)
     print("SentencePiece EncodeAsIds - pybind")
     _run_benchmark(train, sp_model.EncodeAsIds)