Update translator.py

translator.py

@@ -288,9 +288,15 @@ def handle_tts_request(request, output_dir):
             logger.debug(f"Stack trace: {traceback.format_exc()}")
             return jsonify({"error": f"TTS inference failed: {str(e)}"}), 500
 
-        # Save to file
+        # Save to file with a unique name to prevent overwriting
         try:
-            output_filename = os.path.join(output_dir, f"{language}_output.wav")
+            # Create a unique filename using timestamp and text hash
+            import hashlib
+            import time
+            text_hash = hashlib.md5(text_input.encode()).hexdigest()[:8]
+            timestamp = int(time.time())
+            
+            output_filename = os.path.join(output_dir, f"{language}_{text_hash}_{timestamp}.wav")
             sampling_rate = model.config.sampling_rate
             sf.write(output_filename, waveform, sampling_rate)
             logger.info(f"✅ Speech generated! File saved: {output_filename}")
@@ -298,9 +304,10 @@ def handle_tts_request(request, output_dir):
             logger.error(f"❌ Failed to save audio file: {str(e)}")
             return jsonify({"error": f"Failed to save audio file: {str(e)}"}), 500
 
+        # Add cache-busting parameter to URL
         return jsonify({
             "message": "TTS audio generated",
-            "file_url": f"/download/{os.path.basename(output_filename)}",
+            "file_url": f"/download/{os.path.basename(output_filename)}?t={timestamp}",
             "language": language,
             "text_length": len(text_input)
         })
@@ -309,75 +316,134 @@ def handle_tts_request(request, output_dir):
         logger.debug(f"Stack trace: {traceback.format_exc()}")
         return jsonify({"error": f"Internal server error: {str(e)}"}), 500
 
-def handle_tts_request(request, output_dir):
-    """Handle TTS (Text-to-Speech) requests"""
+def handle_translation_request(request):
+    """Handle translation requests"""
     try:
         data = request.get_json()
         if not data:
-            logger.warning("⚠️ TTS endpoint called with no JSON data")
+            logger.warning("⚠️ Translation endpoint called with no JSON data")
             return jsonify({"error": "No JSON data provided"}), 400
 
-        text_input = data.get("text", "").strip()
-        language = data.get("language", "kapampangan").lower()
+        source_text = data.get("text", "").strip()
+        source_language = data.get("source_language", "").lower()
+        target_language = data.get("target_language", "").lower()
 
-        if not text_input:
-            logger.warning("⚠️ TTS request with empty text")
+        if not source_text:
+            logger.warning("⚠️ Translation request with empty text")
             return jsonify({"error": "No text provided"}), 400
 
-        if language not in TTS_MODELS:
-            logger.warning(f"⚠️ TTS requested for unsupported language: {language}")
-            return jsonify({"error": f"Invalid language. Available options: {list(TTS_MODELS.keys())}"}), 400
+        # Map language names to codes
+        source_code = LANGUAGE_CODES.get(source_language, source_language)
+        target_code = LANGUAGE_CODES.get(target_language, target_language)
 
-        if tts_models[language] is None:
-            logger.error(f"❌ TTS model for {language} not loaded")
-            return jsonify({"error": f"TTS model for {language} not available"}), 503
+        logger.info(f"🔄 Translating from {source_language} to {target_language}: '{source_text}'")
 
-        logger.info(f"🔄 Generating TTS for language: {language}, text: '{text_input}'")
+        # Special handling for pam-fil, fil-pam, pam-tgl and tgl-pam using the phi model
+        use_phi_model = False
+        actual_source_code = source_code
+        actual_target_code = target_code
 
-        try:
-            processor = tts_processors[language]
-            model = tts_models[language]
-            inputs = processor(text_input, return_tensors="pt")
-            inputs = {k: v.to(model.device) for k, v in inputs.items()}
-        except Exception as e:
-            logger.error(f"❌ TTS preprocessing failed: {str(e)}")
-            return jsonify({"error": f"TTS preprocessing failed: {str(e)}"}), 500
+        # Check if we need to use the phi model with fil replacement
+        if (source_code == "pam" and target_code == "fil") or (source_code == "fil" and target_code == "pam"):
+            use_phi_model = True
+        elif (source_code == "pam" and target_code == "tgl"):
+            use_phi_model = True
+            actual_target_code = "fil"  # Replace tgl with fil for the phi model
+        elif (source_code == "tgl" and target_code == "pam"):
+            use_phi_model = True
+            actual_source_code = "fil"  # Replace tgl with fil for the phi model
 
-        # Generate speech
-        try:
-            with torch.no_grad():
-                output = model(**inputs).waveform
-                waveform = output.squeeze().cpu().numpy()
-        except Exception as e:
-            logger.error(f"❌ TTS inference failed: {str(e)}")
-            logger.debug(f"Stack trace: {traceback.format_exc()}")
-            return jsonify({"error": f"TTS inference failed: {str(e)}"}), 500
+        if use_phi_model:
+            model_key = "phi"
 
-        # Save to file with a unique name to prevent overwriting
-        try:
-            # Create a unique filename using timestamp and text hash
-            import hashlib
-            import time
-            text_hash = hashlib.md5(text_input.encode()).hexdigest()[:8]
-            timestamp = int(time.time())
-            
-            output_filename = os.path.join(output_dir, f"{language}_{text_hash}_{timestamp}.wav")
-            sampling_rate = model.config.sampling_rate
-            sf.write(output_filename, waveform, sampling_rate)
-            logger.info(f"✅ Speech generated! File saved: {output_filename}")
-        except Exception as e:
-            logger.error(f"❌ Failed to save audio file: {str(e)}")
-            return jsonify({"error": f"Failed to save audio file: {str(e)}"}), 500
+            # Check if we have the phi model
+            if model_key not in translation_models or translation_models[model_key] is None:
+                logger.error(f"❌ Translation model for {model_key} not loaded")
+                return jsonify({"error": f"Translation model not available"}), 503
+
+            try:
+                # Get the phi model and tokenizer
+                model = translation_models[model_key]
+                tokenizer = translation_tokenizers[model_key]
+
+                # Prepend target language token to input
+                input_text = f">>{actual_target_code}<< {source_text}"
+
+                logger.info(f"🔄 Using phi model with input: '{input_text}'")
+
+                # Tokenize the text
+                tokenized = tokenizer(input_text, return_tensors="pt", padding=True)
+                tokenized = {k: v.to(model.device) for k, v in tokenized.items()}
+
+                with torch.no_grad():
+                    translated = model.generate(
+                        **tokenized,
+                        max_length=100,              # Reasonable output length
+                        num_beams=4,                 # Same as in training
+                        length_penalty=0.6,          # Same as in training
+                        early_stopping=True,         # Same as in training
+                        repetition_penalty=1.5,      # Add this to prevent repetition
+                        no_repeat_ngram_size=3       # Add this to prevent repetition
+                    )
+
+                # Decode the translation
+                result = 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
+        else:
+            # Create the regular language pair key for other language pairs
+            lang_pair = f"{source_code}-{target_code}"
+
+            # 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:
+                # Regular translation process for other language pairs
+                model = translation_models[lang_pair]
+                tokenizer = translation_tokenizers[lang_pair]
+
+                # Tokenize the text
+                tokenized = tokenizer(source_text, return_tensors="pt", padding=True)
+                tokenized = {k: v.to(model.device) for k, v in tokenized.items()}
+
+                # Generate translation
+                with torch.no_grad():
+                    translated = model.generate(**tokenized)
+
+                # Decode the translation
+                result = 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
 
-        # Add cache-busting parameter to URL
-        return jsonify({
-            "message": "TTS audio generated",
-            "file_url": f"/download/{os.path.basename(output_filename)}?t={timestamp}",
-            "language": language,
-            "text_length": len(text_input)
-        })
     except Exception as e:
-        logger.error(f"❌ Unhandled exception in TTS endpoint: {str(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