Use Vocab for supervised learning datasets

examples/text_classification/model.py

@@ -15,5 +15,4 @@ def init_weights(self):
         self.fc.bias.data.zero_()
 
     def forward(self, text, offsets):
-        embedded = self.embedding(text, offsets)
-        return self.fc(embedded)
+        return self.fc(self.embedding(text, offsets))

examples/text_classification/train.py

@@ -1,42 +1,38 @@
 import os
 import logging
-import random
 import argparse
 
 import torch
 
-from torchtext.datasets.text_classification import AG_NEWS
+from torchtext.datasets import text_classification
+from torch.utils.data import DataLoader
 
 from model import TextSentiment
 
 
-def generate_offsets(data_batch):
-    offsets = [0]
-    for entry in data_batch:
-        offsets.append(offsets[-1] + len(entry))
-    offsets = torch.tensor(offsets[:-1])
-    return offsets
+def generate_batch(batch):
 
+    def generate_offsets(data_batch):
+        offsets = [0]
+        for entry in data_batch:
+            offsets.append(offsets[-1] + len(entry))
+        offsets = torch.tensor(offsets[:-1])
+        return offsets
 
-def generate_batch(data, labels, i, batch_size):
-    data_batch = data[i:i + batch_size]
-    text = torch.cat(data_batch)
-    offsets = generate_offsets(data_batch)
-    cls = torch.tensor(labels[i:i + batch_size])
-    text, offsets, cls = text.to(device), offsets.to(device), cls.to(device)
+    cls = torch.tensor([entry[0] for entry in batch])
+    text = [entry[1] for entry in batch]
+    offsets = generate_offsets(text)
+    text = torch.cat(text)
     return text, offsets, cls
 
 
-def train(lr_, num_epoch, data, labels):
+def train(lr_, num_epoch, data_):
+    data = DataLoader(data_, batch_size=batch_size, shuffle=True,
+                      collate_fn=generate_batch, num_workers=args.num_workers)
     num_lines = num_epochs * len(data)
     for epoch in range(num_epochs):
-        perm = list(range(len(data)))
-        random.shuffle(perm)
-        data = [data[i] for i in perm]
-        labels = [labels[i] for i in perm]
-
-        for i in range(0, len(data), batch_size):
-            text, offsets, cls = generate_batch(data, labels, i, batch_size)
+        for i, (text, offsets, cls) in enumerate(data):
+            text, offsets, cls = text.to(device), offsets.to(device), cls.to(device)
             output = model(text, offsets)
             loss = criterion(output, cls)
             loss.backward()
@@ -50,11 +46,12 @@ def train(lr_, num_epoch, data, labels):
     print("")
 
 
-def test(data, labels):
+def test(data_):
+    data = DataLoader(data_, batch_size=batch_size, collate_fn=generate_batch)
     total_accuracy = []
-    for i in range(0, len(data), batch_size):
+    for text, offsets, cls in data:
+        text, offsets, cls = text.to(device), offsets.to(device), cls.to(device)
         with torch.no_grad():
-            text, offsets, cls = generate_batch(data, labels, i, batch_size)
             output = model(text, offsets)
             accuracy = (output.argmax(1) == cls).float().mean().item()
             total_accuracy.append(accuracy)
@@ -64,15 +61,18 @@ def test(data, labels):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         description='Train a text classification model on AG_NEWS')
+    parser.add_argument('dataset', choices=text_classification.DATASETS)
     parser.add_argument('--num-epochs', type=int, default=3)
     parser.add_argument('--embed-dim', type=int, default=128)
     parser.add_argument('--batch-size', type=int, default=64)
     parser.add_argument('--lr', type=float, default=64.0)
     parser.add_argument('--ngrams', type=int, default=2)
+    parser.add_argument('--num-workers', type=int, default=1)
     parser.add_argument('--device', default='cpu')
     parser.add_argument('--data', default='.data')
     parser.add_argument('--save-model-path')
-    parser.add_argument('--save-dictionary-path')
+    parser.add_argument('--save-vocab-path')
+    parser.add_argument('--load-vocab-path')
     parser.add_argument('--logging-level', default='WARNING')
     args = parser.parse_args()
 
@@ -82,24 +82,28 @@ def test(data, labels):
     lr = args.lr
     device = args.device
     data = args.data
+    vocab = args.load_vocab_path
 
     logging.basicConfig(level=getattr(logging, args.logging_level))
 
     if not os.path.exists(data):
         print("Creating directory {}".format(data))
         os.mkdir(data)
 
-    dataset = AG_NEWS(root=data, ngrams=args.ngrams)
-    model = TextSentiment(len(dataset.dictionary), embed_dim,
-                          len(set(dataset.labels))).to(device)
+    train_dataset, test_dataset = text_classification.DATASETS[args.dataset](
+        root=data, ngrams=args.ngrams, vocab=vocab)
+
+    if args.save_vocab_path:
+        print("Saving vocab to {}".format(args.save_vocab_path))
+        torch.save(train_dataset.get_vocab, args.save_vocab_path)
+
+    model = TextSentiment(len(train_dataset.get_vocab()),
+                          embed_dim, len(train_dataset.get_labels())).to(device)
     criterion = torch.nn.CrossEntropyLoss().to(device)
 
-    train(lr, num_epochs, dataset.train_data, dataset.train_labels)
-    test(dataset.test_data, dataset.test_labels)
+    train(lr, num_epochs, train_dataset)
+    test(test_dataset)
 
     if args.save_model_path:
         print("Saving model to {}".format(args.save_model_path))
         torch.save(model.to('cpu'), args.save_model_path)
-    if args.save_dictionary_path:
-        print("Saving dictionary to {}".format(args.save_dictionary_path))
-        torch.save(dataset.dictionary, args.save_dictionary_path)

test/test_vocab.py

@@ -356,3 +356,8 @@ def test_vectors_get_vecs(self):
             for dim in ["50", "100", "200", "300"]:
                 conditional_remove(os.path.join(self.project_root, ".vector_cache",
                                                 "glove.6B.{}d.txt".format(dim)))
+
+    def test_has_unk(self):
+        c = Counter({'hello': 4, 'world': 3, 'ᑌᑎIᑕOᗪᕮ_Tᕮ᙭T': 5, 'freq_too_low': 2})
+        v = vocab.Vocab(c)
+        self.assertEqual(v['not_in_it'], 0)

torchtext/data/utils.py

@@ -114,27 +114,25 @@ def dtype_to_attr(dtype):
     return dtype
 
 
-def generate_ngrams(token_list, ngrams):
-    """Generate a list of token up to ngrams.
+def ngrams_iterator(token_list, ngrams):
+    """Return an iterator that yields the given tokens and their ngrams.
 
     Arguments:
         token_list: A list of tokens
         ngrams: the number of ngrams.
 
     Examples:
         >>> token_list = ['here', 'we', 'are']
-        >>> torchtext.data.utils.generate_ngrams(token_list, 2)
+        >>> list(ngrams_iterator(token_list, 2))
         >>> ['here', 'here we', 'we', 'we are', 'are']
     """
 
-    re_list = []
     for i in range(0, len(token_list)):
         x = token_list[i]
-        re_list.append(x)
+        yield x
         for j in range(i + 1, min(i + ngrams, len(token_list))):
             x += ' ' + token_list[j]
-            re_list.append(x)
-    return re_list
+            yield x
 
 
 class RandomShuffler(object):