import gradio as gr
from transformers import pipeline

# English → Bemba
eng2bem = pipeline(
    "translation",
    model="kreasof-ai/nllb-200-600M-eng2bem",
    src_lang="eng_Latn",
    tgt_lang="bem_Latn"
)

# Bemba → English
bem2eng = pipeline(
    "translation",
    model="kreasof-ai/nllb-200-600M-bem2eng-bigc-tatoeba",
    src_lang="bem_Latn",
    tgt_lang="eng_Latn"
)

def translate(direction, text):
    if not text.strip():
        return ""
    if direction == "English → Bemba":
        result = eng2bem(text)
        return result[0]["translation_text"]
    else:  # Bemba → English
        result = bem2eng(text)
        return result[0]["translation_text"]

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🌍 English ↔ Bemba Translator  
        Translate text instantly between **English** and **Bemba (Zambia)**.  
        """
    )

    with gr.Row():
        direction = gr.Radio(
            ["English → Bemba", "Bemba → English"],
            value="English → Bemba",
            label="Translation Direction",
        )

    with gr.Row():
        input_text = gr.Textbox(
            lines=4, placeholder="Enter text here...", label="Input Text"
        )

    with gr.Row():
        output_text = gr.Textbox(
            lines=4, placeholder="Translation will appear here...", label="Output Text"
        )

    translate_btn = gr.Button("🔄 Translate")

    translate_btn.click(
        fn=translate, inputs=[direction, input_text], outputs=output_text
    )

demo.launch()