def test(opts): french_word2id = a2_dataloader.read_word2id_from_file(opts.french_vocab) english_word2id = a2_dataloader.read_word2id_from_file(opts.english_vocab) dataloader = a2_dataloader.HansardDataLoader( opts.testing_dir, french_word2id, english_word2id, opts.source_lang, batch_size=opts.batch_size, pin_memory=(opts.device.type == 'cuda')) del french_word2id, english_word2id model = init(opts, dataloader) state_dict = torch.load(opts.model_path) model.load_state_dict(state_dict) del state_dict model.to(opts.device) model.eval() bleu = a2_training_and_testing.compute_average_bleu_over_dataset( model, dataloader, dataloader.dataset.target_sos, dataloader.dataset.target_eos, opts.device, ) print(f'The average BLEU score over the test set was {bleu}')
def train(opts): torch.manual_seed(opts.seed) french_word2id = a2_dataloader.read_word2id_from_file(opts.french_vocab) english_word2id = a2_dataloader.read_word2id_from_file(opts.english_vocab) train_prefixes = opts.train_prefixes.read().strip().split('\n') train_dataloader = a2_dataloader.HansardDataLoader( opts.training_dir, french_word2id, english_word2id, opts.source_lang, train_prefixes, batch_size=opts.batch_size, shuffle=True, pin_memory=(opts.device.type == 'cuda'), num_workers=1, ) del train_prefixes dev_prefixes = opts.dev_prefixes.read().strip().split('\n') dev_dataloader = a2_dataloader.HansardDataLoader( opts.training_dir, french_word2id, english_word2id, opts.source_lang, dev_prefixes, batch_size=opts.batch_size, pin_memory=(opts.device.type == 'cuda'), num_workers=1, ) del dev_prefixes, french_word2id, english_word2id model = init(opts, train_dataloader) model.to(opts.device) optimizer = torch.optim.Adam(model.parameters()) best_bleu = 0. num_poor = 0 epoch = 1 if opts.patience is None: max_epochs = opts.epochs patience = float('inf') else: max_epochs = float('inf') patience = opts.patience while epoch <= max_epochs and num_poor < patience: model.train() loss = a2_training_and_testing.train_for_epoch( model, train_dataloader, optimizer, opts.device) model.eval() bleu = a2_training_and_testing.compute_average_bleu_over_dataset( model, dev_dataloader, dev_dataloader.dataset.target_sos, dev_dataloader.dataset.target_eos, opts.device, ) print(f'Epoch {epoch}: loss={loss}, BLEU={bleu}') if bleu < best_bleu: num_poor += 1 else: num_poor = 0 best_bleu = bleu epoch += 1 if epoch > max_epochs: print(f'Finished {max_epochs} epochs') else: print(f'BLEU did not improve after {patience} epochs. Done.') model.cpu() torch.save(model.state_dict(), opts.model_path)