from transformers import pipeline
import torch
import gradio as gr


translator = pipeline("translation", model="facebook/nllb-200-distilled-600M")

languages = {
    "English": "eng_Latn",
    "French": "fra_Latn",
    "Spanish": "spa_Latn",
    "German": "deu_Latn",
    "Arabic": "arb_Arab",
    "Chinese": "zho_Hans",
    "Russian": "rus_Cyrl"
}

def translate_text(text, src_lang, tgt_lang):
    if not text.strip():
        return "Please enter some text to translate."
    translation = translator(text, src_lang=languages[src_lang], tgt_lang=languages[tgt_lang])
    return translation[0]["translation_text"]

with gr.Blocks() as demo:
    gr.Markdown("## Translation using NLLB-200")
    gr.Markdown("Select source and target languages, then translate the text.")

    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="Input Text", placeholder="Type here...", lines=3)
            src_lang = gr.Dropdown(choices=list(languages.keys()), value="English", label="Source Language")
            tgt_lang = gr.Dropdown(choices=list(languages.keys()), value="French", label="Target Language")
            with gr.Row():
                clear_btn = gr.Button("Clear")
                submit_btn = gr.Button("Submit", variant="primary")

        with gr.Column():
            output_text = gr.Textbox(label="Translated Text", interactive=False, lines=3)
            flag_btn = gr.Button("Flag")

    submit_btn.click(translate_text, inputs=[input_text, src_lang, tgt_lang], outputs=output_text)
    clear_btn.click(lambda: ("", "English", "French"), outputs=[input_text, src_lang, tgt_lang])

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