New: Phoneme-level speech pathology diagnosis MVP with real-time streaming

- Add Wav2Vec2-XLSR-53 based speech pathology classifier with 8-class output (fluency + articulation)
- Implement phone-level feature extraction with 1-second sliding windows and 10ms hops
- Add grapheme-to-phoneme mapping using g2p_en library with frame alignment
- Create error taxonomy system with substitution/omission/distortion detection and therapy recommendations
- Build FastAPI REST API with batch diagnosis endpoints (/diagnose/file) and WebSocket streaming (/ws/diagnose)
- Add Gradio web interface with audio upload/recording, real-time error display, and detailed reports
- Implement training infrastructure: synthetic data generation, classifier head training, and evaluation scripts
- Add Docker containerization with NLTK data download for phoneme mapping
- Fix model loading issues and improve error handling throughout the pipeline
- Remove unnecessary package-lock.json file (Python project)

Features:
- Real-time streaming analysis (<200ms per file, <50ms per frame target)
- Phoneme-level error detection with visual feedback
- Therapy recommendation system with clinical guidance
- Comprehensive error reporting with severity levels and timelines
- WebSocket-based real-time diagnosis for live audio streams
- REST API for batch processing with detailed JSON responses

Infrastructure:
- Training pipeline for classifier fine-tuning
- Data collection tools for phoneme-level annotation
- Performance testing and integration testing suites
- Production-ready logging and error handling

.gitignore

@@ -7,5 +7,5 @@ __pycache__/
 
 .gradio/
 .cursor/
-package-lock.json
+# package-lock.json
 docker-compose.yml

Dockerfile

@@ -27,6 +27,9 @@ RUN pip install --no-cache-dir \
 # Install the rest of requirements
 RUN pip install --no-cache-dir -r requirements.txt
 
+# Download NLTK data required by g2p_en
+RUN python -c "import nltk; nltk.download('averaged_perceptron_tagger_eng', quiet=True)"
+
 # Copy application files
 COPY . .
 

README_TRAINING.md

@@ -0,0 +1,176 @@
+# Training Guide for Speech Pathology Classifier
+
+This guide explains how to train the classifier head for phoneme-level speech pathology detection.
+
+## Overview
+
+The system uses Wav2Vec2-XLSR-53 as a frozen feature extractor and trains only the classification head (2-3 layer feedforward network) on phoneme-level labeled data.
+
+## Prerequisites
+
+1. **Labeled Data**: 50-100 audio samples with phoneme-level error annotations
+2. **Python Environment**: Python 3.10+ with required dependencies
+3. **GPU** (recommended): For faster training
+
+## Step 1: Data Collection
+
+### Using the Annotation Tool
+
+1. Launch the data collection interface:
+```bash
+python scripts/data_collection.py
+```
+
+2. The Gradio interface will open at `http://localhost:7861`
+
+3. For each sample:
+   - Upload or record audio (5-30 seconds, 16kHz WAV)
+   - Enter expected text/transcript
+   - Extract phonemes (automatic G2P conversion)
+   - Annotate errors at phoneme level:
+     - Frame ID where error occurs
+     - Phoneme with error
+     - Error type (substitution/omission/distortion/stutter)
+     - Wrong sound (for substitutions)
+     - Severity (0-1)
+     - Timestamp
+   - Add notes if needed
+   - Save annotation
+
+4. Annotations are saved to:
+   - Audio files: `data/raw/`
+   - Annotations: `data/annotations.json`
+
+### Export Training Data
+
+After collecting annotations, export for training:
+
+```bash
+python scripts/annotation_helper.py
+```
+
+This creates `data/training_dataset.json` with frame-level labels.
+
+## Step 2: Training
+
+### Configuration
+
+Edit `training/config.yaml` to adjust hyperparameters:
+
+- `batch_size`: 16 (adjust based on GPU memory)
+- `learning_rate`: 0.001
+- `num_epochs`: 50
+- `train_split`: 0.8 (80% for training, 20% for validation)
+
+### Run Training
+
+```bash
+python training/train_classifier_head.py --config training/config.yaml
+```
+
+Training will:
+- Load training dataset
+- Extract Wav2Vec2 features for each sample
+- Train classifier head (Wav2Vec2 frozen)
+- Save best checkpoint to `models/checkpoints/classifier_head_best.pt`
+- Save last checkpoint to `models/checkpoints/classifier_head_trained.pt`
+
+### Monitor Training
+
+Training logs include:
+- Loss per epoch
+- Accuracy per epoch
+- Validation metrics
+- Best model checkpoint saves
+
+## Step 3: Evaluation
+
+Evaluate the trained model:
+
+```bash
+python training/evaluate_classifier.py \
+    --checkpoint models/checkpoints/classifier_head_best.pt \
+    --dataset data/training_dataset.json \
+    --output training/evaluation_results.json \
+    --plot training/confusion_matrix.png
+```
+
+This generates:
+- Overall accuracy, F1 score, precision, recall
+- Per-class accuracy
+- Confusion matrix (saved as PNG)
+- Confidence analysis
+- Detailed metrics JSON
+
+## Step 4: Deployment
+
+Once trained, the model automatically loads trained weights on startup:
+
+1. Place checkpoint in `models/checkpoints/classifier_head_best.pt`
+2. Restart the application
+3. The model will detect and load trained weights automatically
+
+### Verify Training Status
+
+Check API responses for:
+- `model_version`: "wav2vec2-xlsr-53-v2-trained" (if trained) or "wav2vec2-xlsr-53-v2-beta" (if untrained)
+- `model_trained`: true/false
+- `confidence_filter_threshold`: 0.65
+
+## Troubleshooting
+
+### Issue: "No training dataset found"
+
+**Solution**: Run `scripts/annotation_helper.py` to export training data from annotations.
+
+### Issue: "CUDA out of memory"
+
+**Solution**: Reduce `batch_size` in `training/config.yaml` (try 8 or 4).
+
+### Issue: "Poor validation accuracy"
+
+**Solutions**:
+- Collect more training data (aim for 100+ samples)
+- Check data quality (ensure accurate annotations)
+- Adjust learning rate or add data augmentation
+- Use class weights for imbalanced data
+
+### Issue: "Model not loading trained weights"
+
+**Solution**: 
+- Verify checkpoint path: `models/checkpoints/classifier_head_best.pt`
+- Check file permissions
+- Review logs for loading errors
+
+## Best Practices
+
+1. **Data Quality > Quantity**: 50 high-quality samples > 100 poor samples
+2. **Balanced Classes**: Ensure all 8 classes have sufficient examples
+3. **Validation Split**: Use 20% for validation, never train on test data
+4. **Early Stopping**: Enabled by default to prevent overfitting
+5. **Class Weights**: Automatically calculated to handle imbalance
+6. **Checkpointing**: Best model saved automatically
+
+## Expected Results
+
+After training with 50-100 samples:
+- **Frame-level accuracy**: >75%
+- **Phoneme-level F1**: >85%
+- **Per-class accuracy**: >70% for each class
+- **Confidence**: Higher for correct predictions
+
+## Next Steps
+
+1. Collect more data based on error patterns
+2. Fine-tune hyperparameters
+3. Add data augmentation
+4. Deploy and monitor in production
+5. Retrain quarterly with new data
+
+## Support
+
+For issues or questions:
+- Check training logs in console
+- Review `training/evaluation_results.json`
+- Verify data format in `data/annotations.json`
+

api/routes.py

@@ -47,6 +47,16 @@ phoneme_mapper: Optional[PhonemeMapper] = None
 error_mapper: Optional[ErrorMapper] = None
 
 
+def get_phoneme_mapper() -> Optional[PhonemeMapper]:
+    """Get the global PhonemeMapper instance."""
+    return phoneme_mapper
+
+
+def get_error_mapper() -> Optional[ErrorMapper]:
+    """Get the global ErrorMapper instance."""
+    return error_mapper
+
+
 def initialize_routes(
     pipeline: InferencePipeline,
     mapper: Optional[PhonemeMapper] = None,
@@ -263,6 +273,10 @@ async def diagnose_file(
         processing_time_ms = (time.time() - start_time) * 1000
         
         # Create response
+        # Check if model is trained
+        model_trained = inference_pipeline.model.is_trained if hasattr(inference_pipeline.model, 'is_trained') else False
+        model_version = "wav2vec2-xlsr-53-v2-trained" if model_trained else "wav2vec2-xlsr-53-v2-beta"
+        
         response = BatchDiagnosisResponse(
             session_id=session_id,
             filename=audio.filename or "unknown",
@@ -274,7 +288,10 @@ async def diagnose_file(
             summary=summary,
             therapy_plan=therapy_plan,
             processing_time_ms=processing_time_ms,
-            created_at=datetime.now()
+            created_at=datetime.utcnow(),
+            model_version=model_version,
+            model_trained=model_trained,
+            confidence_filter_threshold=0.65
         )
         
         # Store in sessions

api/schemas.py

@@ -76,6 +76,9 @@ class BatchDiagnosisResponse(BaseModel):
     therapy_plan: List[str] = Field(default_factory=list, description="Therapy recommendations")
     processing_time_ms: float = Field(..., ge=0.0, description="Processing time in milliseconds")
     created_at: datetime = Field(default_factory=datetime.now, description="Analysis timestamp")
+    model_version: str = Field(default="wav2vec2-xlsr-53-v2", description="Model version identifier")
+    model_trained: bool = Field(default=False, description="Whether classifier head is trained")
+    confidence_filter_threshold: float = Field(default=0.65, ge=0.0, le=1.0, description="Confidence threshold for filtering predictions")
 
 
 class StreamingDiagnosisRequest(BaseModel):

app.py

@@ -7,7 +7,7 @@ from pathlib import Path
 from datetime import datetime
 from typing import Optional
 
-from fastapi import FastAPI, UploadFile, File, Form, HTTPException, WebSocket, WebSocketDisconnect
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException, WebSocket, WebSocketDisconnect, Query
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 import gradio as gr
@@ -26,8 +26,7 @@ sys.path.insert(0, str(Path(__file__).parent))
 # Import model loaders and inference pipeline
 try:
     from diagnosis.ai_engine.model_loader import (
-        get_stutter_detector,  # Legacy detector
-        get_inference_pipeline  # New inference pipeline
+        get_inference_pipeline  # Wav2Vec2-based inference pipeline
     )
     from ui.gradio_interface import create_gradio_interface
     from config import APIConfig, GradioConfig, default_api_config, default_gradio_config
@@ -53,27 +52,30 @@ app.add_middleware(
 )
 
 # Global instances
-detector = None  # Legacy detector
-inference_pipeline = None  # New inference pipeline
+inference_pipeline = None  # Wav2Vec2-based inference pipeline
 
 @app.on_event("startup")
 async def startup_event():
     """Load models on startup"""
-    global detector, inference_pipeline
+    global inference_pipeline
     try:
         logger.info("🚀 Startup event: Loading AI models...")
         
-        # Load legacy detector (for backward compatibility)
-        try:
-            detector = get_stutter_detector()
-            logger.info("✅ Legacy detector loaded")
-        except Exception as e:
-            logger.warning(f"⚠️ Legacy detector not available: {e}")
-        
-        # Load new inference pipeline
+        # Load Wav2Vec2-based inference pipeline
         try:
             inference_pipeline = get_inference_pipeline()
             logger.info("✅ Inference pipeline loaded")
+            
+            # Initialize API routes with phoneme and error mappers
+            try:
+                from api.routes import initialize_routes
+                from api.streaming import initialize_streaming
+                initialize_routes(inference_pipeline)
+                initialize_streaming(inference_pipeline)
+                logger.info("✅ API routes initialized with phoneme/error mappers")
+            except Exception as e:
+                logger.warning(f"⚠️ API routes initialization failed: {e}", exc_info=True)
+                # Continue without phoneme mapping if it fails
         except Exception as e:
             logger.error(f"❌ Failed to load inference pipeline: {e}", exc_info=True)
             # Don't raise - allow API to start even if new pipeline fails
@@ -83,6 +85,24 @@ async def startup_event():
         logger.error(f"❌ Failed to load models: {e}", exc_info=True)
         raise
 
+# Include API routers
+try:
+    from api.routes import router as diagnose_router
+    app.include_router(diagnose_router)
+    logger.info("✅ Diagnosis router included")
+except Exception as e:
+    logger.warning(f"⚠️ Failed to include diagnosis router: {e}")
+
+# Add WebSocket endpoint
+try:
+    from api.streaming import handle_streaming_websocket
+    @app.websocket("/ws/diagnose")
+    async def websocket_diagnose(websocket: WebSocket, session_id: Optional[str] = None):
+        await handle_streaming_websocket(websocket, session_id)
+    logger.info("✅ WebSocket endpoint registered")
+except Exception as e:
+    logger.warning(f"⚠️ Failed to register WebSocket endpoint: {e}")
+
 # Create and mount new Gradio interface
 try:
     gradio_interface = create_gradio_interface(default_gradio_config)
@@ -103,31 +123,30 @@ async def health_check():
     return {
         "status": "healthy",
         "models_loaded": {
-            "legacy_detector": detector is not None,
-            "inference_pipeline": inference_pipeline is not None
+            "inference_pipeline": inference_pipeline is not None,
+            "model_version": "wav2vec2-xlsr-53-v2"
         },
         "timestamp": datetime.utcnow().isoformat() + "Z"
     }
 
 @app.post("/api/diagnose")
 async def diagnose_speech(
-    audio: UploadFile = File(...)
+    audio: UploadFile = File(...),
+    text: Optional[str] = Query(None, description="Expected text/transcript for phoneme mapping (optional)")
 ):
     """
-    Diagnose speech for fluency and articulation issues.
+    Legacy endpoint for speech diagnosis.
     
-    Uses the new Wav2Vec2-XLSR-53 inference pipeline for phone-level analysis.
+    NOTE: For full phoneme-level error detection with therapy recommendations, 
+    use POST /diagnose/file?text=<expected_text> instead.
+    This endpoint is maintained for backward compatibility.
     
     Parameters:
     - audio: Audio file (WAV, MP3, FLAC, M4A)
+    - text: Optional expected text for phoneme mapping
     
     Returns:
-        Dictionary with:
-        - status: "success" or "error"
-        - fluency_metrics: Fluency statistics
-        - articulation_results: Articulation analysis
-        - confidence: Overall confidence score
-        - processing_time_ms: Processing time in milliseconds
+        Dictionary with diagnosis results (legacy format for backward compatibility)
     """
     if not inference_pipeline:
         raise HTTPException(
@@ -135,11 +154,15 @@ async def diagnose_speech(
             detail="Inference pipeline not loaded yet. Try again in a moment."
         )
     
+    # Import here to avoid circular imports
+    from api.routes import get_phoneme_mapper, get_error_mapper
+    from models.error_taxonomy import ErrorType
+    
     start_time = time.time()
     temp_file = None
     
     try:
-        logger.info(f"📥 Processing diagnosis request: {audio.filename}")
+        logger.info(f"📥 Processing legacy diagnosis request: {audio.filename}")
         
         # Validate file extension
         file_ext = Path(audio.filename).suffix.lower()
@@ -173,7 +196,6 @@ async def diagnose_speech(
         
         # Run inference
         logger.info("🔄 Running inference pipeline...")
-        # Use new phone-level prediction
         result = inference_pipeline.predict_phone_level(
             temp_file,
             return_timestamps=True
@@ -181,28 +203,64 @@ async def diagnose_speech(
         
         processing_time_ms = (time.time() - start_time) * 1000
         
-        # Extract metrics from new PhoneLevelResult format
+        # Get mappers for phoneme/error processing
+        phoneme_mapper = get_phoneme_mapper()
+        error_mapper = get_error_mapper()
+        
+        # Map phonemes if text provided
+        frame_phonemes = []
+        errors = []
+        if text and phoneme_mapper and error_mapper:
+            try:
+                frame_phonemes = phoneme_mapper.map_text_to_frames(
+                    text,
+                    num_frames=result.num_frames,
+                    audio_duration=result.duration
+                )
+                
+                # Process errors
+                for i, frame_pred in enumerate(result.frame_predictions):
+                    phoneme = frame_phonemes[i] if i < len(frame_phonemes) else ''
+                    class_id = frame_pred.articulation_class
+                    if frame_pred.fluency_label == 'stutter':
+                        class_id += 4
+                    
+                    error_detail = error_mapper.map_classifier_output(
+                        class_id=class_id,
+                        confidence=frame_pred.confidence,
+                        phoneme=phoneme if phoneme else 'unknown',
+                        fluency_label=frame_pred.fluency_label
+                    )
+                    
+                    if error_detail.error_type != ErrorType.NORMAL:
+                        errors.append({
+                            "phoneme": error_detail.phoneme,
+                            "time": frame_pred.time,
+                            "error_type": error_detail.error_type.value,
+                            "wrong_sound": error_detail.wrong_sound,
+                            "severity": error_mapper.get_severity_level(error_detail.severity).value,
+                            "therapy": error_detail.therapy
+                        })
+            except Exception as e:
+                logger.warning(f"⚠️ Phoneme/error mapping failed: {e}")
+        
+        # Extract metrics
         aggregate = result.aggregate
         mean_fluency_stutter = aggregate.get("fluency_score", 0.0)
-        fluency_percentage = (1.0 - mean_fluency_stutter) * 100  # Convert stutter prob to fluency percentage
+        fluency_percentage = (1.0 - mean_fluency_stutter) * 100
         
-        # Count fluent frames
         fluent_frames = sum(1 for fp in result.frame_predictions if fp.fluency_label == 'normal')
         fluent_frames_ratio = fluent_frames / result.num_frames if result.num_frames > 0 else 0.0
         
-        # Extract articulation class distribution
         articulation_class_counts = {}
         for fp in result.frame_predictions:
             label = fp.articulation_label
             articulation_class_counts[label] = articulation_class_counts.get(label, 0) + 1
         
-        # Get dominant articulation class
         dominant_articulation = aggregate.get("articulation_label", "normal")
-        
-        # Calculate average confidence
         avg_confidence = sum(fp.confidence for fp in result.frame_predictions) / result.num_frames if result.num_frames > 0 else 0.0
         
-        # Format response
+        # Format response (legacy format with optional error info)
         response = {
             "status": "success",
             "fluency_metrics": {
@@ -225,9 +283,10 @@ async def diagnose_speech(
                         "fluency_label": fp.fluency_label,
                         "articulation_class": fp.articulation_class,
                         "articulation_label": fp.articulation_label,
-                        "confidence": fp.confidence
+                        "confidence": fp.confidence,
+                        "phoneme": frame_phonemes[i] if i < len(frame_phonemes) else ''
                     }
-                    for fp in result.frame_predictions
+                    for i, fp in enumerate(result.frame_predictions)
                 ]
             },
             "confidence": avg_confidence,
@@ -235,8 +294,14 @@ async def diagnose_speech(
             "processing_time_ms": processing_time_ms
         }
         
-        logger.info(f"✅ Diagnosis complete: fluency={response['fluency_metrics']['fluency_percentage']:.1f}%, "
-                   f"confidence={response['confidence_percentage']:.1f}%, "
+        # Add error info if available
+        if errors:
+            response["error_count"] = len(errors)
+            response["errors"] = errors[:10]  # Limit to first 10 for legacy format
+            response["problematic_sounds"] = list(set(err["phoneme"] for err in errors if err["phoneme"]))
+        
+        logger.info(f"✅ Legacy diagnosis complete: fluency={response['fluency_metrics']['fluency_percentage']:.1f}%, "
+                   f"errors={len(errors) if errors else 0}, "
                    f"time={processing_time_ms:.0f}ms")
         
         return response
@@ -257,70 +322,7 @@ async def diagnose_speech(
                 logger.warning(f"Could not clean up {temp_file}: {e}")
 
 
-@app.post("/analyze")
-async def analyze_audio(
-    audio: UploadFile = File(...),
-    transcript: str = Form("")
-):
-    """
-    Legacy endpoint: Analyze audio file for stuttering.
-    
-    Uses the legacy Whisper-based detector for backward compatibility.
-    
-    Parameters:
-    - audio: WAV or MP3 audio file
-    - transcript: Optional expected transcript
-    
-    Returns: Complete stutter analysis results
-    """
-    if not detector:
-        raise HTTPException(
-            status_code=503,
-            detail="Legacy detector not loaded. Use /api/diagnose for new analysis."
-        )
-    
-    temp_file = None
-    try:
-        logger.info(f"📥 Processing (legacy): {audio.filename}")
-        
-        # Create temp directory if needed
-        temp_dir = tempfile.gettempdir()
-        os.makedirs(temp_dir, exist_ok=True)
-        
-        # Save uploaded file
-        temp_file = os.path.join(temp_dir, f"legacy_{int(time.time())}_{audio.filename}")
-        content = await audio.read()
-        
-        with open(temp_file, "wb") as f:
-            f.write(content)
-        
-        logger.info(f"📂 Saved to: {temp_file} ({len(content) / 1024 / 1024:.2f} MB)")
-        
-        # Analyze
-        logger.info(f"🔄 Analyzing audio with transcript: '{transcript[:50] if transcript else '(empty)'}...'")
-        result = detector.analyze_audio(temp_file, transcript)
-        
-        actual = result.get('actual_transcript', '')
-        target = result.get('target_transcript', '')
-        logger.info(f"✅ Analysis complete: severity={result.get('severity', 'N/A')}, "
-                   f"mismatch={result.get('mismatch_percentage', 'N/A')}%")
-        
-        return result
-        
-    except HTTPException:
-        raise
-    except Exception as e:
-        logger.error(f"❌ Error during analysis: {str(e)}", exc_info=True)
-        raise HTTPException(status_code=500, detail=f"Analysis failed: {str(e)}")
-    
-    finally:
-        # Cleanup
-        if temp_file and os.path.exists(temp_file):
-            try:
-                os.remove(temp_file)
-                logger.debug(f"🧹 Cleaned up: {temp_file}")
-            except Exception as e:
-                logger.warning(f"Could not clean up {temp_file}: {e}")
+# Legacy /analyze endpoint removed - use /api/diagnose or /diagnose/file instead
 
 
 @app.websocket("/ws/audio")

inference/__init__.py

@@ -6,15 +6,15 @@ This package contains the inference pipeline for real-time and batch processing.
 
 from .inference_pipeline import (
     InferencePipeline,
-    PredictionResult,
-    BatchPredictionResult,
-    create_inference_pipeline
+    FramePrediction,
+    PhoneLevelResult,
+    create_inference_pipeline,
 )
 
 __all__ = [
     "InferencePipeline",
-    "PredictionResult",
-    "BatchPredictionResult",
+    "FramePrediction",
+    "PhoneLevelResult",
     "create_inference_pipeline",
 ]
 

models/phoneme_mapper.py

@@ -72,6 +72,18 @@ class PhonemeMapper:
                 "g2p_en library is required. Install with: pip install g2p-en"
             )
         
+        # Ensure NLTK data is available (required by g2p_en)
+        try:
+            import nltk
+            try:
+                nltk.data.find('taggers/averaged_perceptron_tagger_eng')
+            except LookupError:
+                logger.info("Downloading NLTK averaged_perceptron_tagger_eng...")
+                nltk.download('averaged_perceptron_tagger_eng', quiet=True)
+                logger.info("✅ NLTK data downloaded")
+        except Exception as e:
+            logger.warning(f"⚠️ Could not download NLTK data: {e}")
+        
         try:
             self.g2p = g2p_en.G2p()
             logger.info("✅ G2P model loaded successfully")

models/speech_pathology_model.py

@@ -210,6 +210,7 @@ class SpeechPathologyClassifier(nn.Module):
         
         self.device = torch.device(device)
         self.use_fp16 = use_fp16 and device == "cuda"
+        self.is_trained = False  # Track if classifier is trained
         
         if classifier_hidden_dims is None:
             classifier_hidden_dims = [256, 128]
@@ -257,6 +258,9 @@ class SpeechPathologyClassifier(nn.Module):
                 num_articulation_classes=num_articulation_classes
             )
             
+            # Try to load trained weights if available (None = try default paths)
+            self._load_trained_weights(None)
+            
             # Move to device
             self.wav2vec2_model = self.wav2vec2_model.to(self.device)
             self.classifier_head = self.classifier_head.to(self.device)
@@ -275,6 +279,56 @@ class SpeechPathologyClassifier(nn.Module):
             logger.error(f"❌ Failed to initialize model: {e}", exc_info=True)
             raise RuntimeError(f"Failed to load Wav2Vec2 model: {e}") from e
     
+    def _load_trained_weights(self, model_path: Optional[str] = None):
+        """
+        Load trained classifier head weights if available.
+        
+        Args:
+            model_path: Optional path to model checkpoint. If None, tries default checkpoint paths.
+        """
+        from pathlib import Path
+        
+        checkpoint_paths = []
+        
+        # Add user-provided path
+        if model_path:
+            checkpoint_paths.append(Path(model_path))
+        
+        # Add default checkpoint paths
+        checkpoint_paths.extend([
+            Path("models/checkpoints/classifier_head_best.pt"),
+            Path("models/checkpoints/classifier_head_trained.pt")
+        ])
+        
+        for checkpoint_path in checkpoint_paths:
+            if checkpoint_path.exists():
+                try:
+                    checkpoint = torch.load(checkpoint_path, map_location=self.device)
+                    
+                    # Handle both full checkpoint dict and state_dict directly
+                    if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
+                        state_dict = checkpoint['model_state_dict']
+                        epoch = checkpoint.get('epoch', 'unknown')
+                        val_acc = checkpoint.get('val_accuracy', 'unknown')
+                    else:
+                        state_dict = checkpoint
+                        epoch = 'unknown'
+                        val_acc = 'unknown'
+                    
+                    self.classifier_head.load_state_dict(state_dict)
+                    logger.info(f"✅ Loaded trained classifier head from {checkpoint_path}")
+                    logger.info(f"   Epoch: {epoch}, Validation Accuracy: {val_acc}")
+                    self.is_trained = True
+                    return
+                except Exception as e:
+                    logger.warning(f"⚠️ Could not load checkpoint {checkpoint_path}: {e}")
+                    continue
+        
+        # No trained weights found
+        logger.warning("⚠️ No trained classifier weights found. Using untrained head (beta mode)")
+        logger.warning("   To train the classifier, run: python training/train_classifier_head.py")
+        self.is_trained = False
+    
     def forward(
         self,
         input_values: torch.Tensor,

requirements.txt

@@ -5,6 +5,7 @@ torchaudio>=2.6.0
 transformers>=4.57.3,<5.0
 numpy>=1.24.0,<2.0.0
 protobuf>=3.20.0
+g2p-en>=2.1.0
 
 # Audio Processing
 librosa>=0.10.0
@@ -25,6 +26,12 @@ gradio==6.1.0
 # Logging
 python-json-logger>=2.0.0
 
+# Training dependencies
+pyyaml>=6.0
+scikit-learn>=1.3.0
+matplotlib>=3.7.0
+seaborn>=0.12.0
+
 # Optional: Legacy/Advanced features
 openai-whisper>=20230314
 praat-parselmouth>=0.4.3

scripts/annotation_helper.py

@@ -0,0 +1,146 @@
+"""
+Annotation Helper Utilities
+
+Helper functions for phoneme-level annotation tasks.
+"""
+
+import json
+import logging
+from pathlib import Path
+from typing import List, Dict, Any, Optional
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+def load_annotations(annotations_file: Path = Path("data/annotations.json")) -> List[Dict[str, Any]]:
+    """Load annotations from JSON file."""
+    if not annotations_file.exists():
+        logger.warning(f"Annotations file not found: {annotations_file}")
+        return []
+    
+    try:
+        with open(annotations_file, 'r', encoding='utf-8') as f:
+            return json.load(f)
+    except Exception as e:
+        logger.error(f"Failed to load annotations: {e}")
+        return []
+
+
+def save_annotations(annotations: List[Dict[str, Any]], annotations_file: Path = Path("data/annotations.json")):
+    """Save annotations to JSON file."""
+    annotations_file.parent.mkdir(parents=True, exist_ok=True)
+    
+    with open(annotations_file, 'w', encoding='utf-8') as f:
+        json.dump(annotations, f, indent=2, ensure_ascii=False)
+    
+    logger.info(f"Saved {len(annotations)} annotations to {annotations_file}")
+
+
+def get_annotation_statistics(annotations: List[Dict[str, Any]]) -> Dict[str, Any]:
+    """Calculate statistics from annotations."""
+    total_samples = len(annotations)
+    total_errors = sum(a.get('total_errors', 0) for a in annotations)
+    
+    error_types = {
+        'substitution': 0,
+        'omission': 0,
+        'distortion': 0,
+        'stutter': 0,
+        'normal': 0
+    }
+    
+    phoneme_errors = {}
+    
+    for ann in annotations:
+        for err in ann.get('phoneme_errors', []):
+            err_type = err.get('error_type', 'normal')
+            error_types[err_type] = error_types.get(err_type, 0) + 1
+            
+            phoneme = err.get('phoneme', 'unknown')
+            if phoneme not in phoneme_errors:
+                phoneme_errors[phoneme] = 0
+            phoneme_errors[phoneme] += 1
+    
+    return {
+        'total_samples': total_samples,
+        'total_errors': total_errors,
+        'error_types': error_types,
+        'phoneme_errors': phoneme_errors,
+        'avg_errors_per_sample': total_errors / total_samples if total_samples > 0 else 0.0
+    }
+
+
+def export_for_training(
+    annotations: List[Dict[str, Any]],
+    output_file: Path = Path("data/training_dataset.json")
+) -> Dict[str, Any]:
+    """Export annotations in training-ready format."""
+    training_data = []
+    
+    for ann in annotations:
+        audio_file = ann.get('audio_file')
+        expected_text = ann.get('expected_text', '')
+        duration = ann.get('duration', 0.0)
+        
+        # Create frame-level labels
+        num_frames = int((duration * 1000) / 20)  # 20ms frames
+        frame_labels = [0] * num_frames  # 0 = normal
+        
+        # Map errors to frames
+        for err in ann.get('phoneme_errors', []):
+            frame_id = err.get('frame_id', 0)
+            err_type = err.get('error_type', 'normal')
+            
+            # Map to 8-class system
+            class_id = {
+                'normal': 0,
+                'substitution': 1,
+                'omission': 2,
+                'distortion': 3,
+                'stutter': 4
+            }.get(err_type, 0)
+            
+            # Check if stutter + articulation error
+            if err_type != 'normal' and err_type != 'stutter':
+                # Check if there's also stutter
+                if any(e.get('error_type') == 'stutter' for e in ann.get('phoneme_errors', []) 
+                       if e.get('frame_id') == frame_id):
+                    class_id += 4  # Add 4 for stutter classes (5-7)
+            
+            if 0 <= frame_id < num_frames:
+                frame_labels[frame_id] = class_id
+        
+        training_data.append({
+            'audio_file': audio_file,
+            'expected_text': expected_text,
+            'duration': duration,
+            'num_frames': num_frames,
+            'frame_labels': frame_labels,
+            'phoneme_errors': ann.get('phoneme_errors', [])
+        })
+    
+    output_file.parent.mkdir(parents=True, exist_ok=True)
+    with open(output_file, 'w', encoding='utf-8') as f:
+        json.dump(training_data, f, indent=2, ensure_ascii=False)
+    
+    logger.info(f"Exported {len(training_data)} samples for training to {output_file}")
+    
+    return {
+        'samples': len(training_data),
+        'output_file': str(output_file)
+    }
+
+
+if __name__ == "__main__":
+    # Example usage
+    annotations = load_annotations()
+    stats = get_annotation_statistics(annotations)
+    
+    print(f"Total samples: {stats['total_samples']}")
+    print(f"Total errors: {stats['total_errors']}")
+    print(f"Error types: {stats['error_types']}")
+    
+    if annotations:
+        export_for_training(annotations)
+

scripts/data_collection.py

@@ -0,0 +1,418 @@
+"""
+Data Collection Tool for Speech Pathology Annotation
+
+This module provides a Gradio-based interface for collecting and annotating
+phoneme-level speech pathology data. Clinicians can record or upload audio,
+then annotate errors at the phoneme level with timestamps.
+
+Usage:
+    python scripts/data_collection.py
+"""
+
+import logging
+import os
+import json
+import time
+import tempfile
+from pathlib import Path
+from typing import Optional, List, Dict, Any, Tuple
+from datetime import datetime
+import numpy as np
+
+import gradio as gr
+import librosa
+import soundfile as sf
+
+from models.phoneme_mapper import PhonemeMapper
+from models.error_taxonomy import ErrorType, SeverityLevel
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Configuration
+DATA_DIR = Path("data/raw")
+ANNOTATIONS_FILE = Path("data/annotations.json")
+SAMPLE_RATE = 16000
+FRAME_DURATION_MS = 20
+
+# Ensure directories exist
+DATA_DIR.mkdir(parents=True, exist_ok=True)
+ANNOTATIONS_FILE.parent.mkdir(parents=True, exist_ok=True)
+
+# Load existing annotations
+annotations_db: List[Dict[str, Any]] = []
+if ANNOTATIONS_FILE.exists():
+    try:
+        with open(ANNOTATIONS_FILE, 'r', encoding='utf-8') as f:
+            annotations_db = json.load(f)
+        logger.info(f"✅ Loaded {len(annotations_db)} existing annotations")
+    except Exception as e:
+        logger.warning(f"⚠️ Could not load annotations: {e}")
+
+
+def save_audio_file(audio_data: Optional[Tuple[int, np.ndarray]], filename: str) -> Optional[str]:
+    """Save uploaded/recorded audio to file."""
+    if audio_data is None:
+        return None
+    
+    sample_rate, audio_array = audio_data
+    
+    # Resample to 16kHz if needed
+    if sample_rate != SAMPLE_RATE:
+        audio_array = librosa.resample(
+            audio_array.astype(np.float32),
+            orig_sr=sample_rate,
+            target_sr=SAMPLE_RATE
+        )
+        sample_rate = SAMPLE_RATE
+    
+    # Normalize
+    if np.max(np.abs(audio_array)) > 0:
+        audio_array = audio_array / np.max(np.abs(audio_array))
+    
+    # Save to data/raw
+    output_path = DATA_DIR / filename
+    sf.write(str(output_path), audio_array, sample_rate)
+    logger.info(f"✅ Saved audio to {output_path}")
+    
+    return str(output_path)
+
+
+def get_phoneme_list(text: str) -> List[str]:
+    """Convert text to phoneme list using PhonemeMapper."""
+    try:
+        mapper = PhonemeMapper(
+            frame_duration_ms=FRAME_DURATION_MS,
+            sample_rate=SAMPLE_RATE
+        )
+        phonemes = mapper.g2p.convert(text)
+        return [p for p in phonemes if p.strip()] if phonemes else []
+    except Exception as e:
+        logger.error(f"❌ G2P conversion failed: {e}")
+        return []
+
+
+def calculate_frame_count(audio_path: str) -> int:
+    """Calculate number of frames for audio file."""
+    try:
+        duration = librosa.get_duration(path=audio_path)
+        frames = int((duration * 1000) / FRAME_DURATION_MS)
+        return max(1, frames)
+    except Exception as e:
+        logger.error(f"❌ Could not calculate frames: {e}")
+        return 0
+
+
+def save_annotation(
+    audio_path: str,
+    expected_text: str,
+    phoneme_errors: List[Dict[str, Any]],
+    annotator_name: str,
+    notes: str
+) -> Dict[str, Any]:
+    """Save annotation to database."""
+    try:
+        duration = librosa.get_duration(path=audio_path)
+        
+        annotation = {
+            'id': f"annot_{int(time.time())}",
+            'audio_file': audio_path,
+            'expected_text': expected_text,
+            'duration': float(duration),
+            'annotator': annotator_name,
+            'notes': notes,
+            'created_at': datetime.utcnow().isoformat() + "Z",
+            'phoneme_errors': phoneme_errors,
+            'total_errors': len(phoneme_errors),
+            'error_types': {
+                'substitution': sum(1 for e in phoneme_errors if e.get('error_type') == 'substitution'),
+                'omission': sum(1 for e in phoneme_errors if e.get('error_type') == 'omission'),
+                'distortion': sum(1 for e in phoneme_errors if e.get('error_type') == 'distortion'),
+                'stutter': sum(1 for e in phoneme_errors if e.get('error_type') == 'stutter'),
+            }
+        }
+        
+        annotations_db.append(annotation)
+        
+        # Save to file
+        with open(ANNOTATIONS_FILE, 'w', encoding='utf-8') as f:
+            json.dump(annotations_db, f, indent=2, ensure_ascii=False)
+        
+        logger.info(f"✅ Saved annotation {annotation['id']} with {len(phoneme_errors)} errors")
+        
+        return {
+            'status': 'success',
+            'annotation_id': annotation['id'],
+            'total_errors': len(phoneme_errors),
+            'message': f"✅ Annotation saved! Total annotations: {len(annotations_db)}"
+        }
+    except Exception as e:
+        logger.error(f"❌ Failed to save annotation: {e}", exc_info=True)
+        return {
+            'status': 'error',
+            'message': f"❌ Failed to save: {str(e)}"
+        }
+
+
+def create_annotation_interface():
+    """Create Gradio interface for data collection."""
+    
+    with gr.Blocks(title="Speech Pathology Data Collection", theme=gr.themes.Soft()) as interface:
+        gr.Markdown("""
+        # 🎤 Speech Pathology Data Collection Tool
+        
+        **Purpose:** Collect and annotate phoneme-level speech pathology data for training.
+        
+        **Instructions:**
+        1. Upload or record audio (5-30 seconds, 16kHz WAV)
+        2. Enter expected text/transcript
+        3. Review phoneme list
+        4. Annotate errors at phoneme level
+        5. Save annotation
+        """)
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown("### 📥 Audio Input")
+                
+                audio_input = gr.Audio(
+                    type="numpy",
+                    label="Record or Upload Audio",
+                    sources=["microphone", "upload"],
+                    format="wav"
+                )
+                
+                expected_text = gr.Textbox(
+                    label="Expected Text/Transcript",
+                    placeholder="Enter the expected text that should be spoken",
+                    lines=3
+                )
+                
+                phoneme_display = gr.Textbox(
+                    label="Phonemes (G2P)",
+                    lines=5,
+                    interactive=False,
+                    info="Phonemes extracted from expected text"
+                )
+                
+                btn_get_phonemes = gr.Button("🔍 Extract Phonemes", variant="secondary")
+            
+            with gr.Column(scale=1):
+                gr.Markdown("### ✏️ Annotation")
+                
+                annotator_name = gr.Textbox(
+                    label="Annotator Name",
+                    placeholder="Your name",
+                    value="clinician"
+                )
+                
+                error_frame_id = gr.Number(
+                    label="Frame ID (0-based)",
+                    value=0,
+                    precision=0,
+                    info="Frame number where error occurs"
+                )
+                
+                error_phoneme = gr.Textbox(
+                    label="Phoneme with Error",
+                    placeholder="/r/",
+                    info="The phoneme that has an error"
+                )
+                
+                error_type = gr.Dropdown(
+                    label="Error Type",
+                    choices=["normal", "substitution", "omission", "distortion", "stutter"],
+                    value="normal",
+                    info="Type of error detected"
+                )
+                
+                wrong_sound = gr.Textbox(
+                    label="Wrong Sound (if substitution)",
+                    placeholder="/w/",
+                    info="What sound was produced instead (for substitutions)"
+                )
+                
+                error_severity = gr.Slider(
+                    label="Severity (0-1)",
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.5,
+                    step=0.1,
+                    info="Severity of the error"
+                )
+                
+                error_timestamp = gr.Number(
+                    label="Timestamp (seconds)",
+                    value=0.0,
+                    precision=2,
+                    info="Time in audio where error occurs"
+                )
+                
+                btn_add_error = gr.Button("➕ Add Error", variant="primary")
+                
+                errors_list = gr.Dataframe(
+                    label="Annotated Errors",
+                    headers=["Frame", "Phoneme", "Type", "Wrong Sound", "Severity", "Time"],
+                    interactive=False,
+                    wrap=True
+                )
+                
+                notes = gr.Textbox(
+                    label="Notes",
+                    placeholder="Additional notes about this sample",
+                    lines=3
+                )
+                
+                btn_save = gr.Button("💾 Save Annotation", variant="primary", size="lg")
+                
+                output_status = gr.Textbox(
+                    label="Status",
+                    interactive=False,
+                    lines=3
+                )
+        
+        # Statistics panel
+        with gr.Row():
+            gr.Markdown("### 📊 Statistics")
+            stats_display = gr.Markdown("**Total Annotations:** 0 | **Total Errors:** 0")
+        
+        # Event handlers
+        errors_data = gr.State(value=[])
+        
+        def extract_phonemes(text: str) -> str:
+            """Extract phonemes from text."""
+            if not text:
+                return "Enter expected text first"
+            phonemes = get_phoneme_list(text)
+            return " ".join([f"/{p}/" for p in phonemes]) if phonemes else "No phonemes found"
+        
+        def add_error(
+            frame_id: int,
+            phoneme: str,
+            error_type: str,
+            wrong_sound: str,
+            severity: float,
+            timestamp: float,
+            current_errors: List[Dict]
+        ) -> Tuple[List[Dict], gr.Dataframe]:
+            """Add an error to the list."""
+            error = {
+                'frame_id': int(frame_id),
+                'phoneme': phoneme.strip(),
+                'error_type': error_type,
+                'wrong_sound': wrong_sound.strip() if wrong_sound else None,
+                'severity': float(severity),
+                'timestamp': float(timestamp),
+                'confidence': 1.0  # Manual annotation is always confident
+            }
+            
+            new_errors = current_errors + [error]
+            
+            # Create dataframe
+            df_data = [
+                [
+                    e['frame_id'],
+                    e['phoneme'],
+                    e['error_type'],
+                    e.get('wrong_sound', 'N/A'),
+                    f"{e['severity']:.2f}",
+                    f"{e['timestamp']:.2f}s"
+                ]
+                for e in new_errors
+            ]
+            
+            return new_errors, df_data
+        
+        def save_annotation_handler(
+            audio_data: Optional[Tuple[int, np.ndarray]],
+            expected_text: str,
+            errors: List[Dict],
+            annotator: str,
+            notes: str
+        ) -> str:
+            """Handle annotation saving."""
+            if audio_data is None:
+                return "❌ Please provide audio first"
+            
+            if not expected_text:
+                return "❌ Please provide expected text"
+            
+            # Save audio
+            filename = f"sample_{int(time.time())}.wav"
+            audio_path = save_audio_file(audio_data, filename)
+            
+            if not audio_path:
+                return "❌ Failed to save audio file"
+            
+            # Save annotation
+            result = save_annotation(
+                audio_path=audio_path,
+                expected_text=expected_text,
+                phoneme_errors=errors,
+                annotator_name=annotator,
+                notes=notes
+            )
+            
+            return result.get('message', 'Unknown status')
+        
+        def update_stats() -> str:
+            """Update statistics display."""
+            total_annotations = len(annotations_db)
+            total_errors = sum(a.get('total_errors', 0) for a in annotations_db)
+            
+            error_breakdown = {}
+            for ann in annotations_db:
+                for err_type, count in ann.get('error_types', {}).items():
+                    error_breakdown[err_type] = error_breakdown.get(err_type, 0) + count
+            
+            stats_text = f"""
+            **Total Annotations:** {total_annotations} | **Total Errors:** {total_errors}
+            
+            **Error Breakdown:**
+            - Substitution: {error_breakdown.get('substitution', 0)}
+            - Omission: {error_breakdown.get('omission', 0)}
+            - Distortion: {error_breakdown.get('distortion', 0)}
+            - Stutter: {error_breakdown.get('stutter', 0)}
+            """
+            return stats_text
+        
+        # Wire up events
+        btn_get_phonemes.click(
+            fn=extract_phonemes,
+            inputs=[expected_text],
+            outputs=[phoneme_display]
+        )
+        
+        btn_add_error.click(
+            fn=add_error,
+            inputs=[
+                error_frame_id,
+                error_phoneme,
+                error_type,
+                wrong_sound,
+                error_severity,
+                error_timestamp,
+                errors_data
+            ],
+            outputs=[errors_data, errors_list]
+        )
+        
+        btn_save.click(
+            fn=save_annotation_handler,
+            inputs=[audio_input, expected_text, errors_data, annotator_name, notes],
+            outputs=[output_status]
+        ).then(
+            fn=update_stats,
+            outputs=[stats_display]
+        )
+        
+        # Load stats on startup
+        interface.load(fn=update_stats, outputs=[stats_display])
+    
+    return interface
+
+
+if __name__ == "__main__":
+    interface = create_annotation_interface()
+    interface.launch(server_name="0.0.0.0", server_port=7861, share=False)
+

training/config.yaml

@@ -0,0 +1,86 @@
+# Training Configuration for Classifier Head
+
+# Data Configuration
+data:
+  annotations_file: "data/annotations.json"
+  training_dataset: "data/training_dataset.json"
+  train_split: 0.8
+  val_split: 0.2
+  test_split: 0.0  # Use validation as test if test_split is 0
+  random_seed: 42
+
+# Model Configuration
+model:
+  input_dim: 1024  # Wav2Vec2-XLSR-53 feature dimension
+  hidden_dims: [512, 256]  # Shared layers
+  dropout: 0.1
+  num_classes: 8  # 8-class output (fluency + articulation combined)
+  use_pretrained_head: false  # Set to true after first training
+
+# Training Configuration
+training:
+  batch_size: 16
+  num_epochs: 50
+  learning_rate: 0.001
+  weight_decay: 0.0001
+  gradient_clip_norm: 1.0
+  
+  # Loss Configuration
+  loss:
+    type: "cross_entropy"  # or "focal" for imbalanced data
+    class_weights: null  # Auto-calculate from data if null
+    focal_alpha: 0.25
+    focal_gamma: 2.0
+  
+  # Optimizer
+  optimizer: "adam"
+  adam_betas: [0.9, 0.999]
+  
+  # Scheduler
+  scheduler: "reduce_on_plateau"
+  scheduler_patience: 5
+  scheduler_factor: 0.5
+  scheduler_min_lr: 0.00001
+  
+  # Early Stopping
+  early_stopping:
+    enabled: true
+    patience: 10
+    min_delta: 0.001
+    monitor: "val_loss"
+  
+  # Data Augmentation
+  augmentation:
+    enabled: false  # Enable after initial training
+    time_stretch: [0.9, 1.1]
+    noise_injection: 0.01
+    pitch_shift: [-2, 2]  # semitones
+
+# Validation Configuration
+validation:
+  metrics: ["accuracy", "f1_score", "precision", "recall", "confusion_matrix"]
+  per_class_metrics: true
+  save_predictions: true
+
+# Checkpoint Configuration
+checkpoint:
+  save_dir: "models/checkpoints"
+  save_best: true
+  save_last: true
+  save_frequency: 5  # Save every N epochs
+  filename: "classifier_head_trained.pt"
+  best_filename: "classifier_head_best.pt"
+
+# Logging Configuration
+logging:
+  log_dir: "training/logs"
+  tensorboard: false  # Enable if tensorboard installed
+  wandb: false  # Enable if wandb installed
+  log_frequency: 10  # Log every N batches
+
+# Device Configuration
+device:
+  use_cuda: true
+  cuda_device: 0
+  mixed_precision: false  # Use FP16 training
+

training/evaluate_classifier.py

@@ -0,0 +1,282 @@
+"""
+Evaluation Script for Trained Classifier Head
+
+Evaluates the trained classifier on test/validation data and generates
+comprehensive metrics including per-class accuracy, confusion matrix, etc.
+
+Usage:
+    python training/evaluate_classifier.py --checkpoint models/checkpoints/classifier_head_best.pt
+"""
+
+import logging
+import argparse
+import json
+import yaml
+from pathlib import Path
+from typing import Dict, List, Any
+import sys
+
+import torch
+import numpy as np
+from sklearn.metrics import (
+    accuracy_score, f1_score, precision_score, recall_score,
+    confusion_matrix, classification_report
+)
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+# Add project root to path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from training.train_classifier_head import PhonemeDataset, collate_fn
+from torch.utils.data import DataLoader
+from models.speech_pathology_model import SpeechPathologyClassifier
+from models.phoneme_mapper import PhonemeMapper
+from inference.inference_pipeline import InferencePipeline
+from config import default_audio_config, default_model_config, default_inference_config
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def load_trained_model(checkpoint_path: Path, config_path: Path = Path("training/config.yaml")) -> torch.nn.Module:
+    """Load trained classifier head from checkpoint."""
+    # Load config
+    with open(config_path, 'r') as f:
+        config = yaml.safe_load(f)
+    
+    # Initialize inference pipeline
+    inference_pipeline = InferencePipeline(
+        audio_config=default_audio_config,
+        model_config=default_model_config,
+        inference_config=default_inference_config
+    )
+    
+    model = inference_pipeline.model
+    
+    # Load checkpoint
+    checkpoint = torch.load(checkpoint_path, map_location='cpu')
+    model.classifier_head.load_state_dict(checkpoint['model_state_dict'])
+    
+    logger.info(f"✅ Loaded checkpoint from epoch {checkpoint.get('epoch', 'unknown')}")
+    logger.info(f"   Validation loss: {checkpoint.get('val_loss', 'unknown'):.4f}")
+    logger.info(f"   Validation accuracy: {checkpoint.get('val_accuracy', 'unknown'):.4f}")
+    
+    return model
+
+
+def evaluate_model(
+    model: torch.nn.Module,
+    dataloader: DataLoader,
+    device: torch.device,
+    class_names: List[str]
+) -> Dict[str, Any]:
+    """Evaluate model and return comprehensive metrics."""
+    model.eval()
+    
+    all_preds = []
+    all_labels = []
+    all_probs = []
+    
+    with torch.no_grad():
+        for batch in dataloader:
+            features = batch['features'].to(device)
+            labels = batch['labels'].to(device)
+            
+            batch_size, seq_len, feat_dim = features.shape
+            features_flat = features.view(-1, feat_dim)
+            labels_flat = labels.view(-1)
+            
+            # Forward pass
+            shared_features = model.classifier_head.shared_layers(features_flat)
+            logits = model.classifier_head.full_head(shared_features)
+            probs = torch.softmax(logits, dim=-1)
+            
+            preds = torch.argmax(logits, dim=-1).cpu().numpy()
+            all_preds.extend(preds)
+            all_labels.extend(labels_flat.cpu().numpy())
+            all_probs.extend(probs.cpu().numpy())
+    
+    # Calculate metrics
+    accuracy = accuracy_score(all_labels, all_preds)
+    f1_macro = f1_score(all_labels, all_preds, average='macro', zero_division=0)
+    f1_weighted = f1_score(all_labels, all_preds, average='weighted', zero_division=0)
+    precision_macro = precision_score(all_labels, all_preds, average='macro', zero_division=0)
+    recall_macro = recall_score(all_labels, all_preds, average='macro', zero_division=0)
+    
+    # Per-class metrics
+    cm = confusion_matrix(all_labels, all_preds, labels=list(range(len(class_names))))
+    
+    # Per-class accuracy
+    per_class_accuracy = cm.diagonal() / cm.sum(axis=1)
+    per_class_accuracy = np.nan_to_num(per_class_accuracy)  # Handle division by zero
+    
+    # Classification report
+    report = classification_report(
+        all_labels, all_preds,
+        target_names=class_names,
+        output_dict=True,
+        zero_division=0
+    )
+    
+    # Confidence analysis
+    all_probs = np.array(all_probs)
+    max_probs = np.max(all_probs, axis=1)
+    correct_mask = np.array(all_preds) == np.array(all_labels)
+    
+    avg_confidence_correct = np.mean(max_probs[correct_mask]) if np.any(correct_mask) else 0.0
+    avg_confidence_incorrect = np.mean(max_probs[~correct_mask]) if np.any(~correct_mask) else 0.0
+    
+    return {
+        'overall_accuracy': float(accuracy),
+        'f1_macro': float(f1_macro),
+        'f1_weighted': float(f1_weighted),
+        'precision_macro': float(precision_macro),
+        'recall_macro': float(recall_macro),
+        'confusion_matrix': cm.tolist(),
+        'per_class_accuracy': per_class_accuracy.tolist(),
+        'classification_report': report,
+        'confidence': {
+            'avg_correct': float(avg_confidence_correct),
+            'avg_incorrect': float(avg_confidence_incorrect),
+            'confidence_distribution': {
+                'mean': float(np.mean(max_probs)),
+                'std': float(np.std(max_probs)),
+                'min': float(np.min(max_probs)),
+                'max': float(np.max(max_probs))
+            }
+        },
+        'num_samples': len(all_labels)
+    }
+
+
+def plot_confusion_matrix(cm: np.ndarray, class_names: List[str], output_path: Path):
+    """Plot and save confusion matrix."""
+    plt.figure(figsize=(10, 8))
+    sns.heatmap(
+        cm,
+        annot=True,
+        fmt='d',
+        cmap='Blues',
+        xticklabels=class_names,
+        yticklabels=class_names
+    )
+    plt.title('Confusion Matrix')
+    plt.ylabel('True Label')
+    plt.xlabel('Predicted Label')
+    plt.tight_layout()
+    plt.savefig(output_path)
+    logger.info(f"✅ Saved confusion matrix to {output_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Evaluate trained classifier")
+    parser.add_argument('--checkpoint', type=str, required=True,
+                       help='Path to checkpoint file')
+    parser.add_argument('--config', type=str, default='training/config.yaml',
+                       help='Path to config file')
+    parser.add_argument('--dataset', type=str, default='data/training_dataset.json',
+                       help='Path to evaluation dataset')
+    parser.add_argument('--output', type=str, default='training/evaluation_results.json',
+                       help='Path to save evaluation results')
+    parser.add_argument('--plot', type=str, default='training/confusion_matrix.png',
+                       help='Path to save confusion matrix plot')
+    args = parser.parse_args()
+    
+    # Load config
+    with open(args.config, 'r') as f:
+        config = yaml.safe_load(f)
+    
+    # Set device
+    device = torch.device('cuda' if torch.cuda.is_available() and config['device']['use_cuda'] else 'cpu')
+    logger.info(f"Using device: {device}")
+    
+    # Load model
+    checkpoint_path = Path(args.checkpoint)
+    if not checkpoint_path.exists():
+        logger.error(f"Checkpoint not found: {checkpoint_path}")
+        return
+    
+    model = load_trained_model(checkpoint_path, Path(args.config))
+    model = model.to(device)
+    
+    # Load evaluation dataset
+    dataset_path = Path(args.dataset)
+    if not dataset_path.exists():
+        logger.error(f"Dataset not found: {dataset_path}")
+        return
+    
+    with open(dataset_path, 'r') as f:
+        eval_data = json.load(f)
+    
+    logger.info(f"Loaded {len(eval_data)} evaluation samples")
+    
+    # Create dataset and dataloader
+    inference_pipeline = InferencePipeline(
+        audio_config=default_audio_config,
+        model_config=default_model_config,
+        inference_config=default_inference_config
+    )
+    
+    phoneme_mapper = PhonemeMapper(frame_duration_ms=20, sample_rate=16000)
+    
+    from training.train_classifier_head import PhonemeDataset
+    dataset = PhonemeDataset(eval_data, inference_pipeline, phoneme_mapper)
+    
+    dataloader = DataLoader(
+        dataset,
+        batch_size=config['training']['batch_size'],
+        shuffle=False,
+        collate_fn=collate_fn
+    )
+    
+    # Class names
+    class_names = [
+        "Normal",
+        "Substitution",
+        "Omission",
+        "Distortion",
+        "Normal+Stutter",
+        "Substitution+Stutter",
+        "Omission+Stutter",
+        "Distortion+Stutter"
+    ]
+    
+    # Evaluate
+    logger.info("Evaluating model...")
+    metrics = evaluate_model(model, dataloader, device, class_names)
+    
+    # Print results
+    logger.info("\n" + "="*50)
+    logger.info("EVALUATION RESULTS")
+    logger.info("="*50)
+    logger.info(f"Overall Accuracy: {metrics['overall_accuracy']:.4f}")
+    logger.info(f"F1 Score (macro): {metrics['f1_macro']:.4f}")
+    logger.info(f"F1 Score (weighted): {metrics['f1_weighted']:.4f}")
+    logger.info(f"Precision (macro): {metrics['precision_macro']:.4f}")
+    logger.info(f"Recall (macro): {metrics['recall_macro']:.4f}")
+    logger.info(f"\nPer-Class Accuracy:")
+    for i, (name, acc) in enumerate(zip(class_names, metrics['per_class_accuracy'])):
+        logger.info(f"  {name}: {acc:.4f}")
+    logger.info(f"\nConfidence Analysis:")
+    logger.info(f"  Avg confidence (correct): {metrics['confidence']['avg_correct']:.4f}")
+    logger.info(f"  Avg confidence (incorrect): {metrics['confidence']['avg_incorrect']:.4f}")
+    
+    # Save results
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(output_path, 'w') as f:
+        json.dump(metrics, f, indent=2)
+    logger.info(f"\n✅ Saved evaluation results to {output_path}")
+    
+    # Plot confusion matrix
+    if args.plot:
+        plot_path = Path(args.plot)
+        plot_path.parent.mkdir(parents=True, exist_ok=True)
+        cm = np.array(metrics['confusion_matrix'])
+        plot_confusion_matrix(cm, class_names, plot_path)
+
+
+if __name__ == "__main__":
+    main()
+

training/train_classifier_head.py

@@ -0,0 +1,469 @@
+"""
+Training Script for Speech Pathology Classifier Head
+
+This script fine-tunes the classification head on phoneme-level labeled data.
+Wav2Vec2 encoder is frozen; only the classifier head is trained.
+
+Usage:
+    python training/train_classifier_head.py --config training/config.yaml
+"""
+
+import logging
+import os
+import sys
+import json
+import yaml
+import argparse
+from pathlib import Path
+from typing import Dict, List, Tuple, Optional, Any
+from datetime import datetime
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader, random_split
+import numpy as np
+from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, confusion_matrix
+import librosa
+import soundfile as sf
+
+# Add project root to path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from models.speech_pathology_model import SpeechPathologyClassifier, MultiTaskClassifierHead
+from models.phoneme_mapper import PhonemeMapper
+from inference.inference_pipeline import InferencePipeline
+from config import default_audio_config, default_model_config, default_inference_config
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+
+class PhonemeDataset(Dataset):
+    """Dataset for phoneme-level speech pathology training."""
+    
+    def __init__(
+        self,
+        training_data: List[Dict[str, Any]],
+        inference_pipeline: InferencePipeline,
+        phoneme_mapper: PhonemeMapper
+    ):
+        """
+        Initialize dataset.
+        
+        Args:
+            training_data: List of training samples with frame labels
+            inference_pipeline: Pipeline for extracting Wav2Vec2 features
+            phoneme_mapper: Mapper for phoneme alignment
+        """
+        self.training_data = training_data
+        self.inference_pipeline = inference_pipeline
+        self.phoneme_mapper = phoneme_mapper
+        
+        logger.info(f"Initialized dataset with {len(training_data)} samples")
+    
+    def __len__(self) -> int:
+        return len(self.training_data)
+    
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        """Get a training sample."""
+        sample = self.training_data[idx]
+        audio_file = sample['audio_file']
+        frame_labels = sample['frame_labels']
+        
+        # Load audio
+        try:
+            audio, sr = librosa.load(audio_file, sr=16000)
+        except Exception as e:
+            logger.error(f"Failed to load {audio_file}: {e}")
+            # Return dummy data
+            return {
+                'features': torch.zeros(1, 1024),
+                'labels': torch.tensor([0], dtype=torch.long),
+                'valid': torch.tensor(False)
+            }
+        
+        # Extract Wav2Vec2 features
+        try:
+            frame_features, frame_times = self.inference_pipeline.get_phone_level_features(audio)
+            
+            # Align labels with features
+            num_features = len(frame_features)
+            num_labels = len(frame_labels)
+            
+            # Pad or truncate labels to match features
+            if num_labels < num_features:
+                frame_labels = frame_labels + [0] * (num_features - num_labels)
+            elif num_labels > num_features:
+                frame_labels = frame_labels[:num_features]
+            
+            # Convert to tensors
+            features_tensor = frame_features  # Already a tensor
+            labels_tensor = torch.tensor(frame_labels[:num_features], dtype=torch.long)
+            
+            return {
+                'features': features_tensor,
+                'labels': labels_tensor,
+                'valid': torch.tensor(True)
+            }
+        except Exception as e:
+            logger.error(f"Failed to extract features from {audio_file}: {e}")
+            return {
+                'features': torch.zeros(1, 1024),
+                'labels': torch.tensor([0], dtype=torch.long),
+                'valid': torch.tensor(False)
+            }
+
+
+def collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
+    """Collate function for DataLoader."""
+    # Filter out invalid samples
+    valid_batch = [b for b in batch if b['valid'].item()]
+    
+    if not valid_batch:
+        # Return dummy batch
+        return {
+            'features': torch.zeros(1, 1, 1024),
+            'labels': torch.zeros(1, 1, dtype=torch.long)
+        }
+    
+    # Stack features and labels
+    features_list = []
+    labels_list = []
+    
+    for item in valid_batch:
+        features_list.append(item['features'])
+        labels_list.append(item['labels'])
+    
+    # Pad to same length
+    max_len = max(f.shape[0] for f in features_list)
+    
+    padded_features = []
+    padded_labels = []
+    
+    for feat, lab in zip(features_list, labels_list):
+        if feat.shape[0] < max_len:
+            padding = max_len - feat.shape[0]
+            feat = torch.cat([feat, torch.zeros(padding, feat.shape[1])])
+            lab = torch.cat([lab, torch.zeros(padding, dtype=torch.long)])
+        padded_features.append(feat)
+        padded_labels.append(lab)
+    
+    return {
+        'features': torch.stack(padded_features),
+        'labels': torch.stack(padded_labels)
+    }
+
+
+def calculate_class_weights(dataset: PhonemeDataset) -> torch.Tensor:
+    """Calculate class weights for imbalanced data."""
+    all_labels = []
+    for i in range(len(dataset)):
+        sample = dataset[i]
+        if sample['valid'].item():
+            all_labels.extend(sample['labels'].tolist())
+    
+    if not all_labels:
+        return torch.ones(8)
+    
+    unique, counts = np.unique(all_labels, return_counts=True)
+    total = len(all_labels)
+    
+    weights = torch.ones(8)
+    for cls, count in zip(unique, counts):
+        if count > 0:
+            weights[int(cls)] = total / (8 * count)  # Inverse frequency weighting
+    
+    logger.info(f"Class weights: {weights.tolist()}")
+    return weights
+
+
+def train_epoch(
+    model: nn.Module,
+    dataloader: DataLoader,
+    optimizer: optim.Optimizer,
+    criterion: nn.Module,
+    device: torch.device,
+    epoch: int
+) -> Dict[str, float]:
+    """Train for one epoch."""
+    model.train()
+    total_loss = 0.0
+    all_preds = []
+    all_labels = []
+    
+    for batch_idx, batch in enumerate(dataloader):
+        features = batch['features'].to(device)  # (batch, seq_len, 1024)
+        labels = batch['labels'].to(device)  # (batch, seq_len)
+        
+        # Flatten for processing
+        batch_size, seq_len, feat_dim = features.shape
+        features_flat = features.view(-1, feat_dim)  # (batch * seq_len, 1024)
+        labels_flat = labels.view(-1)  # (batch * seq_len)
+        
+        # Forward pass
+        optimizer.zero_grad()
+        
+        # Get predictions from full_head
+        shared_features = model.classifier_head.shared_layers(features_flat)
+        logits = model.classifier_head.full_head(shared_features)  # (batch * seq_len, 8)
+        
+        # Calculate loss
+        loss = criterion(logits, labels_flat)
+        
+        # Backward pass
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.classifier_head.parameters(), max_norm=1.0)
+        optimizer.step()
+        
+        # Metrics
+        total_loss += loss.item()
+        preds = torch.argmax(logits, dim=-1).cpu().numpy()
+        all_preds.extend(preds)
+        all_labels.extend(labels_flat.cpu().numpy())
+        
+        if batch_idx % 10 == 0:
+            logger.info(f"Epoch {epoch}, Batch {batch_idx}/{len(dataloader)}, Loss: {loss.item():.4f}")
+    
+    avg_loss = total_loss / len(dataloader)
+    accuracy = accuracy_score(all_labels, all_preds)
+    
+    return {
+        'loss': avg_loss,
+        'accuracy': accuracy
+    }
+
+
+def validate(
+    model: nn.Module,
+    dataloader: DataLoader,
+    criterion: nn.Module,
+    device: torch.device
+) -> Dict[str, float]:
+    """Validate model."""
+    model.eval()
+    total_loss = 0.0
+    all_preds = []
+    all_labels = []
+    
+    with torch.no_grad():
+        for batch in dataloader:
+            features = batch['features'].to(device)
+            labels = batch['labels'].to(device)
+            
+            batch_size, seq_len, feat_dim = features.shape
+            features_flat = features.view(-1, feat_dim)
+            labels_flat = labels.view(-1)
+            
+            # Forward pass
+            shared_features = model.classifier_head.shared_layers(features_flat)
+            logits = model.classifier_head.full_head(shared_features)
+            
+            loss = criterion(logits, labels_flat)
+            total_loss += loss.item()
+            
+            preds = torch.argmax(logits, dim=-1).cpu().numpy()
+            all_preds.extend(preds)
+            all_labels.extend(labels_flat.cpu().numpy())
+    
+    avg_loss = total_loss / len(dataloader)
+    accuracy = accuracy_score(all_labels, all_preds)
+    f1 = f1_score(all_labels, all_preds, average='weighted', zero_division=0)
+    precision = precision_score(all_labels, all_preds, average='weighted', zero_division=0)
+    recall = recall_score(all_labels, all_preds, average='weighted', zero_division=0)
+    
+    # Per-class metrics
+    cm = confusion_matrix(all_labels, all_preds, labels=list(range(8)))
+    
+    return {
+        'loss': avg_loss,
+        'accuracy': accuracy,
+        'f1_score': f1,
+        'precision': precision,
+        'recall': recall,
+        'confusion_matrix': cm.tolist()
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train classifier head")
+    parser.add_argument('--config', type=str, default='training/config.yaml',
+                       help='Path to config file')
+    parser.add_argument('--resume', type=str, default=None,
+                       help='Resume from checkpoint')
+    args = parser.parse_args()
+    
+    # Load config
+    with open(args.config, 'r') as f:
+        config = yaml.safe_load(f)
+    
+    # Set device
+    device = torch.device('cuda' if torch.cuda.is_available() and config['device']['use_cuda'] else 'cpu')
+    logger.info(f"Using device: {device}")
+    
+    # Load training data
+    training_file = Path(config['data']['training_dataset'])
+    if not training_file.exists():
+        logger.error(f"Training dataset not found: {training_file}")
+        logger.info("Run scripts/annotation_helper.py to export training data first")
+        return
+    
+    with open(training_file, 'r') as f:
+        training_data = json.load(f)
+    
+    logger.info(f"Loaded {len(training_data)} training samples")
+    
+    # Initialize inference pipeline for feature extraction
+    inference_pipeline = InferencePipeline(
+        audio_config=default_audio_config,
+        model_config=default_model_config,
+        inference_config=default_inference_config
+    )
+    
+    # Initialize phoneme mapper
+    phoneme_mapper = PhonemeMapper(
+        frame_duration_ms=20,
+        sample_rate=16000
+    )
+    
+    # Create dataset
+    dataset = PhonemeDataset(training_data, inference_pipeline, phoneme_mapper)
+    
+    # Split dataset
+    train_size = int(config['data']['train_split'] * len(dataset))
+    val_size = len(dataset) - train_size
+    
+    train_dataset, val_dataset = random_split(
+        dataset,
+        [train_size, val_size],
+        generator=torch.Generator().manual_seed(config['data']['random_seed'])
+    )
+    
+    logger.info(f"Train samples: {len(train_dataset)}, Val samples: {len(val_dataset)}")
+    
+    # Create data loaders
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=config['training']['batch_size'],
+        shuffle=True,
+        collate_fn=collate_fn
+    )
+    
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=config['training']['batch_size'],
+        shuffle=False,
+        collate_fn=collate_fn
+    )
+    
+    # Load model
+    model = inference_pipeline.model
+    model.train()  # Set to training mode
+    
+    # Freeze Wav2Vec2 (should already be frozen, but ensure it)
+    for param in model.wav2vec2_model.parameters():
+        param.requires_grad = False
+    
+    # Unfreeze classifier head
+    for param in model.classifier_head.parameters():
+        param.requires_grad = True
+    
+    logger.info("Model prepared: Wav2Vec2 frozen, classifier head trainable")
+    
+    # Calculate class weights
+    class_weights = calculate_class_weights(dataset)
+    class_weights = class_weights.to(device)
+    
+    # Loss function
+    if config['training']['loss']['type'] == 'cross_entropy':
+        criterion = nn.CrossEntropyLoss(weight=class_weights)
+    else:
+        # Focal loss implementation would go here
+        criterion = nn.CrossEntropyLoss(weight=class_weights)
+    
+    # Optimizer
+    optimizer = optim.Adam(
+        model.classifier_head.parameters(),
+        lr=config['training']['learning_rate'],
+        weight_decay=config['training']['weight_decay']
+    )
+    
+    # Scheduler
+    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer,
+        mode='min',
+        factor=config['training']['scheduler_factor'],
+        patience=config['training']['scheduler_patience'],
+        min_lr=config['training']['scheduler_min_lr']
+    )
+    
+    # Training loop
+    best_val_loss = float('inf')
+    patience_counter = 0
+    
+    checkpoint_dir = Path(config['checkpoint']['save_dir'])
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    
+    for epoch in range(config['training']['num_epochs']):
+        logger.info(f"\n{'='*50}")
+        logger.info(f"Epoch {epoch+1}/{config['training']['num_epochs']}")
+        logger.info(f"{'='*50}")
+        
+        # Train
+        train_metrics = train_epoch(model, train_loader, optimizer, criterion, device, epoch+1)
+        logger.info(f"Train - Loss: {train_metrics['loss']:.4f}, Accuracy: {train_metrics['accuracy']:.4f}")
+        
+        # Validate
+        val_metrics = validate(model, val_loader, criterion, device)
+        logger.info(f"Val - Loss: {val_metrics['loss']:.4f}, Accuracy: {val_metrics['accuracy']:.4f}, "
+                   f"F1: {val_metrics['f1_score']:.4f}")
+        
+        # Scheduler step
+        scheduler.step(val_metrics['loss'])
+        
+        # Save checkpoint
+        if config['checkpoint']['save_best'] and val_metrics['loss'] < best_val_loss:
+            best_val_loss = val_metrics['loss']
+            checkpoint_path = checkpoint_dir / config['checkpoint']['best_filename']
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.classifier_head.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'val_loss': val_metrics['loss'],
+                'val_accuracy': val_metrics['accuracy'],
+                'config': config
+            }, checkpoint_path)
+            logger.info(f"✅ Saved best checkpoint to {checkpoint_path}")
+            patience_counter = 0
+        else:
+            patience_counter += 1
+        
+        # Early stopping
+        if config['training']['early_stopping']['enabled']:
+            if patience_counter >= config['training']['early_stopping']['patience']:
+                logger.info(f"Early stopping triggered after {epoch+1} epochs")
+                break
+        
+        # Save last checkpoint
+        if config['checkpoint']['save_last'] and (epoch + 1) % config['checkpoint']['save_frequency'] == 0:
+            checkpoint_path = checkpoint_dir / config['checkpoint']['filename']
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.classifier_head.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'val_loss': val_metrics['loss'],
+                'val_accuracy': val_metrics['accuracy'],
+                'config': config
+            }, checkpoint_path)
+            logger.info(f"Saved checkpoint to {checkpoint_path}")
+    
+    logger.info("\n✅ Training complete!")
+    logger.info(f"Best validation loss: {best_val_loss:.4f}")
+
+
+if __name__ == "__main__":
+    main()
+

ui/gradio_interface.py

@@ -17,6 +17,8 @@ import numpy as np
 import gradio as gr
 
 from diagnosis.ai_engine.model_loader import get_inference_pipeline
+from api.routes import get_phoneme_mapper, get_error_mapper
+from models.error_taxonomy import ErrorType, SeverityLevel
 from config import GradioConfig, default_gradio_config
 
 logger = logging.getLogger(__name__)
@@ -83,8 +85,9 @@ def format_articulation_issues(articulation_scores: list) -> str:
 
 def analyze_speech(
     audio_input: Optional[Tuple[int, np.ndarray]],
-    audio_file: Optional[str]
-) -> Tuple[str, str, str, str, Dict[str, Any]]:
+    audio_file: Optional[str],
+    expected_text: Optional[str] = None
+) -> Tuple[str, str, str, str, str, Dict[str, Any]]:
     """
     Analyze speech audio for fluency and articulation issues.
     
@@ -167,6 +170,72 @@ def analyze_speech(
         except: pass
         # #endregion
         
+        # Get phoneme and error mappers
+        phoneme_mapper = get_phoneme_mapper()
+        error_mapper = get_error_mapper()
+        
+        # Map phonemes to frames if text provided
+        frame_phonemes = []
+        if expected_text and phoneme_mapper:
+            try:
+                frame_phonemes = phoneme_mapper.map_text_to_frames(
+                    expected_text,
+                    num_frames=result.num_frames,
+                    audio_duration=result.duration
+                )
+                logger.info(f"✅ Mapped {len(frame_phonemes)} phonemes to frames")
+            except Exception as e:
+                logger.warning(f"⚠️ Phoneme mapping failed: {e}")
+                frame_phonemes = [''] * result.num_frames
+        else:
+            frame_phonemes = [''] * result.num_frames
+        
+        # Process errors with error mapper
+        errors = []
+        error_table_rows = []
+        
+        for i, frame_pred in enumerate(result.frame_predictions):
+            phoneme = frame_phonemes[i] if i < len(frame_phonemes) else ''
+            
+            # Map classifier output to error detail (8-class system)
+            class_id = frame_pred.articulation_class
+            if frame_pred.fluency_label == 'stutter':
+                class_id += 4  # Add 4 for stutter classes (4-7)
+            
+            # Get error detail
+            if error_mapper:
+                try:
+                    error_detail = error_mapper.map_classifier_output(
+                        class_id=class_id,
+                        confidence=frame_pred.confidence,
+                        phoneme=phoneme if phoneme else 'unknown',
+                        fluency_label=frame_pred.fluency_label
+                    )
+                    
+                    if error_detail.error_type != ErrorType.NORMAL:
+                        errors.append((i, frame_pred.time, error_detail))
+                        
+                        # Add to error table
+                        severity_level = error_mapper.get_severity_level(error_detail.severity)
+                        severity_color = {
+                            SeverityLevel.NONE: "green",
+                            SeverityLevel.LOW: "orange",
+                            SeverityLevel.MEDIUM: "orange",
+                            SeverityLevel.HIGH: "red"
+                        }.get(severity_level, "gray")
+                        
+                        error_table_rows.append({
+                            "phoneme": error_detail.phoneme,
+                            "time": f"{frame_pred.time:.2f}s",
+                            "error_type": error_detail.error_type.value,
+                            "wrong_sound": error_detail.wrong_sound or "N/A",
+                            "severity": severity_level.value,
+                            "severity_color": severity_color,
+                            "therapy": error_detail.therapy[:80] + "..." if len(error_detail.therapy) > 80 else error_detail.therapy
+                        })
+                except Exception as e:
+                    logger.warning(f"Error mapping failed for frame {i}: {e}")
+        
         # Calculate processing time
         processing_time_ms = (time.time() - start_time) * 1000
         
@@ -240,7 +309,121 @@ def analyze_speech(
         </div>
         """
         
-        # Create JSON output
+        # Format error table with summary of problematic sounds
+        if error_table_rows:
+            # Group errors by phoneme to show which sounds have issues
+            phoneme_errors = {}
+            for row in error_table_rows:
+                phoneme = row['phoneme']
+                if phoneme not in phoneme_errors:
+                    phoneme_errors[phoneme] = {
+                        'count': 0,
+                        'types': set(),
+                        'severity': 'low',
+                        'examples': []
+                    }
+                phoneme_errors[phoneme]['count'] += 1
+                phoneme_errors[phoneme]['types'].add(row['error_type'])
+                if row['severity'] in ['high', 'medium']:
+                    phoneme_errors[phoneme]['severity'] = row['severity']
+                if len(phoneme_errors[phoneme]['examples']) < 2:
+                    phoneme_errors[phoneme]['examples'].append(row)
+            
+            # Create summary section
+            problematic_sounds = sorted(phoneme_errors.keys())
+            summary_html = f"""
+            <div style='background-color: #fff3cd; border: 2px solid #ffc107; border-radius: 8px; padding: 15px; margin-bottom: 20px;'>
+                <h3 style='color: #856404; margin-top: 0;'>⚠️ Problematic Sounds Detected</h3>
+                <p style='color: #856404; font-size: 14px; margin-bottom: 10px;'>
+                    <strong>{len(problematic_sounds)} sound(s) with issues:</strong> {', '.join([f'<strong style="color: red;">/{p}/</strong>' for p in problematic_sounds[:10]])}
+                    {f'<span style="color: #666;">(+{len(problematic_sounds) - 10} more)</span>' if len(problematic_sounds) > 10 else ''}
+                </p>
+                <div style='display: flex; flex-wrap: wrap; gap: 10px;'>
+            """
+            
+            for phoneme in problematic_sounds[:10]:
+                error_info = phoneme_errors[phoneme]
+                severity_color = 'red' if error_info['severity'] == 'high' else 'orange' if error_info['severity'] == 'medium' else '#666'
+                summary_html += f"""
+                    <div style='background-color: white; border: 1px solid {severity_color}; border-radius: 4px; padding: 8px; min-width: 120px;'>
+                        <strong style='color: {severity_color}; font-size: 18px;'>/{phoneme}/</strong>
+                        <div style='font-size: 12px; color: #666;'>
+                            {error_info['count']} error(s)<br/>
+                            Types: {', '.join(error_info['types'])}
+                        </div>
+                    </div>
+                """
+            
+            summary_html += """
+                </div>
+            </div>
+            """
+            
+            # Create detailed error table
+            error_table_html = summary_html + """
+            <h4 style='color: #333; margin-top: 20px;'>📋 Detailed Error Report</h4>
+            <table style='width: 100%; border-collapse: collapse; margin: 10px 0; font-size: 13px;'>
+                <thead>
+                    <tr style='background-color: #f0f0f0;'>
+                        <th style='padding: 10px; border: 1px solid #ddd; text-align: left; background-color: #e8e8e8;'>Sound</th>
+                        <th style='padding: 10px; border: 1px solid #ddd; text-align: left; background-color: #e8e8e8;'>Time</th>
+                        <th style='padding: 10px; border: 1px solid #ddd; text-align: left; background-color: #e8e8e8;'>Error Type</th>
+                        <th style='padding: 10px; border: 1px solid #ddd; text-align: left; background-color: #e8e8e8;'>Wrong Sound</th>
+                        <th style='padding: 10px; border: 1px solid #ddd; text-align: left; background-color: #e8e8e8;'>Severity</th>
+                        <th style='padding: 10px; border: 1px solid #ddd; text-align: left; background-color: #e8e8e8;'>Therapy Recommendation</th>
+                    </tr>
+                </thead>
+                <tbody>
+            """
+            
+            for row in error_table_rows[:20]:  # Limit to first 20 errors
+                severity_bg = {
+                    'high': '#ffebee',
+                    'medium': '#fff3e0',
+                    'low': '#f3e5f5',
+                    'none': '#e8f5e9'
+                }.get(row['severity'], '#f5f5f5')
+                
+                error_table_html += f"""
+                    <tr style='background-color: {severity_bg};'>
+                        <td style='padding: 10px; border: 1px solid #ddd;'>
+                            <strong style='color: {row['severity_color']}; font-size: 16px;'>/{row['phoneme']}/</strong>
+                        </td>
+                        <td style='padding: 10px; border: 1px solid #ddd;'>{row['time']}</td>
+                        <td style='padding: 10px; border: 1px solid #ddd;'>
+                            <span style='background-color: {row['severity_color']}; color: white; padding: 3px 8px; border-radius: 3px; font-size: 11px;'>
+                                {row['error_type'].upper()}
+                            </span>
+                        </td>
+                        <td style='padding: 10px; border: 1px solid #ddd;'>
+                            {f"<strong style='color: red;'>/{row['wrong_sound']}/</strong>" if row['wrong_sound'] != 'N/A' else '<span style="color: #999;">N/A</span>'}
+                        </td>
+                        <td style='padding: 10px; border: 1px solid #ddd;'>
+                            <strong style='color: {row['severity_color']};'>{row['severity'].upper()}</strong>
+                        </td>
+                        <td style='padding: 10px; border: 1px solid #ddd; font-size: 12px;'>{row['therapy']}</td>
+                    </tr>
+                """
+            
+            error_table_html += """
+                </tbody>
+            </table>
+            """
+            
+            if len(error_table_rows) > 20:
+                error_table_html += f"<p style='color: #666; font-size: 12px; margin-top: 10px;'>📊 Showing first 20 of <strong>{len(error_table_rows)}</strong> total errors detected</p>"
+        else:
+            error_table_html = """
+            <div style='background-color: #d4edda; border: 2px solid #28a745; border-radius: 8px; padding: 20px; text-align: center;'>
+                <h3 style='color: #155724; margin-top: 0;'>✅ No Errors Detected</h3>
+                <p style='color: #155724; font-size: 16px;'>
+                    All sounds/phonemes were produced correctly!<br/>
+                    <span style='font-size: 14px; color: #666;'>Great job! 🎉</span>
+                </p>
+            </div>
+            """
+        
+        # Create JSON output with errors
         json_output = {
             "status": "success",
             "fluency_metrics": {
@@ -259,6 +442,18 @@ def analyze_speech(
             "confidence": avg_confidence,
             "confidence_percentage": confidence_percentage,
             "processing_time_ms": processing_time_ms,
+            "error_count": len(errors),
+            "errors": [
+                {
+                    "phoneme": err[2].phoneme,
+                    "time": err[1],
+                    "error_type": err[2].error_type.value,
+                    "wrong_sound": err[2].wrong_sound,
+                    "severity": error_mapper.get_severity_level(err[2].severity).value if error_mapper else "unknown",
+                    "therapy": err[2].therapy
+                }
+                for err in errors[:20]
+            ] if errors else [],
             "frame_predictions": [
                 {
                     "time": fp.time,
@@ -266,9 +461,10 @@ def analyze_speech(
                     "fluency_label": fp.fluency_label,
                     "articulation_class": fp.articulation_class,
                     "articulation_label": fp.articulation_label,
-                    "confidence": fp.confidence
+                    "confidence": fp.confidence,
+                    "phoneme": frame_phonemes[i] if i < len(frame_phonemes) else ''
                 }
-                for fp in result.frame_predictions[:20]  # First 20 frames for preview
+                for i, fp in enumerate(result.frame_predictions[:20])  # First 20 frames for preview
             ]
         }
         
@@ -289,6 +485,7 @@ def analyze_speech(
             articulation_text,
             confidence_html,
             processing_time_html,
+            error_table_html,
             json_output
         )
         
@@ -302,11 +499,13 @@ def analyze_speech(
         except: pass
         # #endregion
         error_html = f"<p style='color: red;'>❌ Error: {str(e)}</p>"
+        error_table_html = "<p style='color: #999;'>No error details available</p>"
         return (
             error_html,
             f"Error: {str(e)}",
             "N/A",
             "N/A",
+            error_table_html,
             {"error": str(e), "status": "error"}
         )
 
@@ -370,6 +569,13 @@ def create_gradio_interface(gradio_config: Optional[GradioConfig] = None) -> gr.
                     format="wav"
                 )
                 
+                expected_text_input = gr.Textbox(
+                    label="Expected Text (Optional)",
+                    placeholder="Enter the expected text/transcript for phoneme mapping",
+                    lines=2,
+                    info="Provide the expected text to enable phoneme-level error detection"
+                )
+                
                 analyze_btn = gr.Button(
                     "🔍 Analyze Speech",
                     variant="primary",
@@ -409,6 +615,11 @@ def create_gradio_interface(gradio_config: Optional[GradioConfig] = None) -> gr.
                         elem_classes=["output-box"]
                     )
                 
+                error_table_output = gr.HTML(
+                    label="Error Details",
+                    elem_classes=["output-box"]
+                )
+                
                 json_output = gr.JSON(
                     label="Detailed Results (JSON)",
                     elem_classes=["output-box"]
@@ -417,12 +628,13 @@ def create_gradio_interface(gradio_config: Optional[GradioConfig] = None) -> gr.
         # Set up event handlers
         analyze_btn.click(
             fn=analyze_speech,
-            inputs=[audio_mic, audio_file],
+            inputs=[audio_mic, audio_file, expected_text_input],
             outputs=[
                 fluency_output,
                 articulation_output,
                 confidence_output,
                 processing_time_output,
+                error_table_output,
                 json_output
             ]
         )