File size: 661 Bytes
486a17b 846c95f 486a17b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 |
import sys
import time
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
sr= sys.argv[1]
tr=sys.argv[2]
model = M2M100ForConditionalGeneration.from_pretrained("m2m-translator")
tokenizer = M2M100Tokenizer.from_pretrained("m2m-translator")
tokenizer.src_lang =sr
text="I need to buy a tablet for my mum as soon as possible"
#text="Tôi cần mua một viên thuốc cho mẹ tôi càng sớm càng tốt" VIET
encoded = tokenizer(text, return_tensors="pt")
generated_tokens = model.generate(**encoded, forced_bos_token_id=tokenizer.get_lang_id(tr))
res=(tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0])
print(res)
|