translation.py

from transformers import pipeline
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

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


def translation(text, src_lang, tgt_lang):
  trans_pipe = pipeline("translation", model=model, tokenizer=tokenizer, src_lang=src_lang, tgt_lang=tgt_lang, max_length=400)
  return trans_pipe(text)[0]["translation_text"]