Upload ml/model/engine.py with huggingface_hub

ml/model/engine.py

@@ -1,72 +1,62 @@
 """
-ml/model/engine.py - Production-Grade Unified Speech Diagnostics Engine
-========================================================================
-Integrates and optimizes all speech pathology subsystems into a fast,
-stable, singleton pipeline:
-  1. Neural LoRA Stutter & Disfluency Classifier (wav2vec2-base)
-  2. Neural ASR & GOP Phonetic Alignment (wav2vec2-base-960h CTC)
-  3. High-Precision Praat Phonation & Articulation Acoustics
-  4. Explicit Sound Disorder Rules ('R' Rhotacism, 'S' Sigmatism / Lisping)
-  5. Multi-Modal Decision Fusion with Self-Calibration & Silence Guard
+ml/model/engine.py - Complete Multi-Modal Speech Diagnostic Engine
+===================================================================
+Orchestrates:
+1. Signal conditioning & raw-signal Multi-Factor VAD
+2. Acoustic & Voice Phonation Analysis (Praat PointProcess)
+3. Neural ASR & DP Phonetic Alignment (wav2vec2-CTC)
+4. Neural Disfluency Detection (wav2vec2 + LoRA)
+5. Multi-Modal Decision Fusion, Concern Bands & Confidence Margins
 """
 from __future__ import annotations
-import difflib
 import gc
 import json
 import os
-import string
 import time
 from pathlib import Path
-from typing import Optional, Dict, Any, List, Tuple, Generator
+from typing import Any, Dict, Generator, List, Optional
 
 import numpy as np
 import torch
 from peft import PeftModel
 from transformers import (
     Wav2Vec2FeatureExtractor,
-    Wav2Vec2ForSequenceClassification,
     Wav2Vec2ForCTC,
+    Wav2Vec2ForSequenceClassification,
     Wav2Vec2Processor,
 )
 
-from ml.model import pron_eval, fusion
-from ml.model.stutter_trainer import SR, MAX_SECONDS, MODEL_BASE, ID2LABEL, BIN_ID2LABEL
+from ml.model import fusion, pron_eval
 
+SR = 16000
+MAX_SECONDS = 6.0
 CKPT_PATH = "ml/models/stutter/stutter_lora"
+CLASS_MAP_PATH = "ml/models/stutter/class_map.json"
+MODEL_BASE = "facebook/wav2vec2-base"
 CTC_MODEL_NAME = "facebook/wav2vec2-base-960h"
 
-# Constrain PyTorch thread overhead for low-RAM CPU environments (Render / Cloud)
-if not torch.cuda.is_available():
-    torch.set_num_threads(min(2, os.cpu_count() or 1))
-    torch.set_num_interop_threads(1)
-
 
 class SpeechDiagnosticEngine:
-    """Singleton, thread-safe, high-speed speech diagnostic pipeline."""
-
     _instance: Optional[SpeechDiagnosticEngine] = None
 
     def __init__(self, ckpt_dir: str = CKPT_PATH, device: Optional[str] = None):
         self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
         print(f"[SpeechDiagnosticEngine] Initializing on device: {self.device}")
 
-        # 1. Load Neural Stutter Model
-        ckpt = Path(ckpt_dir)
+        # 1. Load LoRA Stutter Classifier
         self.stutter_model = None
         self.stutter_feat = None
-        self.id2label = BIN_ID2LABEL
-
+        self.id2label = {0: "fluent", 1: "stutter"}
+        
+        ckpt = Path(ckpt_dir)
         if ckpt.exists():
             try:
-                cm_path = ckpt.parent / "class_map.json"
-                binary = True
-                if cm_path.exists():
-                    try:
-                        binary = bool(json.loads(cm_path.read_text(encoding="utf-8")).get("binary", True))
-                    except Exception:
-                        binary = True
-                self.id2label = BIN_ID2LABEL if binary else ID2LABEL
-
+                cm_file = ckpt.parent / "class_map.json"
+                if cm_file.exists():
+                    with open(cm_file, "r", encoding="utf-8") as f:
+                        data = json.load(f)
+                        self.id2label = {int(k): v for k, v in data.get("id2label", {}).items()}
+                
                 base = Wav2Vec2ForSequenceClassification.from_pretrained(
                     MODEL_BASE, num_labels=len(self.id2label), ignore_mismatched_sizes=True
                 )
@@ -85,12 +75,11 @@ class SpeechDiagnosticEngine:
         self.asr_model.to(self.device)
         self.asr_model.eval()
         
-        # Clean memory allocation
         gc.collect()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
 
-        print("[SpeechDiagnosticEngine] Engine ready.")
+        print("[SpeechDiagnosticEngine] Engine initialized and ready.")
 
     @classmethod
     def get_instance(cls, ckpt_dir: str = CKPT_PATH) -> SpeechDiagnosticEngine:
@@ -101,9 +90,9 @@ class SpeechDiagnosticEngine:
 
     @torch.inference_mode()
     def transcribe_and_align(self, audio_input: Any, reference: str) -> dict:
-        """Perform neural ASR decoding and word alignment with human tolerance."""
-        arr = pron_eval._load_wave(audio_input)
-        if pron_eval.is_silent_or_empty(arr, SR):
+        """Perform neural ASR decoding and DP phonetic word alignment."""
+        raw_arr = pron_eval._load_wave(audio_input)
+        if pron_eval.is_silent_or_empty(raw_arr, SR):
             ref_norm = pron_eval._norm(reference)
             return {
                 "asr_hypothesis": "",
@@ -118,7 +107,9 @@ class SpeechDiagnosticEngine:
                 "length_warning": None,
             }
 
-        inp = self.asr_processor(arr, sampling_rate=SR, return_tensors="pt")
+        # Normalize audio strictly for neural model forward pass
+        norm_arr = pron_eval.normalize_for_neural_inference(raw_arr)
+        inp = self.asr_processor(norm_arr, sampling_rate=SR, return_tensors="pt")
         inp = {k: v.to(self.device) for k, v in inp.items()}
         logits = self.asr_model(**inp).logits
         pred_ids = torch.argmax(logits, dim=-1)
@@ -126,42 +117,20 @@ class SpeechDiagnosticEngine:
 
         ref = pron_eval._norm(reference)
         alignment = pron_eval.align_words(ref, hypothesis)
-
-        ref_words = ref.split()
-        hyp_words = hypothesis.split()
-        n_ref = max(len(ref_words), 1)
-        n_hyp = max(len(hyp_words), 1)
+        wer, pron_score, counts = pron_eval.compute_standard_wer_and_pron_score(ref, hypothesis, alignment)
 
         length_warning = None
-        if len(hyp_words) == 1 and len(ref_words) >= 4:
-            length_warning = f"You spoke 1 word ('{hypothesis}'), but the target sentence has {len(ref_words)} words."
-
-        correct_count = sum(1 for a in alignment if a["status"] == "correct")
-
-        if len(hyp_words) < len(ref_words) and len(hyp_words) > 0:
-            spoken_precision = correct_count / n_hyp
-            matched_targets = " ".join([a["expected"] for a in alignment if a["spoken"] != "—"])
-            char_sim = difflib.SequenceMatcher(None, hypothesis, matched_targets).ratio()
-            pron_score = 0.70 * spoken_precision + 0.30 * char_sim
-        else:
-            word_acc = correct_count / n_ref
-            char_acc = difflib.SequenceMatcher(None, ref, hypothesis).ratio()
-            pron_score = 0.75 * word_acc + 0.25 * char_acc
-
-        if correct_count == n_ref or (len(hyp_words) == 1 and len(ref_words) == 1 and hyp_words[0] == ref_words[0]):
-            pron_score = 1.0
-
-        errors = sum(1 for a in alignment if a["status"] != "correct")
-        wer = min(1.0, errors / n_ref)
+        if len(hypothesis.split()) == 1 and len(ref.split()) >= 4:
+            length_warning = f"Only 1 word ('{hypothesis}') was recognized out of {len(ref.split())} target words."
 
         return {
             "asr_hypothesis": hypothesis,
             "reference_normalized": ref,
-            "word_error": errors,
-            "n_reference_words": n_ref,
-            "wer": round(wer, 4),
-            "goodness": round(pron_score, 4),
-            "pron_score": round(pron_score, 4),
+            "word_error": counts["substitutions"] + counts["deletions"] + counts["insertions"],
+            "n_reference_words": counts["n_ref"],
+            "wer": wer,
+            "goodness": pron_score,
+            "pron_score": pron_score,
             "alignment": alignment,
             "is_silent": False,
             "length_warning": length_warning,
@@ -173,11 +142,12 @@ class SpeechDiagnosticEngine:
         if self.stutter_model is None:
             return None
 
-        arr = pron_eval._load_wave(audio_input)
-        if pron_eval.is_silent_or_empty(arr, SR):
+        raw_arr = pron_eval._load_wave(audio_input)
+        if pron_eval.is_silent_or_empty(raw_arr, SR):
             return None
 
-        arr_clipped = arr[: int(SR * MAX_SECONDS)]
+        norm_arr = pron_eval.normalize_for_neural_inference(raw_arr)
+        arr_clipped = norm_arr[: int(SR * MAX_SECONDS)]
         inp = self.stutter_feat(arr_clipped, sampling_rate=SR, return_tensors="pt", padding=True)
         inp = {k: v.to(self.device) for k, v in inp.items()}
         logits = self.stutter_model(**inp).logits
@@ -192,23 +162,23 @@ class SpeechDiagnosticEngine:
         """Complete, unified diagnostic pipeline with latency timing."""
         t0 = time.perf_counter()
 
-        # 1. Load and condition audio
-        arr = pron_eval._load_wave(audio_input)
-        is_silent = pron_eval.is_silent_or_empty(arr, SR)
+        # 1. Load raw audio and check VAD
+        raw_arr = pron_eval._load_wave(audio_input)
+        is_silent = pron_eval.is_silent_or_empty(raw_arr, SR)
 
-        # 2. Articulation Acoustics (Praat)
-        artic = pron_eval.praat_metrics_arr(arr, SR)
+        # 2. Voice Acoustics & Phonation (Praat on raw waveform)
+        artic = pron_eval.praat_metrics_arr(raw_arr, SR)
 
-        # 3. Neural ASR Pronunciation & Word Alignment
-        pron = self.transcribe_and_align(arr, target_phrase)
+        # 3. Neural ASR & DP Alignment
+        pron = self.transcribe_and_align(raw_arr, target_phrase)
 
-        # 4. Neural Stutter Classification
-        stut_probs = self.predict_stutter_probs(arr) if not is_silent else None
+        # 4. Neural Disfluency Detection
+        stut_probs = self.predict_stutter_probs(raw_arr) if not is_silent else None
         p_stut = float(stut_probs[1]) if (stut_probs and len(stut_probs) > 1) else (
             float(np.sum(stut_probs[1:])) if stut_probs else 0.0
         )
 
-        # 5. Sound Disorder Analysis ('R', 'S', Substitutions, Phonation)
+        # 5. Acoustic-Phonetic Flaw Rules
         flaws = pron_eval.analyze_speech_flaws(
             reference=target_phrase,
             hypothesis=pron.get("asr_hypothesis", ""),
@@ -218,7 +188,7 @@ class SpeechDiagnosticEngine:
             is_silent=is_silent,
         )
 
-        # 6. Self-Calibration (if provided)
+        # 6. Baseline Normalization ("My Normal")
         cal = None
         if normal_calibration_audio is not None and self.stutter_model is not None:
             norm_probs = self.predict_stutter_probs(normal_calibration_audio)
@@ -227,7 +197,6 @@ class SpeechDiagnosticEngine:
 
         # 7. Multi-Modal Decision Fusion
         decision = fusion.diag_statistics(stut_probs, pron, artic, cal)
-
         latency_ms = round((time.perf_counter() - t0) * 1000, 1)
 
         return {
@@ -238,7 +207,8 @@ class SpeechDiagnosticEngine:
             "articulation": artic,
             "stutter_probs": stut_probs,
             "latency_ms": latency_ms,
-            "duration_s": round(len(arr) / SR, 2),
+            "duration_s": round(len(raw_arr) / SR, 2),
+            "device": self.device,
         }
 
     def diagnose_audio_stream(
@@ -250,49 +220,51 @@ class SpeechDiagnosticEngine:
         """Streaming generator yielding step-by-step progress and telemetry."""
         t0 = time.perf_counter()
 
-        # Step 1: Conditioning
+        # Step 1: Conditioning & Raw-Signal VAD
         t_step = time.perf_counter()
-        arr = pron_eval._load_wave(audio_input)
-        is_silent = pron_eval.is_silent_or_empty(arr, SR)
+        raw_arr = pron_eval._load_wave(audio_input)
+        is_silent = pron_eval.is_silent_or_empty(raw_arr, SR)
         t_s1 = round((time.perf_counter() - t_step) * 1000, 1)
         yield {
             "step": 1,
             "total": 5,
             "label": "Acoustic Signal Preconditioning",
-            "detail": f"16kHz PCM Resampling, 60Hz Butterworth High-Pass, Silence Check ({t_s1} ms)",
+            "detail": f"16kHz PCM Resampling, 60Hz High-Pass, Multi-Factor VAD ({t_s1} ms)",
             "progress": 0.20,
             "elapsed_ms": round((time.perf_counter() - t0) * 1000, 1),
         }
 
-        # Step 2: Praat Phonation
+        # Step 2: Phonation & Voice Analysis (Praat)
         t_step = time.perf_counter()
-        artic = pron_eval.praat_metrics_arr(arr, SR)
+        artic = pron_eval.praat_metrics_arr(raw_arr, SR)
         t_s2 = round((time.perf_counter() - t_step) * 1000, 1)
+        f0_val = f"{artic.get('f0_median_hz', 0):.1f}Hz" if artic.get('f0_median_hz') is not None else "N/A"
+        hnr_val = f"{artic.get('hnr_db', 0):.1f}dB" if artic.get('hnr_db') is not None else "N/A"
         yield {
             "step": 2,
             "total": 5,
-            "label": "Biomechanical Phonation Tracking",
-            "detail": f"Praat PointProcess Pitch F0={artic.get('f0_median_hz',0):.1f}Hz, Jitter={artic.get('jitter',0)*100:.2f}%, HNR={artic.get('hnr_db',0):.1f}dB ({t_s2} ms)",
+            "label": "Acoustic & Voice Phonation Analysis",
+            "detail": f"Praat PointProcess Pitch F0={f0_val}, HNR={hnr_val} ({t_s2} ms)",
             "progress": 0.40,
             "elapsed_ms": round((time.perf_counter() - t0) * 1000, 1),
         }
 
-        # Step 3: Neural CTC ASR Alignment
+        # Step 3: Neural ASR & DP Phonetic Alignment
         t_step = time.perf_counter()
-        pron = self.transcribe_and_align(arr, target_phrase)
+        pron = self.transcribe_and_align(raw_arr, target_phrase)
         t_s3 = round((time.perf_counter() - t_step) * 1000, 1)
         yield {
             "step": 3,
             "total": 5,
-            "label": "Neural ASR & Dynamic Alignment",
-            "detail": f"wav2vec2-CTC Transcribed: \"{pron.get('asr_hypothesis','')}\" | WER: {pron.get('wer',0)*100:.1f}% ({t_s3} ms)",
+            "label": "Neural ASR & Phonetic Alignment",
+            "detail": f"wav2vec2-CTC: \"{pron.get('asr_hypothesis','')}\" | WER: {pron.get('wer',0)*100:.1f}% ({t_s3} ms)",
             "progress": 0.60,
             "elapsed_ms": round((time.perf_counter() - t0) * 1000, 1),
         }
 
-        # Step 4: Neural Disfluency Classification
+        # Step 4: Neural Disfluency Detection
         t_step = time.perf_counter()
-        stut_probs = self.predict_stutter_probs(arr) if not is_silent else None
+        stut_probs = self.predict_stutter_probs(raw_arr) if not is_silent else None
         p_stut = float(stut_probs[1]) if (stut_probs and len(stut_probs) > 1) else (
             float(np.sum(stut_probs[1:])) if stut_probs else 0.0
         )
@@ -335,7 +307,8 @@ class SpeechDiagnosticEngine:
             "articulation": artic,
             "stutter_probs": stut_probs,
             "latency_ms": latency_ms,
-            "duration_s": round(len(arr) / SR, 2),
+            "duration_s": round(len(raw_arr) / SR, 2),
+            "device": self.device,
             "step_timings_ms": {
                 "preconditioning": t_s1,
                 "phonation_praat": t_s2,
@@ -349,7 +322,7 @@ class SpeechDiagnosticEngine:
             "step": 5,
             "total": 5,
             "label": "Multi-Modal Decision Fusion & Clinical Report",
-            "detail": f"Fluency Index: {decision.get('fluency_100', 0)}/100 | Severity: {decision['buckets']['overall'].upper()} ({t_s5} ms)",
+            "detail": f"Screening Index: {decision.get('fluency_100', 0)}/100 | Concern Band: {decision['buckets']['overall'].upper()} ({t_s5} ms)",
             "progress": 1.0,
             "elapsed_ms": latency_ms,
             "final_result": final_res,