lang-autoencoder/translate.py at main · timmy-feng/lang-autoencoder · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
import sentencepiece as spm
import argparse
import yaml

import torch
import torch.nn.functional as F

from utils.model import Autoencoder

UNK_ID, SOS_ID, EOS_ID, PAD_ID = 0, 1, 2, 3

parser = argparse.ArgumentParser()
parser.add_argument('config', type=str)
args = parser.parse_args()

config = yaml.safe_load(open(args.config, 'r'))

num_epochs = config['train']['epochs']

input_vocab_size = config['lang']['input']['vocab_size']
output_vocab_size = config['lang']['output']['vocab_size']
input_len = config['lang']['input']['seq_len']

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = Autoencoder(
    input_vocab_size = config['lang']['input']['vocab_size'],
    output_vocab_size = config['lang']['output']['vocab_size'],
    input_len = config['lang']['input']['seq_len'],
    output_len = config['lang']['output']['seq_len'],
    d_model = config['model']['d_model'],
    nhead = config['model']['nhead'],
    num_layers = config['model']['num_encoder_layers'],
    dim_feedforward = config['model']['dim_feedforward'],
    dropout = 0,
    discrete = True,
    temperature = config['predict']['temperature'],
)

model.load_state_dict(torch.load(config['model']['save_path'], map_location=device))
sp = spm.SentencePieceProcessor(model_file=config['predict']['tokenizer'])
model.to(device)
model.eval()

print('Model loaded. Ready for inference.\n')

outputs = {
    'tokenized': False,
    'translation': True,
    'backtranslation': True,
}

while True:
    print('>>> ', end = '')
    sentence = input().strip()
    print()

    # empty line; skip
    if len(sentence) == 0:
        continue

    # toggling option
    if sentence[0] == '!':
        sentence = sentence.strip('! ')
        if sentence == 'quit':
            exit()
        if sentence not in outputs:
            options = list(outputs.keys()) + ['quit']
            print(f'Invalid option. Options include: {options}')
        else:
            outputs[sentence] = not outputs[sentence]
            print(f'Options updated: {outputs}\n')
        continue

    tokens = sp.encode(sentence)
    src = torch.tensor([[SOS_ID] + tokens + [EOS_ID] + (input_len - len(tokens) - 2) * [PAD_ID]])
    src = F.one_hot(src, input_vocab_size).float().to(device)

    translation = model.translate(src)
    translation_one_hot = F.one_hot(translation, output_vocab_size).float()
    backtranslation = model.backtranslate(translation_one_hot)

    if outputs['tokenized']:
        print(f'Tokenized: ', end = '')
        print(*tokens)
        print()

    if outputs['translation']:
        print(f'Translation: ', end = '')
        print(*map(int, translation[0]))
        print()

    if outputs['backtranslation']:
        backtranslation = list(filter(lambda token: token != PAD_ID, map(int, backtranslation[0])))
        print(f'Backtranslation: {sp.decode(backtranslation)}\n')