Update app.py

app.py

@@ -1,47 +1,55 @@
 import gradio as gr
 from transformers import MarianMTModel, MarianTokenizer
 
-# Specify the model name from the Hugging Face Hub, for example, an English to French model by the University of Helsinki
-model_name = "Helsinki-NLP/opus-mt-en-fr"
-
-# Load the tokenizer and model
-tokenizer = MarianTokenizer.from_pretrained(model_name)
-model = MarianMTModel.from_pretrained(model_name)
-
-# Function to handle translation
-def translate_text(text, target_language):
-    # Adjust the model_name based on the target language
-    # Note: You'd need to find the exact model names for each language pair you want to support
-    model_name_map = {
-        "French": "Helsinki-NLP/opus-mt-en-fr",
-        "German": "Helsinki-NLP/opus-mt-en-de",
-        "Spanish": "Helsinki-NLP/opus-mt-en-es",
-    }
-    
-    selected_model_name = model_name_map.get(target_language, "Helsinki-NLP/opus-mt-en-fr")
-    
-    # Load the selected model and tokenizer
-    tokenizer = MarianTokenizer.from_pretrained(selected_model_name)
-    model = MarianMTModel.from_pretrained(selected_model_name)
-    
-    # Prepare the text for translation
-    encoded_text = tokenizer.prepare_seq2seq_batch([text], return_tensors="pt")
-    
-    # Perform the translation
-    translated = model.generate(**encoded_text)
-    
-    # Decode the translated text
+# 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.Textbox(lines=2, placeholder="Enter text to translate..."), gr.Dropdown(["French", "German", "Spanish"], label="Select Language")],
-    outputs=[gr.Textbox()],
-    title="Text Translator with Helsinki NLP Models",
-    description="Select a language to translate English text into using University of Helsinki models."
+    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()