Views
No views yet
1import torch
2
3
4model = TransformerMT(
5 vocab_size_src=VOCAB_SIZE,
6 vocab_size_tgt=VOCAB_SIZE,
7 d_model=D_MODEL,
8 n_layers=NUM_LAYERS,
9 n_heads=N_HEADS,
10 d_ff=D_FF,
11 dropout=DROPOUT
12)
13ckpt = torch.load(transformer_en_fr.pth, map_location=DEVICE)
14model.load_state_dict(ckpt["model"])
15model.eval()
16
17sentence = "how are you today"
18output = beam_search(model, sentence, beam_size=4, max_len=100)
19print(output)transformer_en_fr.pth — trained model checkpoint