Update app.py

app.py

@@ -23,6 +23,9 @@ import torch
 from phonemizer import phonemize
 from faster_whisper import WhisperModel
 from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
+import edge_tts
+import asyncio
+import base64
 
 app = FastAPI()
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
@@ -32,8 +35,14 @@ phoneme_processor = Wav2Vec2Processor.from_pretrained("vitouphy/wav2vec2-xls-r-3
 phoneme_model = Wav2Vec2ForCTC.from_pretrained("vitouphy/wav2vec2-xls-r-300m-timit-phoneme")
 whisper_model = WhisperModel("small", compute_type="float32")
 
-# Cache for phoneme lookups
+# Cache for phoneme lookups and TTS audio
 phoneme_cache = {}
+tts_cache = {}
+
+# TTS configuration
+TTS_VOICE = "en-US-AriaNeural"  # High quality English voice
+TTS_RATE = "+0%"  # Normal speed
+TTS_PITCH = "+0Hz"  # Normal pitch
 
 def log(msg):
     print(f"[{datetime.now().strftime('%H:%M:%S')}] {msg}")
@@ -55,6 +64,75 @@ def words_sim(a: str, b: str) -> float:
     """Cached similarity calculation"""
     return SequenceMatcher(None, a.lower(), b.lower()).ratio()
 
+async def generate_expected_audio(word: str) -> str:
+    """Generate TTS audio for expected pronunciation and return as base64"""
+    # Check cache first
+    cache_key = f"{word.lower()}_{TTS_VOICE}_{TTS_RATE}_{TTS_PITCH}"
+    if cache_key in tts_cache:
+        log(f"TTS cache hit for: {word}")
+        return tts_cache[cache_key]
+    
+    log(f"Generating TTS for: {word}")
+    
+    try:
+        # Generate TTS audio
+        communicate = edge_tts.Communicate(word, TTS_VOICE, rate=TTS_RATE, pitch=TTS_PITCH)
+        
+        # Collect audio data
+        audio_data = b""
+        async for chunk in communicate.stream():
+            if chunk["type"] == "audio":
+                audio_data += chunk["data"]
+        
+        if not audio_data:
+            log(f"No audio data generated for: {word}")
+            return ""
+        
+        # Convert to base64
+        audio_b64 = base64.b64encode(audio_data).decode('utf-8')
+        
+        # Cache the result
+        tts_cache[cache_key] = audio_b64
+        log(f"TTS generated and cached for: {word} ({len(audio_data)} bytes)")
+        
+        return audio_b64
+        
+    except Exception as e:
+        log(f"TTS generation failed for '{word}': {str(e)}")
+        return ""
+
+def extract_user_audio_segment(waveform: torch.Tensor, sample_rate: int, 
+                              start_time: float, end_time: float) -> str:
+    """Extract user's pronunciation segment and return as base64 WAV"""
+    try:
+        # Add small buffer around the word
+        buffer_time = 0.05  # 50ms buffer
+        adj_start = max(0, start_time - buffer_time)
+        adj_end = end_time + buffer_time
+        
+        # Convert to sample indices
+        start_sample = int(adj_start * sample_rate)
+        end_sample = int(adj_end * sample_rate)
+        end_sample = min(waveform.shape[-1], end_sample)
+        
+        # Extract segment
+        segment = waveform[:, start_sample:end_sample]
+        
+        # Convert to bytes
+        buffer = io.BytesIO()
+        torchaudio.save(buffer, segment, sample_rate, format="wav")
+        buffer.seek(0)
+        
+        # Convert to base64
+        audio_b64 = base64.b64encode(buffer.read()).decode('utf-8')
+        
+        log(f"Extracted user audio segment: {start_time:.2f}s-{end_time:.2f}s ({len(audio_b64)} chars)")
+        return audio_b64
+        
+    except Exception as e:
+        log(f"Failed to extract audio segment {start_time}-{end_time}: {str(e)}")
+        return ""
+
 def get_reference_phonemes(words: List[str]) -> List[str]:
     """Get reference phonemes for all words at once with caching"""
     cache_key = tuple(words)
@@ -316,6 +394,10 @@ async def transcribe(audio: UploadFile = File(...), similarity: float = Form(0.4
     # Find best phoneme matches considering all variants
     scores = find_best_phoneme_matches(segment_variants, reference_phonemes, sample_rate)
     
+    # Generate audio data for playback
+    log("Generating audio data for playback...")
+    audio_data = await generate_audio_data(words, word_texts, waveform, sample_rate)
+    
     # Format output
     full_text = " ".join(word_texts)
     resolved_output = []
@@ -336,18 +418,58 @@ async def transcribe(audio: UploadFile = File(...), similarity: float = Form(0.4
     return {
         "transcript": full_text,
         "resolved": " ".join(resolved_output),
-        "resolved_colored": " ".join(resolved_colored)
+        "resolved_colored": " ".join(resolved_colored),
+        "audio_data": audio_data,
+        "debug_info": {
+            "total_words": len(words),
+            "audio_segments_generated": len([a for a in audio_data if a["user_audio"]]),
+            "tts_segments_generated": len([a for a in audio_data if a["expected_audio"]]),
+            "cache_stats": {
+                "phoneme_cache_size": len(phoneme_cache),
+                "tts_cache_size": len(tts_cache)
+            }
+        }
     }
 
 @app.get("/")
 def root():
     return "fonetik running (optimized)"
 
-# Optional: Add an endpoint to clear caches if needed
+# Optional: Add endpoints for debugging and cache management
 @app.post("/api/clear-cache")
 def clear_cache():
-    global phoneme_cache
+    global phoneme_cache, tts_cache
     phoneme_cache.clear()
+    tts_cache.clear()
     normalize_phoneme_string.cache_clear()
     words_sim.cache_clear()
-    return {"message": "Cache cleared"}
+    return {"message": "All caches cleared"}
+
+@app.get("/api/debug/cache-stats")
+def get_cache_stats():
+    return {
+        "phoneme_cache_size": len(phoneme_cache),
+        "tts_cache_size": len(tts_cache),
+        "lru_cache_info": {
+            "normalize_phoneme_string": normalize_phoneme_string.cache_info()._asdict(),
+            "words_sim": words_sim.cache_info()._asdict()
+        }
+    }
+
+@app.post("/api/debug/test-tts")
+async def test_tts(word: str = Form(...)):
+    """Test TTS generation for a single word"""
+    try:
+        audio_b64 = await generate_expected_audio(word)
+        return {
+            "word": word,
+            "success": len(audio_b64) > 0,
+            "audio_length": len(audio_b64),
+            "audio_data": audio_b64 if len(audio_b64) > 0 else None
+        }
+    except Exception as e:
+        return {
+            "word": word,
+            "success": False,
+            "error": str(e)
+        }