mirror of
https://github.com/wassname/Castor.git
synced 2026-10-30 12:50:26 +08:00
Achyudh Ram
1b817d3e24
* Add ReutersTrainer, ReutersEvaluator options in Factory classes * Add Reuters to Kim-CNN command line arguments * Fix SST dataset path according to changes in Kim-CNN args The dataset path in args.py was made to point at the dataset folder rather than dataset/SST folder. Hence SST folder was added to paths in the SST dataset class * Add Reuters dataset class, and support in __main__ * Add Reuters dataset trainers and evaluators * Remove debug print statement in reuters_evaluator * Fix rounding bug in reuters_trainer and reuters_evaluator
46 lines
1.6 KiB
Python
46 lines
1.6 KiB
Python
from .trainers.sick_trainer import SICKTrainer
|
|
from .trainers.msrvid_trainer import MSRVIDTrainer
|
|
from .trainers.trecqa_trainer import TRECQATrainer
|
|
from .trainers.wikiqa_trainer import WikiQATrainer
|
|
from .trainers.pit2015_trainer import PIT2015Trainer
|
|
from .trainers.sst_trainer import SSTTrainer
|
|
from .trainers.reuters_trainer import ReutersTrainer
|
|
from nce.nce_pairwise_mp.trainers.trecqa_trainer import TRECQATrainerNCE
|
|
from nce.nce_pairwise_mp.trainers.wikiqa_trainer import WikiQATrainerNCE
|
|
|
|
|
|
class TrainerFactory(object):
|
|
"""
|
|
Get the corresponding Trainer class for a particular dataset.
|
|
"""
|
|
trainer_map = {
|
|
'sick': SICKTrainer,
|
|
'msrvid': MSRVIDTrainer,
|
|
'SST-1': SSTTrainer,
|
|
'SST-2': SSTTrainer,
|
|
'trecqa': TRECQATrainer,
|
|
'wikiqa': WikiQATrainer,
|
|
'pit2015': PIT2015Trainer,
|
|
'twitterurl': PIT2015Trainer,
|
|
'Reuters': ReutersTrainer
|
|
}
|
|
|
|
trainer_map_nce = {
|
|
'trecqa': TRECQATrainerNCE,
|
|
'wikiqa': WikiQATrainerNCE
|
|
}
|
|
|
|
@staticmethod
|
|
def get_trainer(dataset_name, model, embedding, train_loader, trainer_config, train_evaluator, test_evaluator, dev_evaluator=None, nce=False):
|
|
if nce:
|
|
trainer_map = TrainerFactory.trainer_map_nce
|
|
else:
|
|
trainer_map = TrainerFactory.trainer_map
|
|
|
|
if dataset_name not in trainer_map:
|
|
raise ValueError('{} is not implemented.'.format(dataset_name))
|
|
|
|
return trainer_map[dataset_name](
|
|
model, embedding, train_loader, trainer_config, train_evaluator, test_evaluator, dev_evaluator
|
|
)
|