from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
import gradio as gr

# Create translation function
tokenizer = AutoTokenizer.from_pretrained("claudiatang/flan-t5-base-eng-hwp")
model = AutoModelForSeq2SeqLM.from_pretrained("claudiatang/flan-t5-base-eng-hwp")

def translate_eng_hwp(eng_input):
    translator = pipeline("translation", model=model, tokenizer=tokenizer)
    eng_input = "translate English to Hawaiian Pidgin: " + str(eng_input)
    output = translator(eng_input)
    hwp_output = output[0]["translation_text"]
    return hwp_output

# Create Gradio interface
with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="blue")) as demo: 
    gr.HTML("""<br><h1 style="text-align:center; font-weight: bold"> English-Hawaiian Pidgin Translator 🤙 </h1>""")
    with gr.Row():
        with gr.Column():
            english = gr.Textbox(label="English Text", placeholder="Enter English text here...")
        with gr.Column():
            hawaiian_pidgin = gr.Textbox(label="Hawaiian Pidgin Text", placeholder= "Hawaiian Pidgin translation will appear here.")
    with gr.Row():
        translate_btn = gr.Button(value="Translate", variant="primary")
        translate_btn.click(translate_eng_hwp, inputs=english, outputs=hawaiian_pidgin)
    examples = gr.Examples(examples=["I went shopping at Ala Moana yesterday.", "We baked a cake today."],
                           inputs=[english])
    gr.Markdown("""For more information, check out the [model page](https://huggingface.co/claudiatang/flan-t5-base-eng-hwp).""")

demo.launch()