# -*- coding: utf-8 -*-
"""app.py

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1UBKp9Q5WTAiaXvsRK_6KupUXBI3i05wH
"""

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import gradio as gr

# Dictionary mapping target languages to Helsinki models
LANG_MODELS = {
    "German": "Helsinki-NLP/opus-mt-en-de",
    "French": "Helsinki-NLP/opus-mt-en-fr",
    "Spanish": "Helsinki-NLP/opus-mt-en-es",
    "Italian": "Helsinki-NLP/opus-mt-en-it",
    "Japanese": "Helsinki-NLP/opus-mt-en-jap",
}

# Preload models for speed (optional, you can also load on demand)
model_cache = {}
tokenizer_cache = {}

def translate_text_multi(text, target_lang):
    model_name = LANG_MODELS.get(target_lang)
    if not model_name:
        return "Error: language not supported."

    # Load model/tokenizer if not already loaded
    if model_name not in model_cache:
        tokenizer_cache[model_name] = AutoTokenizer.from_pretrained(model_name)
        model_cache[model_name] = AutoModelForSeq2SeqLM.from_pretrained(model_name)

    tokenizer = tokenizer_cache[model_name]
    model = model_cache[model_name]

    inputs = tokenizer(text, return_tensors="pt")
    outputs = model.generate(**inputs)
    translated = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return translated

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🌐 Multi-language Translator")
    input_text = gr.Textbox(label="Enter English text", lines=4)
    target_language = gr.Dropdown(choices=list(LANG_MODELS.keys()), value="German", label="Target Language")
    output_text = gr.Textbox(label="Translation", lines=4)

    translate_btn = gr.Button("Translate")
    translate_btn.click(fn=translate_text_multi, inputs=[input_text, target_language], outputs=output_text)

demo.launch()