Upload 3 files

childfluency.py

@@ -0,0 +1,140 @@
+# src/models/childfluency.py
+"""
+ChildFluency-Net — Full Model
+
+Combines WavLM acoustic encoder + LPE linguistic encoder
+into a unified architecture with dual SLD/OD classification heads.
+
+Architecture:
+    [WavLM → 256-dim] + [LPE → 32-dim]
+          ↓ concat [288-dim]
+        Shared MLP [64-dim]
+          ↓
+    [SLD Head]   [OD Head]
+          ↓
+    Mean pool across windows → SLD rate % (clinical output)
+
+Clinical output:
+    SLD rate > 3% → refer for specialist assessment
+"""
+
+import torch
+import torch.nn as nn
+from typing import Dict
+from wavlm_encoder import WavLMEncoder
+from lpe_module import LPEModule
+
+
+class ChildFluencyNet(nn.Module):
+
+    def __init__(
+        self,
+        wavlm_name   : str   = "microsoft/wavlm-large",
+        acoustic_dim : int   = 256,
+        lpe_dim      : int   = 32,
+        hidden_dim   : int   = 128,
+        dropout      : float = 0.3
+    ):
+        super().__init__()
+
+        # Branch A — acoustic
+        self.acoustic_encoder = WavLMEncoder(
+            model_name = wavlm_name,
+            output_dim = acoustic_dim
+        )
+
+        # Branch B — linguistic position
+        self.lpe_module = LPEModule(
+            input_dim  = 6,
+            output_dim = lpe_dim
+        )
+
+        fusion_dim = acoustic_dim + lpe_dim   # 256 + 32 = 288
+
+        # Shared MLP after fusion
+        self.shared_mlp = nn.Sequential(
+            nn.Linear(fusion_dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, hidden_dim // 2)  # → 64-dim
+        )
+
+        # Head 1 — SLD (Stuttering-Like Disfluencies)
+        # Detects: blocks, prolongations, sound/syllable repetitions
+        # These are PATHOLOGICAL — what clinicians diagnose
+        self.sld_head = nn.Linear(hidden_dim // 2, 1)
+
+        # Head 2 — OD (Other Disfluencies)
+        # Detects: word repetitions, interjections (um, uh)
+        # These are NORMAL — present in all speakers
+        self.od_head  = nn.Linear(hidden_dim // 2, 1)
+
+    def forward(
+        self,
+        audio        : torch.Tensor,   # [batch, 64000]
+        lpe_features : torch.Tensor    # [batch, 6]
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Args:
+            audio        : raw waveform, 4 seconds at 16kHz
+            lpe_features : 6 linguistic position features per window
+
+        Returns dict with:
+            sld_logit : [batch]   raw score for SLD detection
+            od_logit  : [batch]   raw score for OD detection
+            embedding : [batch, 64]  shared representation
+        """
+        # Branch A
+        acoustic_emb = self.acoustic_encoder(audio)          # [batch, 256]
+
+        # Branch B
+        lpe_emb = self.lpe_module(lpe_features)              # [batch, 32]
+
+        # Fuse
+        fused  = torch.cat([acoustic_emb, lpe_emb], dim=-1)  # [batch, 288]
+        shared = self.shared_mlp(fused)                       # [batch, 64]
+
+        return {
+            "sld_logit" : self.sld_head(shared).squeeze(-1),  # [batch]
+            "od_logit"  : self.od_head(shared).squeeze(-1),   # [batch]
+            "embedding" : shared                               # [batch, 64]
+        }
+
+    @torch.no_grad()
+    def predict_recording(self, window_outputs: list) -> Dict:
+        """
+        Aggregates window-level predictions into recording-level diagnosis.
+        Called during inference on a full recording.
+
+        Args:
+            window_outputs : list of dicts from forward(), one per window
+
+        Returns:
+            sld_rate    : float  SLD rate % (clinical diagnostic metric)
+            has_stutter : bool   True if SLD rate > 3%
+            sld_timeline: list   per-window SLD probability (for visualization)
+        """
+        sld_probs = torch.sigmoid(
+            torch.stack([o["sld_logit"] for o in window_outputs])
+        )
+        od_probs = torch.sigmoid(
+            torch.stack([o["od_logit"] for o in window_outputs])
+        )
+
+        sld_rate = (sld_probs > 0.5).float().mean().item() * 100.0
+
+        return {
+            "sld_rate"     : round(sld_rate, 2),
+            "has_stutter"  : sld_rate > 3.0,
+            "mean_sld_prob": round(sld_probs.mean().item(), 4),
+            "mean_od_prob" : round(od_probs.mean().item(), 4),
+            "sld_timeline" : sld_probs.cpu().tolist()
+        }
+
+    def count_params(self):
+        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        total     = sum(p.numel() for p in self.parameters())
+        print(f"ChildFluencyNet")
+        print(f"  Trainable : {trainable:,}")
+        print(f"  Total     : {total:,}")
+        print(f"  Frozen    : {total - trainable:,}")

lpe_module.py

@@ -0,0 +1,58 @@
+# src/models/lpe_module.py
+"""
+Linguistic Position Encoding (LPE) Module
+
+Core novel contribution of ChildFluency-Net.
+
+Encodes 6 loci-of-stuttering features into a 32-dim vector.
+These features are pre-computed by build_features.py using
+WhisperX transcription + spaCy linguistic parsing.
+
+Clinical basis (Bloodstein 1960, Bernstein-Ratner 1997):
+    Stuttering is 3-5x more likely on:
+    - Sentence-initial positions
+    - Content words (nouns, verbs) vs function words
+    - Longer, rarer words
+    - Near clause boundaries (cognitive planning spikes)
+    - In longer, more complex utterances (higher MLU)
+
+Input features [6 scalars, all in 0-1 range]:
+    0: sentence position     (0=initial=high risk, 1=final)
+    1: word class            (1=content, 0=function)
+    2: syllable count        (normalized, max=6 syllables)
+    3: clause boundary prox  (1=near boundary)
+    4: word rarity           (1=rare=high risk, 0=common)
+    5: MLU                   (sentence length / 20)
+"""
+
+import torch
+import torch.nn as nn
+
+
+class LPEModule(nn.Module):
+
+    def __init__(
+        self,
+        input_dim  : int = 6,
+        output_dim : int = 32
+    ):
+        super().__init__()
+
+        # Small MLP: 6 → 16 → 32
+        # Kept intentionally small — 6 features don't need
+        # a large network, just a learned non-linear transformation
+        self.encoder = nn.Sequential(
+            nn.Linear(input_dim, 16),
+            nn.ReLU(),
+            nn.Linear(16, output_dim)
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x : [batch, 6]   pre-computed LPE feature vectors
+
+        Returns:
+              : [batch, 32]  LPE embeddings
+        """
+        return self.encoder(x)

wavlm_encoder.py

@@ -0,0 +1,71 @@
+# src/models/wavlm_encoder.py
+"""
+WavLM Acoustic Encoder
+
+Takes raw 4-second audio → outputs 256-dim acoustic embedding.
+
+WavLM-Large is chosen over Wav2Vec2 because:
+- Trained with masked speech prediction + denoising objective
+- More robust to background noise in clinical recordings
+- State-of-the-art on SUPERB benchmark across all speech tasks
+"""
+
+import torch
+import torch.nn as nn
+from transformers import WavLMModel
+
+
+class WavLMEncoder(nn.Module):
+
+    def __init__(
+        self,
+        model_name : str = "microsoft/wavlm-large",
+        output_dim : int = 256,
+        freeze_cnn : bool = True
+    ):
+        super().__init__()
+
+        # Load pretrained WavLM-Large
+        self.wavlm = WavLMModel.from_pretrained(model_name)
+
+        # Freeze the CNN feature extractor (low-level, no benefit fine-tuning)
+        # Only the transformer layers get fine-tuned
+        if freeze_cnn:
+            for param in self.wavlm.feature_extractor.parameters():
+                param.requires_grad = False
+
+        # WavLM-Large hidden size = 1024
+        wavlm_dim = self.wavlm.config.hidden_size
+
+        # Project 1024 → 256
+        self.projection = nn.Sequential(
+            nn.Linear(wavlm_dim, 512),
+            nn.GELU(),
+            nn.Dropout(0.1),
+            nn.Linear(512, output_dim)
+        )
+
+    def forward(self, audio: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            audio : [batch, time_samples]  raw waveform at 16kHz
+                    4 seconds = 64,000 samples
+
+        Returns:
+            embedding : [batch, 256]
+        """
+        # WavLM transformer layers → frame-level features
+        out    = self.wavlm(input_values=audio)
+        hidden = out.last_hidden_state        # [batch, frames, 1024]
+
+        # Mean pool all frames into one vector per clip
+        pooled = hidden.mean(dim=1)           # [batch, 1024]
+
+        # Project to output dim
+        return self.projection(pooled)        # [batch, 256]
+
+    def trainable_params(self):
+        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        total     = sum(p.numel() for p in self.parameters())
+        print(f"WavLMEncoder — trainable: {trainable:,} / total: {total:,} "
+              f"({100*trainable/total:.1f}%)")