import gradio as gr
from transformers import AutoModelForSeq2SeqLM, NllbTokenizerFast

model_repo = "wtarit/nllb-600M-th-en"

model = AutoModelForSeq2SeqLM.from_pretrained(model_repo)
tokenizer = NllbTokenizerFast.from_pretrained(model_repo, src_lang="tha_Thai", tgt_lang="eng_Latn")

def translate(Text):
    inputs = tokenizer(Text, return_tensors="pt")
    translated_tokens = model.generate(
        **inputs, forced_bos_token_id=tokenizer.convert_tokens_to_ids("eng_Latn"), max_length=64
    )
    return tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]

demo = gr.Interface(
    fn=translate, 
    inputs=[
        gr.components.Textbox(placeholder="Enter Thai text here...")
    ], 
    outputs=["text"],
    title="NLLB TH-EN Translation",
    allow_flagging="never",
    )

demo.launch()