Update app.py

app.py

@@ -40,6 +40,7 @@ try:
     from flask import Flask, request, jsonify, send_file
     from flask_cors import CORS
     from transformers import Wav2Vec2ForCTC, AutoProcessor, VitsModel, AutoTokenizer
+    from transformers import MarianMTModel, MarianTokenizer
     logger.info("✅ All required libraries imported successfully")
 except ImportError as e:
     logger.critical(f"❌ Failed to import necessary libraries: {str(e)}")
@@ -120,6 +121,31 @@ 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 = "Helsinki-NLP/opus-mt-tc-bible-big-phi-en"
+logger.info(f"🔄 Loading Translation model: {TRANSLATION_MODEL_ID}")
+
+# Initialize translation model and tokenizer (add this after your other model initializations)
+translation_model = None
+translation_tokenizer = None
+
+try:
+    translation_tokenizer = MarianTokenizer.from_pretrained(
+        TRANSLATION_MODEL_ID,
+        cache_dir=cache_dirs["TRANSFORMERS_CACHE"]
+    )
+    logger.info("✅ Translation tokenizer loaded successfully")
+    
+    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()}")
+
 # Constants
 SAMPLE_RATE = 16000
 OUTPUT_DIR = "/tmp/audio_outputs"
@@ -140,6 +166,7 @@ def health_check():
         "asr_model": "loaded" if asr_model is not None else "failed",
         "tts_models": {lang: "loaded" if model is not None else "failed" 
                       for lang, model in tts_models.items()},
+        "translation_model": "loaded" if translation_model is not None else "failed",
         "device": device
     }
     return jsonify(health_status)
@@ -321,6 +348,57 @@ def download_audio(filename):
     logger.warning(f"⚠️ Requested file not found: {file_path}")
     return jsonify({"error": "File not found"}), 404
 
+@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:
+            logger.warning("⚠️ Translation endpoint called with no JSON data")
+            return jsonify({"error": "No JSON data provided"}), 400
+            
+        source_text = data.get("text", "").strip()
+        source_language = data.get("source_language", "").lower()
+        target_language = data.get("target_language", "").lower()
+
+        if not source_text:
+            logger.warning("⚠️ Translation request with empty text")
+            return jsonify({"error": "No text provided"}), 400
+            
+        logger.info(f"🔄 Translating from {source_language} to {target_language}: '{source_text}'")
+        
+        try:
+            # Tokenize the text
+            tokenized = translation_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)
+            
+            # Decode the translation
+            result = translation_tokenizer.decode(translated[0], skip_special_tokens=True)
+            
+            logger.info(f"✅ Translation result: '{result}'")
+            
+            return jsonify({
+                "translated_text": result,
+                "source_language": source_language,
+                "target_language": target_language
+            })
+        except Exception as e:
+            logger.error(f"❌ Translation processing failed: {str(e)}")
+            logger.debug(f"Stack trace: {traceback.format_exc()}")
+            return jsonify({"error": f"Translation processing failed: {str(e)}"}), 500
+            
+    except Exception as e:
+        logger.error(f"❌ Unhandled exception in translation endpoint: {str(e)}")
+        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")