import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Dictionary to hold models and tokenizers for each language
models = {}
tokenizers = {}

# List of language pairs
language_pairs = {
    "English to French": "Helsinki-NLP/opus-mt-en-fr",
    "English to Chinese": "Helsinki-NLP/opus-mt-en-zh",
    "English to German": "Helsinki-NLP/opus-mt-en-de",
    "English to Urdu": "Helsinki-NLP/opus-mt-en-ur"
}

# Load models and tokenizers for each language pair
for lang, model_name in language_pairs.items():
    tokenizers[lang] = AutoTokenizer.from_pretrained(model_name)
    models[lang] = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Function to perform translation
def translate_text(text, language_choice):
    # Select the appropriate tokenizer and model based on the chosen language
    tokenizer = tokenizers[language_choice]
    model = models[language_choice]
    
    # Tokenize the input text
    inputs = tokenizer(text, return_tensors="pt", truncation=True)
    
    # Generate the translation
    outputs = model.generate(**inputs)
    
    # Decode the translated text
    translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    return translated_text

# Define the Gradio interface
def gradio_interface(text, language_choice):
    translated_text = translate_text(text, language_choice)
    return translated_text

# Create a list of language choices for the dropdown
language_choices = list(language_pairs.keys())

# Set up the Gradio app
interface = gr.Interface(
    fn=gradio_interface,
    inputs=[gr.Textbox(lines=2, placeholder="Enter text here..."), 
            gr.Dropdown(choices=language_choices, label="Select Target Language")],
    outputs=gr.Textbox(label="Translated Text"),
    title="Multi-Language Translation App"
)

# Launch the app
interface.launch()