diff --git a/common/dataset.py b/common/dataset.py index 9acf486..e766148 100644 --- a/common/dataset.py +++ b/common/dataset.py @@ -8,7 +8,7 @@ from datasets.msrvid import MSRVID from datasets.trecqa import TRECQA from datasets.wikiqa import WikiQA from datasets.pit2015 import PIT2015 - +from datasets.reuters import Reuters class UnknownWordVecCache(object): """ @@ -66,6 +66,11 @@ class DatasetFactory(object): train_loader, dev_loader, test_loader = PIT2015.iters(dataset_root, word_vectors_file, word_vectors_dir, batch_size, device=device, unk_init=UnknownWordVecCache.unk) embedding = nn.Embedding.from_pretrained(PIT2015.TEXT_FIELD.vocab.vectors) return PIT2015, embedding, train_loader, test_loader, dev_loader + elif dataset_name == 'reuters': + dataset_root = os.path.join(castor_dir, os.pardir, 'Castor-data', 'datasets', 'Reuters-21578/') + train_loader, dev_loader, test_loader = Reuters.iters(dataset_root, word_vectors_file, word_vectors_dir, batch_size, device=device, unk_init=UnknownWordVecCache.unk) + embedding = nn.Embedding.from_pretrained(Reuters.TEXT_FIELD.vocab.vectors) + return Reuters, embedding, train_loader, test_loader, dev_loader else: raise ValueError('{} is not a valid dataset.'.format(dataset_name))