Training in progress - step 1000

asr_config.py

@@ -4,139 +4,120 @@ import transformers
 
 
 class ASRConfig(transformers.PretrainedConfig):
-    """Configuration class for the ASR model.
-
-    This config combines settings for:
-    - Audio encoder (GLM-ASR/Whisper)
-    - Text decoder (Qwen)
-    - Projector (MLP, MOSA, MoE, QFormer)
-    - Generation parameters
-    - Training options (SpecAugment, LoRA)
-    """
+    """Configuration class for the ASR model."""
 
     model_type = "asr_model"
     is_composition = True
 
+    # Generation defaults
+    GENERATION_DEFAULTS = {
+        "num_beams": 1,
+        "max_new_tokens": 128,
+        "min_new_tokens": 0,
+        "repetition_penalty": 1.0,
+        "length_penalty": 1.0,
+        "no_repeat_ngram_size": 0,
+        "use_cache": True,
+        "do_sample": False,
+        "temperature": None,
+        "top_p": None,
+        "top_k": None,
+    }
+
     def __init__(
         self,
+        # Model IDs
         audio_model_id: str = "zai-org/GLM-ASR-Nano-2512",
         text_model_id: str = "Qwen/Qwen3-0.6B",
+        # Model settings
         attn_implementation: str = "sdpa",
         model_dtype: str = "bfloat16",
-        num_beams: Optional[int] = None,
         system_prompt: str = "You are a helpful assistant.",
+        enable_thinking: bool = False,
+        # Encoder settings (auto-detected if None)
         encoder_dim: Optional[int] = None,
         llm_dim: Optional[int] = None,
-        # Encoder conv layers: list of (padding, kernel_size, stride) tuples
-        # Default is Whisper/GLM-ASR structure: conv1(k=3,s=1,p=1) + conv2(k=3,s=2,p=1)
         encoder_conv_layers: Optional[list] = None,
         audio_sample_rate: int = 16000,
+        # Projector settings
+        projector_type: str = "mlp",
         projector_pool_stride: int = 4,
-        downsample_rate: int = 5,  # Granite default
         projector_hidden_dim: Optional[int] = None,
-        projector_type: str = "mlp",  # "mlp", "mosa", "moe", "qformer"
-        projector_num_layers: int = 2,  # Number of layers in MLP projector
-        projector_init_std: float = 0.02,  # Weight initialization std
-        projector_dropout: float = 0.0,  # Dropout rate for projector layers
-        # MoE-specific configuration
-        num_experts: int = 4,  # Number of experts in MoE projectors
-        num_experts_per_tok: int = 2,  # Top-k experts per token
-        router_aux_loss_coef: float = 0.01,  # Auxiliary loss coefficient for load balancing
-        # QFormer-specific configuration (Granite defaults)
-        qformer_window_size: int = 15,  # Window size for QFormer processing
-        qformer_hidden_size: Optional[int] = None,  # QFormer hidden size (defaults to encoder_dim)
-        qformer_num_layers: int = 2,  # Number of QFormer transformer layers
-        qformer_num_heads: int = 16,  # Number of attention heads in QFormer
-        qformer_intermediate_size: Optional[int] = None,  # FFN size (defaults to 4x hidden)
-        label_smoothing: float = 0.0,  # Label smoothing for cross-entropy loss
-        inference_warmup_tokens: int = 10,
-        # SpecAugment settings
+        projector_num_layers: int = 2,
+        projector_init_std: float = 0.02,
+        projector_dropout: float = 0.0,
+        # MoE projector settings
+        num_experts: int = 4,
+        num_experts_per_tok: int = 2,
+        router_aux_loss_coef: float = 0.01,
+        # QFormer projector settings
+        qformer_window_size: int = 15,
+        qformer_hidden_size: Optional[int] = None,
+        qformer_num_layers: int = 2,
+        qformer_num_heads: int = 16,
+        qformer_intermediate_size: Optional[int] = None,
+        downsample_rate: int = 5,
+        # Training settings (not saved to config.json for inference)
         use_specaugment: bool = False,
         num_time_masks: int = 2,
         time_mask_length: int = 10,
         num_freq_masks: int = 0,
         freq_mask_length: int = 10,
-        # LoRA configuration (for Stage 2 fine-tuning)
         use_lora: bool = False,
-        lora_rank: int = 8,  # SALMONN default
-        lora_alpha: int = 32,  # SALMONN default (scaling factor 4.0)
+        lora_rank: int = 8,
+        lora_alpha: int = 32,
         lora_dropout: float = 0.0,
-        lora_target_modules: Optional[list] = None,  # Default: all linear layers
-        freeze_projector: bool = False,  # True for Stage 2 (LoRA-only training)
-        do_sample: bool = False,
-        enable_thinking: bool = False,  # Enable Qwen3 thinking mode for omni models
-        temperature: Optional[float] = None,
-        top_p: Optional[float] = None,
-        top_k: Optional[int] = None,
-        max_new_tokens: Optional[int] = None,
-        min_new_tokens: Optional[int] = None,
-        repetition_penalty: Optional[float] = None,
-        length_penalty: Optional[float] = None,
-        no_repeat_ngram_size: Optional[int] = None,
-        use_cache: Optional[bool] = None,
+        lora_target_modules: Optional[list] = None,
+        freeze_projector: bool = False,
+        label_smoothing: float = 0.0,
         **kwargs,
     ):
-        """Initialize ASR model configuration.
-
-        Args:
-            audio_model_id: HuggingFace model ID for audio encoder (GLM-ASR/Whisper)
-            text_model_id: HuggingFace model ID for text decoder (Qwen)
-            attn_implementation: Attention implementation ("sdpa", "flash_attention_2", "eager")
-            model_dtype: Model dtype ("bfloat16", "float16", "float32")
-            projector_type: Projector architecture ("mlp", "mosa", "moe", "qformer")
-            use_lora: Enable LoRA adapters for Stage 2 fine-tuning
-            use_specaugment: Enable SpecAugment data augmentation
-        """
-        # Set default generation parameters (greedy decoding only)
-        generation_defaults = {
-            "num_beams": 1,
-            "max_new_tokens": 128,
-            "min_new_tokens": 0,
-            "repetition_penalty": 1.0,
-            "length_penalty": 1.0,
-            "no_repeat_ngram_size": 0,  # Prevent repeating 3-grams like "so so so"
-            "use_cache": True,
-        }
-
-        # Apply defaults (config.json values take precedence)
-        kwargs = {**generation_defaults, **kwargs}
+        # Merge generation defaults with kwargs (kwargs takes precedence)
+        for key, default in self.GENERATION_DEFAULTS.items():
+            if key not in kwargs:
+                kwargs[key] = default
 
+        # Core model settings
         self.audio_model_id = audio_model_id
         self.text_model_id = text_model_id
         self.attn_implementation = attn_implementation
         self.model_dtype = model_dtype
         self.system_prompt = system_prompt
+        self.enable_thinking = enable_thinking
+
+        # Encoder settings
         self.encoder_dim = encoder_dim
         self.llm_dim = llm_dim
-        # Default conv layers for Whisper/GLM-ASR: [(pad, kernel, stride), ...]
         self.encoder_conv_layers = encoder_conv_layers or [(1, 3, 1), (1, 3, 2)]
         self.audio_sample_rate = audio_sample_rate
-        self.projector_init_std = projector_init_std
+
+        # Projector settings
+        self.projector_type = projector_type
         self.projector_pool_stride = projector_pool_stride
-        self.downsample_rate = downsample_rate
         self.projector_hidden_dim = projector_hidden_dim
-        self.projector_type = projector_type
         self.projector_num_layers = projector_num_layers
+        self.projector_init_std = projector_init_std
         self.projector_dropout = projector_dropout
-        # MoE-specific configuration
+
+        # MoE settings
         self.num_experts = num_experts
         self.num_experts_per_tok = num_experts_per_tok
         self.router_aux_loss_coef = router_aux_loss_coef
-        # QFormer-specific configuration
+
+        # QFormer settings
         self.qformer_window_size = qformer_window_size
         self.qformer_hidden_size = qformer_hidden_size
         self.qformer_num_layers = qformer_num_layers
         self.qformer_num_heads = qformer_num_heads
         self.qformer_intermediate_size = qformer_intermediate_size
-        self.label_smoothing = label_smoothing
-        self.inference_warmup_tokens = inference_warmup_tokens
-        # SpecAugment configuration
+        self.downsample_rate = downsample_rate
+
+        # Training settings
         self.use_specaugment = use_specaugment
         self.num_time_masks = num_time_masks
         self.time_mask_length = time_mask_length
         self.num_freq_masks = num_freq_masks
         self.freq_mask_length = freq_mask_length
-        # LoRA configuration
         self.use_lora = use_lora
         self.lora_rank = lora_rank
         self.lora_alpha = lora_alpha
@@ -151,69 +132,48 @@ class ASRConfig(transformers.PretrainedConfig):
             "down_proj",
         ]
         self.freeze_projector = freeze_projector
+        self.label_smoothing = label_smoothing
 
-        # Generation parameters (use explicit value if provided, else use default)
-        self.num_beams = num_beams if num_beams is not None else generation_defaults["num_beams"]
-        self.max_new_tokens = (
-            max_new_tokens if max_new_tokens is not None else generation_defaults["max_new_tokens"]
-        )
-        self.min_new_tokens = (
-            min_new_tokens if min_new_tokens is not None else generation_defaults["min_new_tokens"]
-        )
-        self.repetition_penalty = (
-            repetition_penalty
-            if repetition_penalty is not None
-            else generation_defaults["repetition_penalty"]
-        )
-        self.length_penalty = (
-            length_penalty if length_penalty is not None else generation_defaults["length_penalty"]
-        )
-        self.no_repeat_ngram_size = (
-            no_repeat_ngram_size
-            if no_repeat_ngram_size is not None
-            else generation_defaults["no_repeat_ngram_size"]
-        )
-        self.use_cache = use_cache if use_cache is not None else generation_defaults["use_cache"]
-        self.do_sample = do_sample
-        self.enable_thinking = enable_thinking
-        self.temperature = temperature
-        self.top_p = top_p
-        self.top_k = top_k
-
-        if "audio_config" not in kwargs:
+        # Generation parameters (from kwargs after merge with defaults)
+        self.num_beams = kwargs.pop("num_beams")
+        self.max_new_tokens = kwargs.pop("max_new_tokens")
+        self.min_new_tokens = kwargs.pop("min_new_tokens")
+        self.repetition_penalty = kwargs.pop("repetition_penalty")
+        self.length_penalty = kwargs.pop("length_penalty")
+        self.no_repeat_ngram_size = kwargs.pop("no_repeat_ngram_size")
+        self.use_cache = kwargs.pop("use_cache")
+        self.do_sample = kwargs.pop("do_sample")
+        self.temperature = kwargs.pop("temperature")
+        self.top_p = kwargs.pop("top_p")
+        self.top_k = kwargs.pop("top_k")
+
+        # Load sub-configs
+        self.audio_config = kwargs.pop("audio_config", None)
+        if self.audio_config is None:
             self.audio_config = transformers.AutoConfig.from_pretrained(audio_model_id)
-            # Override dtype to match model_dtype
             self.audio_config.dtype = model_dtype
-        else:
-            self.audio_config = kwargs.pop("audio_config")
-
-        if "text_config" not in kwargs:
+        elif isinstance(self.audio_config, dict) and self.audio_config.get("model_type"):
+            config_class = transformers.AutoConfig.for_model(
+                self.audio_config["model_type"]
+            ).__class__
+            self.audio_config = config_class(**self.audio_config)
+
+        self.text_config = kwargs.pop("text_config", None)
+        if self.text_config is None:
             self.text_config = transformers.AutoConfig.from_pretrained(
                 text_model_id, trust_remote_code=True
             )
-            # Override dtype to match model_dtype
             self.text_config.dtype = model_dtype
-        else:
-            self.text_config = kwargs.pop("text_config")
-
-        if isinstance(self.text_config, dict):
-            # Reconstruct config from dict using the model_type stored in the dict
-            model_type = self.text_config["model_type"]
-            config_class = transformers.AutoConfig.for_model(model_type).__class__
+        elif isinstance(self.text_config, dict):
+            config_class = transformers.AutoConfig.for_model(
+                self.text_config["model_type"]
+            ).__class__
             self.text_config = config_class(**self.text_config)
 
-        if isinstance(self.audio_config, dict):
-            model_type = self.audio_config.get("model_type")
-            if model_type:
-                config_class = transformers.AutoConfig.for_model(model_type).__class__
-                self.audio_config = config_class(**self.audio_config)
-
         super().__init__(**kwargs)
 
-        # Point encoder to audio_config so pipeline uses correct feature extractor
-        # The pipeline looks for config.encoder._name_or_path for feature extractor
+        # Pipeline configuration
         self.encoder = self.audio_config
-
         self.auto_map = {
             "AutoConfig": "asr_config.ASRConfig",
             "AutoModel": "asr_modeling.ASRModel",

asr_pipeline.py

@@ -18,7 +18,26 @@ except ImportError:
     from diarization import SpeakerDiarizer  # type: ignore[no-redef]
 
 # Re-export for backwards compatibility
-__all__ = ["ForcedAligner", "SpeakerDiarizer", "ASRPipeline"]
+__all__ = ["ForcedAligner", "SpeakerDiarizer", "ASRPipeline", "strip_thinking"]
+
+# Default TTS voice for Kokoro
+DEFAULT_TTS_VOICE = "af_heart"
+TTS_SAMPLE_RATE = 24000
+
+
+def strip_thinking(text: str) -> str:
+    """Remove <think>...</think> tags from model output.
+
+    Args:
+        text: Model output text that may contain thinking tags
+
+    Returns:
+        Text with thinking content removed
+    """
+    if not text:
+        return text
+    text = re.sub(r"<think>.*?</think>\s*", "", text, flags=re.DOTALL)
+    return text.strip()
 
 
 class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
@@ -43,6 +62,44 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             model=model, feature_extractor=feature_extractor, tokenizer=tokenizer, **kwargs
         )
         self._current_audio = None
+        self._tts_pipeline = None
+
+    @property
+    def tts_pipeline(self):
+        """Lazy-load Kokoro TTS pipeline on first use."""
+        if self._tts_pipeline is None:
+            try:
+                from kokoro import KPipeline
+
+                self._tts_pipeline = KPipeline(lang_code="a", repo_id="hexgrad/Kokoro-82M")
+            except ImportError as e:
+                raise ImportError(
+                    "Kokoro TTS is required for audio output. "
+                    "Install with: pip install kokoro>=0.9.2\n"
+                    "Also requires espeak-ng: apt-get install espeak-ng"
+                ) from e
+        return self._tts_pipeline
+
+    def text_to_speech(self, text: str, voice: str = DEFAULT_TTS_VOICE) -> dict[str, Any]:
+        """Convert text to speech using Kokoro TTS.
+
+        Args:
+            text: Text to synthesize
+            voice: Kokoro voice ID (default: "af_heart")
+
+        Returns:
+            Dict with 'audio' (numpy array) and 'sample_rate' keys
+        """
+        if not text or not text.strip():
+            return {"audio": np.array([], dtype=np.float32), "sample_rate": TTS_SAMPLE_RATE}
+
+        # Generate audio chunks and concatenate
+        audio_chunks = []
+        for _, _, audio in self.tts_pipeline(text, voice=voice):
+            audio_chunks.append(audio)
+
+        audio = np.concatenate(audio_chunks) if audio_chunks else np.array([], dtype=np.float32)
+        return {"audio": audio, "sample_rate": TTS_SAMPLE_RATE}
 
     def _sanitize_parameters(self, **kwargs):
         """Intercept our custom parameters before parent class validates them."""
@@ -55,6 +112,9 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         kwargs.pop("hf_token", None)
         kwargs.pop("user_prompt", None)
         kwargs.pop("diarization_backend", None)
+        # TTS parameters
+        kwargs.pop("return_audio", None)
+        kwargs.pop("tts_voice", None)
 
         return super()._sanitize_parameters(**kwargs)
 
@@ -69,6 +129,8 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             inputs: Audio input (file path, dict with array/sampling_rate, etc.)
             return_timestamps: If True, return word-level timestamps using forced alignment
             return_speakers: If True, return speaker labels for each word
+            return_audio: If True, synthesize transcription as speech using Kokoro TTS
+            tts_voice: Kokoro voice ID for TTS output (default: "af_heart")
             user_prompt: Custom transcription prompt (default: "Transcribe: ")
             num_speakers: Exact number of speakers (if known, for diarization)
             min_speakers: Minimum number of speakers (for diarization)
@@ -77,11 +139,14 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
 
         Returns:
             Dict with 'text' key, 'words' key if return_timestamps=True,
-            and speaker labels on words if return_speakers=True
+            speaker labels on words if return_speakers=True,
+            and 'audio'/'sample_rate' keys if return_audio=True
         """
         # Extract our params before super().__call__ (which will also call _sanitize_parameters)
         return_timestamps = kwargs.pop("return_timestamps", False)
         return_speakers = kwargs.pop("return_speakers", False)
+        return_audio = kwargs.pop("return_audio", False)
+        tts_voice = kwargs.pop("tts_voice", DEFAULT_TTS_VOICE)
         user_prompt = kwargs.pop("user_prompt", None)
         diarization_params = {
             "num_speakers": kwargs.pop("num_speakers", None),
@@ -143,6 +208,18 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
                 result["speaker_segments"] = []
                 result["diarization_error"] = str(e)
 
+        # Synthesize transcription as speech if requested
+        if return_audio:
+            text = result.get("text", "")
+            try:
+                tts_result = self.text_to_speech(text, voice=tts_voice)
+                result["audio"] = tts_result["audio"]
+                result["sample_rate"] = tts_result["sample_rate"]
+            except Exception as e:
+                result["audio"] = np.array([], dtype=np.float32)
+                result["sample_rate"] = TTS_SAMPLE_RATE
+                result["tts_error"] = str(e)
+
         # Clean up
         self._current_audio = None
         if original_prompt is not None:
@@ -257,7 +334,7 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
 
         text = self.tokenizer.decode(tokens, skip_special_tokens=True).strip()
         # Strip <think>...</think> tags (Qwen3 doesn't respect /no_think prompt)
-        text = re.sub(r"<think>.*?</think>\s*", "", text, flags=re.DOTALL).strip()
+        text = strip_thinking(text)
         # Truncate repetitions at end of text
         text = _truncate_repetitions(text)
         return {"text": text}