import gradio
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer


model_name = "zaanind/nllb-ensi-v1-tuning"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

tokenizer.src_lang = "eng_Latn"

def translate(text):
    inputs = tokenizer(text=text, return_tensors="pt")
    translated_tokens = model.generate(
        **inputs, forced_bos_token_id=tokenizer.lang_code_to_id["sin_Sinh"]
    )
    translation = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
    return translation

    
def nmtapifunc(text):
    text = translate(text)
    return text 

gradio_interface = gradio.Interface(
  fn=nmtapifunc,
  inputs="text",
  outputs="text",
  title="En-Si NMT",
  description="",
  article="© zaanind 2024"
)
gradio_interface.launch()