From 10e4e568172528baa1f622afdaf34a9b0213f373 Mon Sep 17 00:00:00 2001
From: Victor Yang <547106693@qq.com>
Date: Sat, 8 Sep 2018 11:16:23 +0800
Subject: [PATCH] Add twitter url dataset (#145)

* add prediction/qrel files dump option

* fix comment

* add twitter-url dataset, minor refactor

* fix minor error
---
 common/dataset.py                      | 9 ++++++---
 common/evaluation.py                   | 3 ++-
 common/evaluators/pit2015_evaluator.py | 6 ++++--
 common/train.py                        | 3 ++-
 4 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/common/dataset.py b/common/dataset.py
index 5cffeeb..9acf486 100644
--- a/common/dataset.py
+++ b/common/dataset.py
@@ -51,18 +51,21 @@ class DatasetFactory(object):
             return TRECQA, embedding, train_loader, test_loader, dev_loader
         elif dataset_name == 'wikiqa':
             if not os.path.exists(os.path.join(castor_dir, utils_trecqa)):
-                raise FileNotFoundError('TrecQA requires the trec_eval tool to run. Please run get_trec_eval.sh inside Castor/utils (as working directory) before continuing.')
+                raise FileNotFoundError('WikiQA requires the trec_eval tool to run. Please run get_trec_eval.sh inside Castor/utils (as working directory) before continuing.')
             dataset_root = os.path.join(castor_dir, os.pardir, 'Castor-data', 'datasets', 'WikiQA/')
             train_loader, dev_loader, test_loader = WikiQA.iters(dataset_root, word_vectors_file, word_vectors_dir, batch_size, device=device, unk_init=UnknownWordVecCache.unk)
             embedding = nn.Embedding.from_pretrained(WikiQA.TEXT_FIELD.vocab.vectors)
             return WikiQA, embedding, train_loader, test_loader, dev_loader
         elif dataset_name == 'pit2015':
-            if not os.path.exists(os.path.join(castor_dir, utils_trecqa)):
-                raise FileNotFoundError('TrecQA requires the trec_eval tool to run. Please run get_trec_eval.sh inside Castor/utils (as working directory) before continuing.')
             dataset_root = os.path.join(castor_dir, os.pardir, 'Castor-data', 'datasets', 'SemEval-PIT2015/')
             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 == 'twitterurl':
+            dataset_root = os.path.join(castor_dir, os.pardir, 'Castor-data', 'datasets', 'Twitter-URL/')
+            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
         else:
             raise ValueError('{} is not a valid dataset.'.format(dataset_name))
 
diff --git a/common/evaluation.py b/common/evaluation.py
index 2fa467d..8c7f650 100644
--- a/common/evaluation.py
+++ b/common/evaluation.py
@@ -19,7 +19,8 @@ class EvaluatorFactory(object):
         'SST-2': SSTEvaluator,
         'trecqa': TRECQAEvaluator,
         'wikiqa': WikiQAEvaluator,
-        'pit2015': PIT2015Evaluator
+        'pit2015': PIT2015Evaluator,
+        'twitterurl': PIT2015Evaluator
     }
 
     evaluator_map_nce = {
diff --git a/common/evaluators/pit2015_evaluator.py b/common/evaluators/pit2015_evaluator.py
index 5ae2186..b6d6a78 100644
--- a/common/evaluators/pit2015_evaluator.py
+++ b/common/evaluators/pit2015_evaluator.py
@@ -21,8 +21,10 @@ class PIT2015Evaluator(Evaluator):
             n_dev_correct += (prediction == gold_label).sum().item()
             acc_total += ((prediction == batch.label.data) * (prediction == 1)).sum().item()
             total_loss += F.nll_loss(scores, batch.label, size_average=False).item()
-            rel_total += batch.label.data.sum().item()
-            pre_total += torch.max(scores, 1)[1].view(batch.label.size()).data.sum().item()
+            rel_total += gold_label.sum().item()
+            pre_total += prediction.sum().item()
+
+            del scores
 
         precision = acc_total / pre_total
         recall = acc_total / rel_total
diff --git a/common/train.py b/common/train.py
index bf8a54f..3099dfd 100644
--- a/common/train.py
+++ b/common/train.py
@@ -19,7 +19,8 @@ class TrainerFactory(object):
         'SST-2': SSTTrainer,
         'trecqa': TRECQATrainer,
         'wikiqa': WikiQATrainer,
-        'pit2015': PIT2015Trainer
+        'pit2015': PIT2015Trainer,
+        'twitterurl': PIT2015Trainer
     }
 
     trainer_map_nce = {