Views
No views yet
1import torch
2from joeynmt.config import load_config, parse_global_args
3from joeynmt.prediction import predict, prepare
4from huggingface_hub import snapshot_download
5
6# Download model
7snapshot_download(
8 repo_id="Koleshjr/dyu-fr-joeynmt-316-epochs_2_layers_8heads_128_384_plateau_2000_7_95_21_48",
9 local_dir="/path/to/save/locally"
10)
11
12# Define model interface
13class JoeyNMTModel:
14 '''
15 JoeyNMTModel which load JoeyNMT model for inference.
16
17 :param config_path: Path to YAML config file
18 :param n_best: return this many hypotheses, <= beam (currently only 1)
19 '''
20 def __init__(self, config_path: str, n_best: int = 1):
21 seed = 42
22 torch.manual_seed(seed)
23 cfg = load_config(config_path)
24 args = parse_global_args(cfg, rank=0, mode="translate")
25 self.args = args._replace(test=args.test._replace(n_best=n_best))
26 # build model
27 self.model, _, _, self.test_data = prepare(self.args, rank=0, mode="translate")
28
29 def _translate_data(self):
30 _, _, hypotheses, trg_tokens, trg_scores, _ = predict(
31 model=self.model,
32 data=self.test_data,
33 compute_loss=False,
34 device=self.args.device,
35 rank=0,
36 n_gpu=self.args.n_gpu,
37 normalization="none",
38 num_workers=self.args.num_workers,
39 args=self.args.test,
40 autocast=self.args.autocast,
41 )
42 return hypotheses, trg_tokens, trg_scores
43
44 def translate(self, sentence) -> list:
45 '''
46 Translate the given sentence.
47
48 :param sentence: Sentence to be translated
49 :return:
50 - translations: (list of str) possible translations of the sentence.
51 '''
52 self.test_data.set_item(sentence.strip())
53 translations, _, _ = self._translate_data()
54 assert len(translations) == len(self.test_data) * self.args.test.n_best
55 self.test_data.reset_cache()
56 return translations
57
58# Load model
59config_path = "/path/to/lean_model/config_local.yaml" # Change this to the path to your model congig file
60model = JoeyNMTModel(config_path=config_path, n_best=1)
61
62# Translate
63model.translate(sentence="i tɔgɔ bi cogodɔ")