import gradio as gr
import os
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

HF_TOKEN = os.environ.get("HF_TOKEN")

MODEL_ID = "toiar/nllb-finetuned-english-pnar"

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    token=HF_TOKEN
)
tokenizer.src_lang = "eng_Latn"

# Load model
model = AutoModelForSeq2SeqLM.from_pretrained(
    MODEL_ID,
    token=HF_TOKEN,
    torch_dtype="auto",
    low_cpu_mem_usage=True,
    device_map="cpu"
)

def translate(text):
    if not text or not text.strip():
        return ""

    inputs = tokenizer(text, return_tensors="pt")

    output = model.generate(
        **inputs,
        forced_bos_token_id=tokenizer.convert_tokens_to_ids("pbv_Latn"),
        max_length=128,
        num_beams=5,
    )

    return tokenizer.decode(output[0], skip_special_tokens=True)

examples = [
    ["Please close the door before you leave."],
    ["She forgot her umbrella at home."],
    ["We will meet again after the festival ends."],
    ["He speaks calmly even when he is angry."],
    ["The river becomes wider during the rainy season."],
    ["They are learning new skills to improve their future."]
]

with gr.Blocks(theme=gr.themes.Soft()) as demo:

    gr.Markdown(
        """
        <div style="text-align: center;">
            <h1>English → Pnar Translator</h1>
            <p><i>Powered by a fine-tuned NLLB-200 model</i></p>
        </div>
        """
    )

    with gr.Row():
        eng_in = gr.Textbox(
            label="English",
            placeholder="Type your English text here...",
            lines=8
        )
        pnar_out = gr.Textbox(
            label="Pnar",
            placeholder="Translation will appear here...",
            lines=8,
            interactive=False
        )

    with gr.Row():
        trans_btn = gr.Button("Translate", variant="primary")
        clear_btn = gr.Button("Clear")

    gr.Markdown("### Examples")
    gr.Examples(
        examples=examples,
        inputs=[eng_in],
        cache_examples=False
    )

    # Events
    trans_btn.click(translate, inputs=eng_in, outputs=pnar_out)
    eng_in.submit(translate, inputs=eng_in, outputs=pnar_out)
    clear_btn.click(lambda: ("", ""), outputs=[eng_in, pnar_out])

demo.launch()