app.py

import gradio
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
import torch

model = AutoModelForSeq2SeqLM.from_pretrained("Jayyydyyy/m2m100_418m_tokipona")
tokenizer = AutoTokenizer.from_pretrained("facebook/m2m100_418M")
device = "cuda:0" if torch.cuda.is_available() else "cpu"
model.to(device)
LANG_CODES = {"English": "en", "toki pona": "tl"}


def translate(text):
    """
    Translate the text from source lang to target lang
    """

    # src = LANG_CODES.get(src_lang)
    # tgt = LANG_CODES.get(tgt_lang)

    tokenizer.src_lang = "en"
    tokenizer.tgt_lang = "tl"
    ins = tokenizer(text, return_tensors="pt").to(device)

    gen_args = {
        "return_dict_in_generate": True,
        "output_scores": True,
        "output_hidden_states": True,
        "length_penalty": 0.0,  # don't encourage longer or shorter output,
        "num_return_sequences": 1,
        "num_beams": 1,
        "forced_bos_token_id": tokenizer.lang_code_to_id["tl"],
    }

    outs = model.generate(**{**ins, **gen_args})
    output = tokenizer.batch_decode(outs.sequences, skip_special_tokens=True)

    return "\n".join(output)

    
gradio_interface = gradio.Interface(
  fn = translate,
  inputs = "text",
  outputs = "text"
)
gradio_interface.launch()