import gradio as gr
from transformers import pipeline

# Load NLLB distill model (English ↔ many, including Malayalam)
# If you prefer mBART‑50, change model to "facebook/mbart-50-many-to-many-mmt"
translator = pipeline(
    "translation",
    model="facebook/nllb-200-distilled-600M",
    src_lang="eng_Latn",
    tgt_lang="mal_Deva",
    max_length=400,
)

LANGS = {
    "English → Malayalam": ("eng_Latn", "mal_Deva"),
    "Malayalam → English": ("mal_Deva", "eng_Latn"),
}

def translate_text(text, direction):
    src, tgt = LANGS[direction]
    # Re‑configure pipeline for the chosen direction
    t = pipeline(
        "translation",
        model="facebook/nllb-200-distilled-600M",
        src_lang=src,
        tgt_lang=tgt,
        max_length=400,
        num_beams=4,
        early_stopping=True,
    )
    result = t(text)
    return result[0]["translation_text"]

demo = gr.Interface(
    fn=translate_text,
    inputs=[
        gr.Textbox(label="Enter text", placeholder="Type here..."),
        gr.Radio(
            list(LANGS.keys()),
            label="Translation direction",
            value="English → Malayalam",
        ),
    ],
    outputs=gr.Textbox(label="Translated text"),
    title="English ↔ Malayalam Translator",
    description="Translate between English and Malayalam using NLLB‑200.",
    examples=[
        [
            "Hello, how are you?",
            "English → Malayalam"
        ],
        [
            "ഞാൻ നല്ലതുപോലെ ഇരിക്കുന്നു.",
            "Malayalam → English"
        ],
    ],
    cache_examples=False,
)

demo.launch()