import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Model name (your fine-tuned model)
MODEL_NAME = "Sazid2/assamese-english-translator"

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

def translate(text):
    """Translate Assamese → English"""
    if not text.strip():
        return "⚠️ Please enter Assamese text."
    
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    outputs = model.generate(**inputs, max_length=128)
    translated = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return translated

title = "🌐 Assamese → English Translator"
description = """
### 🧠 Fine-tuned Neural Machine Translation  
This model translates **Assamese sentences to English** using a custom **Assamese-English parallel corpus (~20k sentences)**.  
It is built on top of the `Helsinki-NLP/opus-mt-mul-en` architecture.

- 🔤 **Source language:** Assamese  
- 🌍 **Target language:** English  
- 🧩 **BLEU Score:** 38.02  
- 🧠 **Framework:** Hugging Face Transformers  
"""

examples = [
    ["মই কামলৈ গৈ আছো।"],
    ["তুমি ক'ত আছা?"],
    ["তেও অত্যন্ত ধুনীয়া।"]
]

# Create Gradio interface
demo = gr.Interface(
    fn=translate,
    inputs=gr.Textbox(label="Enter Assamese text"),
    outputs=gr.Textbox(label="English Translation"),
    title=title,
    description=description,
    examples=examples,
    article="---\n**Fine-tuned by:** Abu Sazid Ahmed 🧑‍💻",
)

if __name__ == "__main__":
    demo.launch()