import gradio as gr
from transformers import MarianMTModel, MarianTokenizer

# Function to dynamically load the model and tokenizer based on selected languages
def translate_text(text, source_language, target_language):
    # Construct model name based on selected languages
    model_name = f"Helsinki-NLP/opus-mt-{source_language}-{target_language}"

    # Load tokenizer and model
    try:
        tokenizer = MarianTokenizer.from_pretrained(model_name)
        model = MarianMTModel.from_pretrained(model_name)
    except Exception as e:
        return f"Failed to load model for {source_language} to {target_language}: {str(e)}"

    # Translate text
    translated = model.generate(**tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512))
    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
    
    return translated_text

# Define language options (ISO 639-1 codes and names)
# Note: This is a simplified subset for demonstration. Expand based on available models.
language_options = [
    ('en', 'English'),
    ('es', 'Spanish'),
    ('fr', 'French'),
    ('de', 'German'),
    ('zh', 'Chinese'),
    ('ru', 'Russian'),
    ('ar', 'Arabic'),
    ('it', 'Italian'),
    ('pt', 'Portuguese'),
    ('nl', 'Dutch'),
    # Add more languages as needed
]

# Convert language options to the format expected by the dropdown
language_dropdown_options = [(code, f"{name} ({code})") for code, name in language_options]

# Create dropdowns for source and target languages
source_language_dropdown = gr.inputs.Dropdown(choices=language_dropdown_options, label="Source Language")
target_language_dropdown = gr.inputs.Dropdown(choices=language_dropdown_options, label="Target Language")

# Define the interface
iface = gr.Interface(
    fn=translate_text,
    inputs=[gr.inputs.Textbox(lines=2, placeholder="Enter text to translate..."), source_language_dropdown, target_language_dropdown],
    outputs=gr.outputs.Textbox(),
    title="Text Translator with Dynamic Helsinki NLP Models",
    description="Select source and target languages to translate text using Helsinki NLP models."
)

# Launch the app
iface.launch()