Views
No views yet
bfloat16.1from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration
2from threading import Lock
3
4model_name = 'Didier/m2m100-12B-avg-5-ckpt'
5
6device = 'mps' # if on Apple silicon
7tokenizer = M2M100Tokenizer.from_pretrained(model_name)
8model = M2M100ForConditionalGeneration.from_pretrained(
9 model_name, device_map=device, low_cpu_mem_usage=True)
10lock = Lock()
11
12def translate(text: str, src_lang: str, tgt_lang: str) -> str:
13 # Acquire lock to set src_lang and tokenize atomically
14 with lock:
15 tokenizer.src_lang = src_lang
16 input_ids = tokenizer([text,], return_tensors="pt").input_ids.to(model.device)
17
18 # Generate translation (outside the lock to allow parallel model
19 outputs = model.generate(
20 input_ids=input_ids,
21 forced_bos_token_id=tokenizer.get_lang_id(tgt_lang))
22 translation = tokenizer.batch_decode(
23 outputs, skip_special_tokens=True)[0]
24
25 return translation
26
27
28text = "ist der Ruf erst ruiniert, lebt es sich ganz ungeniert."
29src_lang = 'de'
30tgt_lang = 'en'
31
32translation = translate(text, src_lang, tgt_lang)
33print(f"{translation=}")
34
35# --> "Once your reputation is ruined, you can live quite freely."