m2m-translator / usage.py
datasetsANDmodels's picture
Rename m2m.py to usage.py
846c95f verified
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)