Update app.py

app.py

@@ -121,30 +121,42 @@ for lang, model_id in TTS_MODELS.items():
         logger.debug(f"Stack trace: {traceback.format_exc()}")
         tts_models[lang] = None
 
-# Add this with your other model configurations
-TRANSLATION_MODEL_ID = "Coco-18/opus-mt-pam-en"
-logger.info(f"🔄 Loading Translation model: {TRANSLATION_MODEL_ID}")
+# Replace the single translation model with a dictionary of models
+TRANSLATION_MODELS = {
+    "pam-eng": "Coco-18/opus-mt-pam-en",
+    "eng-pam": "Coco-18/opus-mt-en-pam",
+    "tgl-eng": "Helsinki-NLP/opus-mt-tl-en",
+    "eng-tgl": "Helsinki-NLP/opus-mt-en-tl"
+    # pam-tgl and tgl-pam will be added later
+}
 
-# Initialize translation model and tokenizer (add this after your other model initializations)
-translation_model = None
-translation_tokenizer = None
+logger.info(f"🔄 Loading Translation model: {TRANSLATION_MODELS}")
 
-try:
-    translation_tokenizer = MarianTokenizer.from_pretrained(
-        TRANSLATION_MODEL_ID,
-        cache_dir=cache_dirs["TRANSFORMERS_CACHE"]
-    )
-    logger.info("✅ Translation tokenizer loaded successfully")
+# Replace the single model initialization with:
+translation_models = {}
+translation_tokenizers = {}
+
+for lang_pair, model_id in TRANSLATION_MODELS.items():
+    logger.info(f"🔄 Loading Translation model: {model_id}")
     
-    translation_model = MarianMTModel.from_pretrained(
-        TRANSLATION_MODEL_ID,
-        cache_dir=cache_dirs["TRANSFORMERS_CACHE"]
-    )
-    translation_model.to(device)
-    logger.info(f"✅ Translation model loaded successfully on {device}")
-except Exception as e:
-    logger.error(f"❌ Error loading Translation model: {str(e)}")
-    logger.debug(f"Stack trace: {traceback.format_exc()}")
+    try:
+        translation_tokenizers[lang_pair] = MarianTokenizer.from_pretrained(
+            model_id,
+            cache_dir=cache_dirs["TRANSFORMERS_CACHE"]
+        )
+        logger.info(f"✅ Translation tokenizer loaded successfully for {lang_pair}")
+        
+        translation_models[lang_pair] = MarianMTModel.from_pretrained(
+            model_id,
+            cache_dir=cache_dirs["TRANSFORMERS_CACHE"]
+        )
+        translation_models[lang_pair].to(device)
+        logger.info(f"✅ Translation model loaded successfully on {device} for {lang_pair}")
+    except Exception as e:
+        logger.error(f"❌ Error loading Translation model for {lang_pair}: {str(e)}")
+        logger.debug(f"Stack trace: {traceback.format_exc()}")
+
+
 
 # Constants
 SAMPLE_RATE = 16000
@@ -350,10 +362,6 @@ def download_audio(filename):
 
 @app.route("/translate", methods=["POST"])
 def translate_text():
-    if translation_model is None or translation_tokenizer is None:
-        logger.error("❌ Translation endpoint called but models aren't loaded")
-        return jsonify({"error": "Translation model not available"}), 503
-        
     try:
         data = request.get_json()
         if not data:
@@ -367,20 +375,40 @@ def translate_text():
         if not source_text:
             logger.warning("⚠️ Translation request with empty text")
             return jsonify({"error": "No text provided"}), 400
+        
+        # Map language names to codes
+        source_code = LANGUAGE_CODES.get(source_language, source_language)
+        target_code = LANGUAGE_CODES.get(target_language, target_language)
+        
+        # Create the language pair key
+        lang_pair = f"{source_code}-{target_code}"
             
         logger.info(f"🔄 Translating from {source_language} to {target_language}: '{source_text}'")
         
+        # Check if we have a model for this language pair
+        if lang_pair not in translation_models:
+            logger.warning(f"⚠️ No translation model available for {lang_pair}")
+            return jsonify({"error": f"Translation from {source_language} to {target_language} is not supported yet"}), 400
+        
+        if translation_models[lang_pair] is None or translation_tokenizers[lang_pair] is None:
+            logger.error(f"❌ Translation model for {lang_pair} not loaded")
+            return jsonify({"error": f"Translation model not available"}), 503
+            
         try:
+            # Get the appropriate model and tokenizer
+            model = translation_models[lang_pair]
+            tokenizer = translation_tokenizers[lang_pair]
+            
             # Tokenize the text
-            tokenized = translation_tokenizer(source_text, return_tensors="pt", padding=True)
+            tokenized = tokenizer(source_text, return_tensors="pt", padding=True)
             tokenized = {k: v.to(device) for k, v in tokenized.items()}
             
             # Generate translation
             with torch.no_grad():
-                translated = translation_model.generate(**tokenized)
+                translated = model.generate(**tokenized)
             
             # Decode the translation
-            result = translation_tokenizer.decode(translated[0], skip_special_tokens=True)
+            result = tokenizer.decode(translated[0], skip_special_tokens=True)
             
             logger.info(f"✅ Translation result: '{result}'")
             
@@ -399,7 +427,6 @@ def translate_text():
         logger.debug(f"Stack trace: {traceback.format_exc()}")
         return jsonify({"error": f"Internal server error: {str(e)}"}), 500
 
-
 if __name__ == "__main__":
     logger.info("🚀 Starting Speech API server")
     logger.info(f"📊 System status: ASR model: {'✅' if asr_model else '❌'}")