import gradio as gr
from transformers import pipeline
import torch

translator = pipeline(task="translation",
                      model="facebook/nllb-200-distilled-600M",
                      torch_dtype=torch.bfloat16)

languages = {
    "English": "eng_Latn",
    "Arabic": "arb_Arab",
    "German": "deu_Latn",
    "Chinese": "zho_Hans",
    "French": "fra_Latn",
    "Italian": "ita_Latn"
}

def translate_text(text, src_lang, tgt_lang):
    result = translator(text, src_lang=languages[src_lang], tgt_lang=languages[tgt_lang])
    return result[0]['translation_text']


demo = gr.Interface(
    fn=translate_text,
    inputs=[
        gr.Textbox(label="Enter text"),
        gr.Dropdown(choices=list(languages.keys()),value="English", label="Source Language"),
        gr.Dropdown(choices=list(languages.keys()), value="French",label="Target Language")
    ],
    outputs=gr.Textbox(label="Translated Text"),
)

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