import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
import torch

# Load the multilingual translation model
model_name = "facebook/nllb-200-distilled-600M"
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
device = 0 if torch.cuda.is_available() else -1

# Available languages (extend this list for more languages)
LANGS = {
    "English": "eng_Latn",
    "Tamil": "tam_Taml",
    "Malayalam": "mal_Mlym",
    "Hindi": "hin_Deva",
    "French": "fra_Latn",
    "Spanish": "spa_Latn",
    # Add more languages here as needed
}

def translate(text, src_lang, tgt_lang):
    if not text.strip():
        return ""
    translation_pipeline = pipeline(
        "translation",
        model=model,
        tokenizer=tokenizer,
        src_lang=src_lang,
        tgt_lang=tgt_lang,
        max_length=400,
        device=device,
    )
    result = translation_pipeline(text)
    return result[0]['translation_text']

with gr.Blocks() as demo:
    gr.Markdown("# Multilingual Translator")
    with gr.Row():
        input_text = gr.Textbox(label="Input Text")
    with gr.Row():
        source_lang = gr.Dropdown(label="Source Language", choices=list(LANGS.keys()), value="English")
        target_lang = gr.Dropdown(label="Target Language", choices=list(LANGS.keys()), value="Tamil")
    output_text = gr.Textbox(label="Translated Text")
    translate_btn = gr.Button("Translate")

    def on_translate_click(text, src, tgt):
        src_code = LANGS[src]
        tgt_code = LANGS[tgt]
        return translate(text, src_code, tgt_code)
    
    translate_btn.click(on_translate_click, inputs=[input_text, source_lang, target_lang], outputs=output_text)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)