Training in progress - step 1000

asr_pipeline.py

@@ -485,6 +485,8 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         if not text:
             return ""
 
+        original_len = len(text.split())
+
         # 1. LOWERCASE
         text = text.lower()
 
@@ -502,6 +504,7 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             if repeat_count >= 1:
                 words = words[: idx + n]
                 text = " ".join(words)
+                print(f"[DEBUG] Truncated repetition: {original_len} -> {len(words)} words (n={n}, repeats={repeat_count})")
                 break
 
         # 3. COMBINE ACRONYMS

projectors.py

@@ -1,9 +1,10 @@
 """Audio projector modules for bridging encoder and decoder embeddings.
 
 This module contains all projector architectures:
-- MLPAudioProjector: Simple 2-layer MLP with frame stacking downsampling
-- MOSAProjector: MOSA-style dense mixture of experts
-- SharedMoEAudioProjector: Shared expert + sparse routed experts
+- LinearProjector: Simple avg pool + linear (Chinese Dialects paper, best for Stage 1)
+- MLPAudioProjector: 2-layer MLP with frame stacking downsampling
+- MOSAProjector: MOSA-style dense mixture of experts (arXiv:2508.18998)
+- MoEAudioProjector: Shared expert + sparse routed experts
 - QFormerAudioProjector: BLIP-2 QFormer with learnable queries (Granite-style)
 """
 
@@ -15,6 +16,51 @@ import torch.nn.functional as F  # noqa: N812
 from transformers import AutoModel, Blip2QFormerConfig
 from transformers.models.llama.modeling_llama import LlamaRMSNorm
 
+# =============================================================================
+# Linear Projector (Chinese Dialects paper style)
+# =============================================================================
+
+
+class LinearProjector(nn.Module):
+    """Simple linear projector with average pooling downsampling.
+
+    Based on Chinese Dialects paper (arXiv:2505.21138) which found this
+    outperformed Conv1D, Transformer, and Q-Former in Stage 1 (projector-only).
+
+    Architecture: AvgPool(4x) -> Linear(encoder_dim, llm_dim)
+    """
+
+    def __init__(self, config):
+        super().__init__()
+
+        self.encoder_dim = getattr(config, "encoder_dim", 768)
+        self.llm_dim = getattr(config, "llm_dim", 2048)
+        self.pool_stride = getattr(config, "projector_pool_stride", 4)  # 4x = 12.5Hz
+
+        # Single linear projection (no hidden layers, no activation)
+        self.linear = nn.Linear(self.encoder_dim, self.llm_dim)
+
+    def get_output_length(self, input_length: int) -> int:
+        """Calculate output sequence length given input length."""
+        return input_length // self.pool_stride
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: [batch, seq_len, encoder_dim]
+
+        Returns:
+            [batch, seq_len // pool_stride, llm_dim]
+        """
+        # Average pooling for downsampling (better than frame stacking for linear)
+        # Transpose for avg_pool1d: [B, S, D] -> [B, D, S]
+        x = x.transpose(1, 2)
+        x = F.avg_pool1d(x, kernel_size=self.pool_stride, stride=self.pool_stride)
+        x = x.transpose(1, 2)  # [B, S//k, D]
+
+        return self.linear(x)
+
+
 # =============================================================================
 # MLP Projector
 # =============================================================================
@@ -63,12 +109,12 @@ class MLPAudioProjector(nn.Module):
 
 
 class SimpleAdapter(nn.Module):
-    """Simple 2-layer GELU adapter (from MOSA paper)."""
+    """Simple 2-layer ReLU adapter (from MOSA paper, arXiv:2508.18998)."""
 
     def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
         super().__init__()
         self.fc1 = nn.Linear(input_dim, hidden_dim)
-        self.act = nn.GELU()
+        self.act = nn.ReLU()
         self.fc2 = nn.Linear(hidden_dim, output_dim)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -93,24 +139,27 @@ class MOSAProjector(nn.Module):
 
     Based on "MOSA: Mixtures of Simple Adapters" (arXiv:2508.18998).
     Uses softmax gating over all experts (dense MoE) with only cross-entropy loss.
-    Uses frame-stacking for downsampling (like MLP projector).
+    Uses conv-based downsampling (2x Conv1d stride-2) as described in the paper.
     """
 
     def __init__(self, config):
         super().__init__()
         self.encoder_dim = getattr(config, "encoder_dim", None) or 1280
         self.llm_dim = getattr(config, "llm_dim", None) or 2048
-        self.k = getattr(config, "projector_pool_stride", 4)
         self.num_experts = getattr(config, "num_experts", None) or 4  # MOSA-Base uses 4
         adapter_hidden = getattr(config, "adapter_hidden_dim", None) or 4096
 
-        # Optional pre-norm before projection
-        self.use_pre_norm = getattr(config, "projector_pre_norm", False)
-        if self.use_pre_norm:
-            self.pre_norm = LlamaRMSNorm(self.encoder_dim, eps=1e-8)
+        # --- Conv-based downsampling (paper: 2x Conv1d, kernel=3, stride=2) ---
+        # Total 4x downsampling: 50Hz -> 12.5Hz
+        self.conv_downsample = nn.Sequential(
+            nn.Conv1d(self.encoder_dim, self.encoder_dim, kernel_size=3, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Conv1d(self.encoder_dim, self.encoder_dim, kernel_size=3, stride=2, padding=1),
+            nn.ReLU(),
+        )
 
-        # Frame stacking: concat k adjacent frames then project
-        in_dim = self.encoder_dim * self.k
+        # Input dim to adapters is now just encoder_dim (not encoder_dim * k)
+        in_dim = self.encoder_dim
 
         # --- 1. Simple Router (MOSA-Base: 2 layers with ReLU) ---
         # Maps encoder_dim -> 512 -> num_experts
@@ -122,7 +171,7 @@ class MOSAProjector(nn.Module):
         )
 
         # --- 2. Experts (Simple 2-layer GELU adapters) ---
-        # Each expert: in_dim (stacked frames) -> hidden -> llm_dim
+        # Each expert: encoder_dim -> hidden -> llm_dim
         self.experts = nn.ModuleList(
             [SimpleAdapter(in_dim, adapter_hidden, self.llm_dim) for _ in range(self.num_experts)]
         )
@@ -131,31 +180,27 @@ class MOSAProjector(nn.Module):
         # x: (B, S, encoder_dim)
         batch_size, seq_len, dim = x.shape
 
-        # Apply pre-norm if enabled
-        if self.use_pre_norm:
-            x = self.pre_norm(x)
-
-        # --- 1. Router Branch ---
-        # Mean pool encoder outputs for routing decisions
-        x_pooled = x.reshape(batch_size, -1, self.k, self.encoder_dim).mean(dim=2)  # (B, S//k, D)
+        # --- 1. Conv downsampling ---
+        # Conv1d expects (B, C, S), so transpose
+        x_conv = x.transpose(1, 2)  # (B, D, S)
+        x_conv = self.conv_downsample(x_conv)  # (B, D, S//4)
+        x_downsampled = x_conv.transpose(1, 2)  # (B, S//4, D)
 
-        # Router logits and softmax gating (dense MoE)
-        routing_weights = F.softmax(self.router(x_pooled), dim=-1)  # (B, S//k, num_experts)
-
-        # --- 2. Frame stacking for experts ---
-        # Reshape to combine k frames: [B, S, D] -> [B, S//k, D*k]
-        x_stacked = x.reshape(batch_size, -1, dim * self.k)
+        # --- 2. Router Branch ---
+        # Router operates on downsampled features
+        routing_weights = F.softmax(self.router(x_downsampled), dim=-1)  # (B, S//4, num_experts)
 
         # --- 3. Expert Mixture (Dense Execution) ---
         # Run all experts and compute weighted sum
         expert_outputs = torch.stack(
-            [expert(x_stacked) for expert in self.experts]
-        )  # (E, B, S//k, D)
+            [expert(x_downsampled) for expert in self.experts]
+        )  # (E, B, S//4, D)
         return torch.einsum("ebsd, bse -> bsd", expert_outputs, routing_weights)
 
     def get_output_length(self, input_length: int) -> int:
         """Calculate output sequence length given input length."""
-        return input_length // self.k
+        # Two stride-2 convs = 4x downsampling
+        return input_length // 4
 
 
 # =============================================================================
@@ -452,6 +497,7 @@ class QFormerAudioProjector(nn.Module):
 # =============================================================================
 
 PROJECTOR_CLASSES = {
+    "linear": LinearProjector,
     "mlp": MLPAudioProjector,
     "mosa": MOSAProjector,
     "moe": MoEAudioProjector,