class TestPredictor(unittest.TestCase):
@classmethod
def setUpClass(self):
test_path = os.path.dirname(os.path.realpath(__file__))
src = SourceField(batch_first=True)
trg = TargetField(batch_first=True)
dataset = torchtext.data.TabularDataset(
path=os.path.join(test_path, 'data/eng-fra.txt'),
format='tsv',
fields=[('src', src), ('trg', trg)],
)
src.build_vocab(dataset)
trg.build_vocab(dataset)
encoder = EncoderRNN(len(src.vocab), 5, 10, 10, rnn_cell='lstm')
decoder = DecoderRNN(len(trg.vocab),
10,
10,
trg.sos_id,
trg.eos_id,
rnn_cell='lstm')
seq2seq = Seq2seq(encoder, decoder)
self.predictor = Predictor(seq2seq, src.vocab, trg.vocab)
def test_predict(self):
src_seq = ["I", "am", "fat"]
tgt_seq = self.predictor.predict(src_seq)
for tok in tgt_seq:
self.assertTrue(tok in self.predictor.tgt_vocab.stoi)
def setUpClass(self):
test_path = os.path.dirname(os.path.realpath(__file__))
src = SourceField(batch_first=True)
trg = TargetField(batch_first=True)
dataset = torchtext.data.TabularDataset(
path=os.path.join(test_path, 'data/eng-fra.txt'),
format='tsv',
fields=[('src', src), ('trg', trg)],
)
src.build_vocab(dataset)
trg.build_vocab(dataset)
encoder = EncoderRNN(len(src.vocab), 5, 10, 10, rnn_cell='lstm')
decoder = DecoderRNN(len(trg.vocab),
10,
10,
trg.sos_id,
trg.eos_id,
rnn_cell='lstm')
seq2seq = Seq2seq(encoder, decoder)
self.predictor = Predictor(seq2seq, src.vocab, trg.vocab)
logging.info(opt)
if torch.cuda.is_available():
logging.info("Cuda device set to %i" % opt.cuda_device)
torch.cuda.set_device(opt.cuda_device)
##########################################################################
# load model
logging.info("loading checkpoint from {}".format(
os.path.join(opt.checkpoint_path)))
checkpoint = Checkpoint.load(opt.checkpoint_path)
seq2seq = checkpoint.model
input_vocab = checkpoint.input_vocab
output_vocab = checkpoint.output_vocab
##########################################################################
# Generate predictor
predictor = Predictor(seq2seq, input_vocab, output_vocab)
if opt.debug:
exit()
while True:
seq_str = raw_input("\n\nType in a source sequence: ")
if seq_str == 'q':
exit()
seq = seq_str.strip().split()
print(predictor.predict(seq))