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_code, target_language_code):
    # Construct model name using ISO 639-1 codes
    model_name = f"Helsinki-NLP/opus-mt-{source_language_code}-{target_language_code}"

    # Check if source and target languages are the same, which is not supported for translation
    if source_language_code == target_language_code:
        return "Translation between the same languages is not supported."

    # 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_code} to {target_language_code}: {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)
language_options = [
    ('en', 'English (en)'),
    ('es', 'Spanish (es)'),
    ('fr', 'French (fr)'),
    ('de', 'German (de)'),
    ('zh', 'Chinese (zh)'),
    ('ru', 'Russian (ru)'),
    ('ar', 'Arabic (ar)'),
    ('it', 'Italian (it)'),
    ('pt', 'Portuguese (pt)'),
    ('nl', 'Dutch (nl)'),
    # Add more languages as needed
]

# Create dropdowns for source and target languages, using only the codes for value
source_language_dropdown = gr.inputs.Dropdown(choices=[(code, name) for code, name in language_options], label="Source Language")
target_language_dropdown = gr.inputs.Dropdown(choices=[(code, name) for code, name in language_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()