import gradio as gr
from transformers import M2M100ForConditionalGeneration
from tokenization_small100 import SMALL100Tokenizer

def th2en(th_text: str) -> str: 
    """
    Translates the input text from Thai to English.
    """
    encoded_th_text = tokenizer(th_text, return_tensors="pt")
    generated_tokens = model.generate(**encoded_th_text)
    translated_text = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
    return translated_text

if __name__ == "__main__":
    model_checkpoint = "kimmchii/small100-th"
    TARGET_LANG = "en"
    
    # Initialize model
    model = M2M100ForConditionalGeneration.from_pretrained(model_checkpoint)
    tokenizer = SMALL100Tokenizer.from_pretrained(model_checkpoint)
    tokenizer.tgt_lang = TARGET_LANG

    # Web app section
    with gr.Blocks() as demo:
        gr.Markdown("Translates Thai to English")

        text_input = gr.Textbox(placeholder="Thai Text Here...")
        text_output = gr.Textbox()
        text_button = gr.Button("Translate")

        text_button.click(th2en, inputs=text_input, outputs=text_output)

    demo.launch(debug=True)