NLLB FB 200 Language Translator

nllb_200.py

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline

model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")

#for source and target lang check this - https://github.com/facebookresearch/flores/tree/main/flores200#languages-in-flores-200

def translation(text,src_lang,tgt_lang):
    translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=src_lang, tgt_lang=tgt_lang, max_length = 400)
    return translator(text)