Update app.py

app.py

@@ -1,38 +1,40 @@
 import gradio as gr
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
 
-# Load the translation model and tokenizer from Hugging Face
-model_name = "Helsinki-NLP/opus-mt-en-fr"  # Example model for English to French
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-
-# Define the translation function
-def translate_text(text, target_language):
-    model_name = f"Helsinki-NLP/opus-mt-en-{target_language}"
-    translator = pipeline("translation", model=model_name)
-    translation = translator(text)[0]['translation_text']
-    return translation
-
-# Define available languages
+# Define available languages with their corresponding model suffixes
 languages = {
-    "english": "English",
-    "spanish": "Spanish",
-    "french": "French",
-    "german": "German",
-    "chinese": "Chinese (Mandarin)",
-    "japanese": "Japanese",
-    "korean": "Korean",
-    "italian": "Italian",
-    "portuguese": "Portuguese",
-    "russian": "Russian",
-    "hindi": "Hindi",
-    "arabic": "Arabic",
-    "dutch": "Dutch",
-    "turkish": "Turkish",
-    "greek": "Greek",
-    "urdu": "Urdu"
+    "english": "en",
+    "spanish": "es",
+    "french": "fr",
+    "german": "de",
+    "chinese": "zh",
+    "japanese": "ja",
+    "korean": "ko",
+    "italian": "it",
+    "portuguese": "pt",
+    "russian": "ru",
+    "hindi": "hi",
+    "arabic": "ar",
+    "dutch": "nl",
+    "turkish": "tr",
+    "greek": "el",
+    "urdu": "ur"
 }
 
+# Define the translation function
+def translate_text(text, target_language):
+    # Get the correct model name based on the target language
+    target_language_code = languages.get(target_language.lower())
+    
+    if target_language_code:
+        model_name = f"Helsinki-NLP/opus-mt-en-{target_language_code}"
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        translator = pipeline("translation", model=model, tokenizer=tokenizer)
+        translation = translator(text)[0]['translation_text']
+        return translation
+    else:
+        return "Error: Language not supported or incorrect."
 
 # Set up the Gradio interface with a submit button and side-by-side layout
 with gr.Blocks() as iface: