class TestReCoSa(unittest.TestCase):
def setUp(self):
self.config = Config()
self.config.add_model("./conf/model/ReCoSa_test.yml")
self.config.add_trainer("./conf/trainer/ReCoSa_test.yml")
self.data = UbuntuDataSet(
"./" + "tests/resources/Ubuntu/",
"sample.csv",
self.config.model.max_seq,
"Ubuntu",
self.config.model.max_turns,
)
self.trainDataLoader = UbuntuDataLoader(
self.data,
**self.config.trainer.data,
collate_fn=collate,
)
self.valDataLoader = UbuntuDataLoader(
self.data,
**self.config.trainer.data,
collate_fn=collate,
)
self.model = RecoSAPL(self.config, len(self.data))
self.trainer = pl.Trainer(**self.config.trainer.pl)
def test_trainer(self):
self.assertFalse(self.config.trainer.data.shuffle)
_ = self.trainer.fit(self.model, self.trainDataLoader,
self.valDataLoader)
def main(
config_data_file: str,
config_model_file: str,
config_trainer_file: str,
config_api_file: str,
version: str,
) -> None:
# TODO: to be removed
_ = build({"data_config": config_data_file, "version": version})
cfg = Config()
cfg.add_dataset(config_data_file)
cfg.add_model(config_model_file)
cfg.add_api(config_api_file)
cfg.add_trainer(config_trainer_file)
val_data = UbuntuDataSet(
cfg.dataset.root + cfg.dataset.target,
cfg.dataset.raw.val,
cfg.model.max_seq,
cfg.dataset.target,
cfg.model.max_turns,
)
val_dataloader = UbuntuDataLoader(
val_data,
batch_size=cfg.model.batch_size,
shuffle=False,
num_workers=8,
collate_fn=collate,
)
model = RecoSAPL.load_from_checkpoint(checkpoint_path=cfg.api.model_path,
config=cfg)
cfg.trainer.pl.max_epochs = 1
trainer = pl.Trainer(**cfg.trainer.pl,
logger=False,
checkpoint_callback=False)
test_result = trainer.test(model, test_dataloaders=val_dataloader)
logger.info(test_result)
bleu_score_4 = bleuS_4(model.pred, model.target)
bleu_score_2 = bleuS_2(model.pred, model.target)
logger.info(bleu_score_4)
logger.info(bleu_score_2)