Update translation.py

translation.py

@@ -30,14 +30,18 @@ def translate(text, source_lang, target_lang):
     if not text:
         return "Please provide text to translate."
     
-    src_code = LANGUAGES.get(source_lang)
-    tgt_code = LANGUAGES.get(target_lang)
+    src_code = LANGUAGES.get(source_lang, "en")
+    tgt_code = LANGUAGES.get(target_lang, "fr")
+    
+    # Check if the language pair is supported
+    if tgt_code not in SUPPORTED_PAIRS.get(src_code, []):
+        raise Exception(f"Translation from {source_lang} to {target_lang} is not supported. Supported pairs: {SUPPORTED_PAIRS.get(src_code, [])}")
     
     # Use preloaded model if en-fr, else load dynamically
     if src_code == "en" and tgt_code == "fr":
         tokenizer, model = DEFAULT_TOKENIZER, DEFAULT_MODEL
     else:
-        tokenizer, model = load_model(src_code, tgt_code) 
+        tokenizer, model = load_model(src_code, tgt_code)
     
     # Perform translation
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=400)
@@ -55,4 +59,17 @@ LANGUAGES = {
     "Russian": "ru",
     "Hindi": "hi",
     "Japanese": "ja"
+}
+
+# Dictionary of supported source-target pairs (based on available MarianMT models)
+SUPPORTED_PAIRS = {
+    "en": ["fr", "es", "de", "zh", "ru"],  # English to French, Spanish, German, Chinese, Russian
+    "fr": ["en"],  # French to English (limited support)
+    "es": ["en"],  # Spanish to English
+    "de": ["en"],  # German to English
+    "zh": ["en"],  # Chinese to English
+    "ru": ["en"],  # Russian to English
+    "hi": [],      # Hindi not supported as source
+    "ar": [],      # Arabic not supported as source
+    "ja": [],      # Japanese not supported as source
 }