Training in progress - step 25000

.gitattributes

@@ -1,3 +1,4 @@
 *.safetensors filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
 tokenizer_config.json -filter -diff -merge text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

asr_config.py

@@ -70,12 +70,6 @@ class ASRConfig(transformers.PretrainedConfig):
         lora_target_modules: Optional[list] = None,
         freeze_projector: bool = False,
         label_smoothing: float = 0.0,
-        # Audio Head settings (Freeze-Omni style AR decoder)
-        use_audio_head: bool = False,
-        audio_head_hidden_dim: int = 512,  # AR decoder hidden dimension
-        codebook_size: int = 2048,  # Mimi codec vocabulary size
-        num_codebooks: int = 1,  # Number of codebooks to predict (first 1-2 most important)
-        freeze_audio_head: bool = False,  # Freeze entire audio head
         **kwargs,
     ):
         # Merge generation defaults with kwargs (kwargs takes precedence)
@@ -140,13 +134,6 @@ class ASRConfig(transformers.PretrainedConfig):
         self.freeze_projector = freeze_projector
         self.label_smoothing = label_smoothing
 
-        # Audio Head settings (Freeze-Omni style AR decoder)
-        self.use_audio_head = use_audio_head
-        self.audio_head_hidden_dim = audio_head_hidden_dim
-        self.codebook_size = codebook_size
-        self.num_codebooks = num_codebooks
-        self.freeze_audio_head = freeze_audio_head
-
         # Generation parameters (from kwargs after merge with defaults)
         self.num_beams = kwargs.pop("num_beams")
         self.max_new_tokens = kwargs.pop("max_new_tokens")
@@ -163,9 +150,7 @@ class ASRConfig(transformers.PretrainedConfig):
         # 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, trust_remote_code=True
-            )
+            self.audio_config = transformers.AutoConfig.from_pretrained(audio_model_id)
             self.audio_config.dtype = model_dtype
         elif isinstance(self.audio_config, dict) and self.audio_config.get("model_type"):
             config_class = transformers.AutoConfig.for_model(

asr_modeling.py

@@ -181,19 +181,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         else:
             self.spec_augment = None
 
-        # Audio Head for S2S (trainable)
-        if getattr(config, "use_audio_head", False):
-            from .audio_head import AudioHead
-
-            self.audio_head = AudioHead(config).to(
-                device=next(self.language_model.parameters()).device,
-                dtype=target_dtype,
-            )
-            if getattr(config, "freeze_audio_head", False):
-                self.audio_head.requires_grad_(False)
-        else:
-            self.audio_head = None
-
         # For model parallelism
         self._no_split_modules = getattr(self.language_model, "_no_split_modules", [])
 
@@ -378,11 +365,8 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         )
 
     def state_dict(self, *args, **kwargs) -> dict[str, torch.Tensor]:
-        """Save trainable weights (projector + audio_head if present)."""
-        state = {f"projector.{k}": v for k, v in self.projector.state_dict().items()}
-        if self.audio_head is not None:
-            state.update({f"audio_head.{k}": v for k, v in self.audio_head.state_dict().items()})
-        return state
+        """Only save trainable projector weights."""
+        return {f"projector.{k}": v for k, v in self.projector.state_dict().items()}
 
     def _compute_encoder_output_lengths(
         self,
@@ -476,8 +460,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         labels: Optional[torch.Tensor] = None,
         use_cache: Optional[bool] = None,
         cache_position: Optional[torch.Tensor] = None,
-        codec_targets: Optional[torch.Tensor] = None,
-        codec_lengths: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> CausalLMOutputWithPast:
         """Forward pass for training and inference."""
@@ -505,10 +487,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
                 audio_embeds.to(inputs_embeds.device, dtype=inputs_embeds.dtype),
             )
 
-        # Request hidden states if training audio head with codec targets
-        if self.audio_head is not None and codec_targets is not None:
-            kwargs["output_hidden_states"] = True
-
         # Run through language model (let it compute loss if labels provided)
         outputs = self.language_model(
             attention_mask=attention_mask,
@@ -527,29 +505,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             if aux_loss is not None and aux_loss.numel() > 0:
                 outputs.loss = outputs.loss + aux_loss.to(outputs.loss.device)
 
-        # Compute audio head loss if training S2S with codec targets
-        if self.audio_head is not None and codec_targets is not None:
-            hidden_states = outputs.hidden_states[-1]  # Last layer hidden states
-            # No detach needed: LLM is frozen (requires_grad=False), so gradients
-            # naturally stop there. Hidden states keep their grad_fn for proper backprop.
-            audio_head_loss = self.audio_head(
-                hidden_states,
-                codec_targets=codec_targets,
-                codec_lengths=codec_lengths,
-            )
-            # Add audio_head_loss directly to outputs.loss
-            # (CausalLMOutputWithPast doesn't preserve custom attributes through Accelerator)
-            if outputs.loss is not None:
-                outputs.loss = outputs.loss + audio_head_loss
-            else:
-                # S2S-only training: audio head loss is the only loss
-                outputs.loss = audio_head_loss
-        else:
-            print(
-                f"DEBUG: audio_head branch NOT taken: audio_head={self.audio_head is not None}, codec_targets={codec_targets is not None}"
-            )
-
-        print(f"DEBUG: returning outputs.loss={outputs.loss}")
         return outputs
 
     def prepare_inputs_for_generation(self, *args, **kwargs):
@@ -833,139 +788,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         response = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
         return response.strip()
 
-    @torch.no_grad()
-    def generate_with_audio(
-        self,
-        input_features: torch.Tensor,
-        audio_attention_mask: torch.Tensor,
-        **generate_kwargs,
-    ) -> dict[str, torch.Tensor]:
-        """Generate text and NeuCodec tokens for Speech-to-Speech.
-
-        Args:
-            input_features: Mel spectrogram features (batch, n_mels, mel_len)
-            audio_attention_mask: Mask for real vs padded mel frames (batch, mel_len)
-            **generate_kwargs: Additional generation arguments
-
-        Returns:
-            Dict with:
-                - text_ids: Generated text token IDs (batch, seq_len)
-                - text: Decoded text strings (list of str)
-                - codec_tokens: Predicted NeuCodec tokens (batch, audio_len)
-        """
-        if self.audio_head is None:
-            raise ValueError("Audio head not configured. Set use_audio_head=True in config.")
-
-        device = input_features.device
-        batch_size = input_features.shape[0]
-
-        # Encode audio -> flattened embeddings
-        audio_embeds = self._encode_audio(input_features, audio_attention_mask)
-
-        # Build prompt with correct number of audio tokens
-        num_audio_tokens = self._get_num_audio_tokens(audio_attention_mask)
-        audio_placeholder = "<audio>" * num_audio_tokens
-
-        messages: list[dict[str, str]] = []
-        if self.system_prompt:
-            messages.append({"role": "system", "content": self.system_prompt})
-        user_content = audio_placeholder
-        if self.TRANSCRIBE_PROMPT:
-            user_content += " " + self.TRANSCRIBE_PROMPT
-        messages.append({"role": "user", "content": user_content})
-
-        chat_result = self.tokenizer.apply_chat_template(
-            messages,
-            tokenize=True,
-            add_generation_prompt=True,
-            return_tensors="pt",
-            enable_thinking=getattr(self.config, "enable_thinking", False),
-        )
-        input_ids = chat_result.input_ids.to(device)
-
-        if input_ids.dim() == 1:
-            input_ids = input_ids.unsqueeze(0)
-        if input_ids.shape[0] == 1 and batch_size > 1:
-            input_ids = input_ids.expand(batch_size, -1)
-
-        attention_mask = torch.ones_like(input_ids)
-
-        # Get text embeddings and replace audio tokens with audio embeddings
-        inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
-        audio_token_mask = (input_ids == self.audio_token_id).unsqueeze(-1)
-        inputs_embeds = inputs_embeds.masked_scatter(
-            audio_token_mask.to(inputs_embeds.device),
-            audio_embeds.to(inputs_embeds.device, dtype=inputs_embeds.dtype),
-        )
-
-        # Generate with hidden states
-        output = self.language_model.generate(
-            input_ids=input_ids,
-            inputs_embeds=inputs_embeds,
-            attention_mask=attention_mask,
-            generation_config=self.generation_config,
-            output_hidden_states=True,
-            return_dict_in_generate=True,
-            **generate_kwargs,
-        )
-
-        # Extract generated text
-        text_ids = output.sequences[:, input_ids.shape[1] :]
-        text = self.tokenizer.batch_decode(text_ids, skip_special_tokens=True)
-
-        # Extract hidden states from generation steps and concatenate
-        # output.hidden_states is tuple of (step,) where each step is tuple of (layer,)
-        # Each layer tensor is (batch, 1, hidden_dim) for generated tokens
-        last_layer_states = []
-        for step_hidden in output.hidden_states:
-            # step_hidden is tuple of (num_layers,) tensors
-            # Get last layer: shape (batch, 1, hidden_dim)
-            last_layer_states.append(step_hidden[-1])
-
-        # Concatenate across generation steps: (batch, gen_seq_len, hidden_dim)
-        hidden_states = torch.cat(last_layer_states, dim=1)
-
-        # Predict codec tokens (uses inference heuristic for duration)
-        # WavTokenizer: single codebook, shape (batch, audio_len)
-        codec_tokens = self.audio_head(hidden_states)
-
-        return {
-            "text_ids": text_ids,
-            "text": text,
-            "codec_tokens": codec_tokens,
-        }
-
-    def decode_audio(
-        self,
-        codec_tokens: torch.Tensor,
-        codec_model_id: str = "neuphonic/neucodec",
-    ) -> torch.Tensor:
-        """Decode NeuCodec tokens to waveform.
-
-        Args:
-            codec_tokens: Codec token indices (batch, audio_len)
-            codec_model_id: HuggingFace model ID for NeuCodec
-
-        Returns:
-            Waveform tensor (batch, 1, samples) at 24kHz
-        """
-        try:
-            from neucodec import NeuCodec
-        except ImportError as e:
-            raise ImportError(
-                "NeuCodec required for audio decoding. Install with: pip install neucodec"
-            ) from e
-
-        model = NeuCodec.from_pretrained(codec_model_id)
-        model = model.to(codec_tokens.device)
-        model.eval()
-
-        # NeuCodec decode expects (batch, 1, seq_len)
-        codes = codec_tokens.unsqueeze(1)
-
-        with torch.no_grad():
-            return model.decode_code(codes)
-
     def save_pretrained(self, save_directory: Union[str, Path], **kwargs) -> None:
         """Save model, tokenizer, and processor."""
         import shutil

asr_pipeline.py

@@ -2,7 +2,7 @@
 
 import re
 from pathlib import Path
-from typing import Any, Iterator, Union
+from typing import Any
 
 import numpy as np
 import torch
@@ -101,142 +101,6 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         audio = np.concatenate(audio_chunks) if audio_chunks else np.array([], dtype=np.float32)
         return {"audio": audio, "sample_rate": TTS_SAMPLE_RATE}
 
-    def transcribe_streaming(
-        self,
-        inputs: Union[str, bytes, np.ndarray, dict],
-        system_prompt: str | None = None,
-    ) -> Iterator[str]:
-        """Transcribe audio with streaming token output for low-latency applications.
-
-        Yields partial transcript strings as tokens are generated, reducing
-        time-to-first-word compared to waiting for full transcription.
-
-        Args:
-            inputs: Audio input in any supported format:
-                - str: File path to audio file
-                - bytes: Raw audio bytes
-                - np.ndarray: Audio samples as numpy array
-                - dict: {"array": np.ndarray, "sampling_rate": int}
-            system_prompt: Optional system prompt override (uses model's default if not provided)
-
-        Yields:
-            Partial transcript text as each token is generated
-
-        Example:
-            >>> for partial in pipeline.transcribe_streaming("audio.wav"):
-            ...     print(partial, end="", flush=True)
-        """
-        # Extract audio array from various input formats
-        audio_data = self._extract_audio(inputs)
-        if audio_data is None:
-            return
-
-        audio_array = audio_data["array"]
-        sample_rate = audio_data.get("sampling_rate", 16000)
-
-        # Preprocess audio through feature extractor
-        model_inputs = self.feature_extractor(
-            audio_array,
-            sampling_rate=sample_rate,
-            return_tensors="pt",
-            return_attention_mask=True,
-        )
-
-        # Get model dtype and device, cast inputs to match
-        device = self.model.device
-        model_dtype = next(self.model.parameters()).dtype
-        input_features = model_inputs.input_features.to(device, dtype=model_dtype)
-        attention_mask = model_inputs.attention_mask.to(device)
-
-        # Stream tokens from model
-        yield from self.model.generate_streaming(
-            input_features=input_features,
-            audio_attention_mask=attention_mask,
-            system_prompt=system_prompt,
-        )
-
-    def transcribe_streaming_with_audio(
-        self,
-        inputs: Union[str, bytes, np.ndarray, dict],
-        voice: str = DEFAULT_TTS_VOICE,
-        system_prompt: str | None = None,
-    ) -> Iterator[dict[str, Any]]:
-        """Transcribe audio with streaming text AND audio output.
-
-        Yields partial text as tokens are generated, and audio chunks
-        as complete sentences are detected. This enables low-latency
-        voice agents that can start speaking before transcription completes.
-
-        Args:
-            inputs: Audio input (same formats as transcribe_streaming)
-            voice: Kokoro TTS voice ID
-            system_prompt: Optional system prompt override (uses model's default if not provided)
-
-        Yields:
-            Dicts with either:
-            - {"type": "text", "text": str, "interim": bool} for text updates
-            - {"type": "audio", "audio": np.ndarray, "sample_rate": int} for audio chunks
-
-        Example:
-            >>> for chunk in pipeline.transcribe_streaming_with_audio(audio):
-            ...     if chunk["type"] == "text":
-            ...         print(chunk["text"], end="", flush=True)
-            ...     elif chunk["type"] == "audio":
-            ...         play_audio(chunk["audio"], chunk["sample_rate"])
-        """
-        import re
-
-        sentence_buffer = ""
-        full_text = ""
-
-        # Sentence-ending patterns (handles ., !, ?, and common abbreviations)
-        sentence_end_pattern = re.compile(r"[.!?](?:\s|$)")
-
-        for token_text in self.transcribe_streaming(inputs, system_prompt=system_prompt):
-            full_text += token_text
-            sentence_buffer += token_text
-
-            # Yield text update
-            yield {"type": "text", "text": full_text, "interim": True}
-
-            # Check for complete sentence
-            match = sentence_end_pattern.search(sentence_buffer)
-            if match:
-                # Extract complete sentence(s)
-                end_pos = match.end()
-                complete_text = sentence_buffer[:end_pos].strip()
-                sentence_buffer = sentence_buffer[end_pos:]
-
-                # Generate audio for the complete sentence
-                if complete_text:
-                    try:
-                        tts_result = self.text_to_speech(complete_text, voice=voice)
-                        if tts_result["audio"] is not None and len(tts_result["audio"]) > 0:
-                            yield {
-                                "type": "audio",
-                                "audio": tts_result["audio"],
-                                "sample_rate": tts_result["sample_rate"],
-                            }
-                    except Exception:
-                        pass  # Skip audio on TTS errors
-
-        # Final text update (not interim)
-        yield {"type": "text", "text": full_text, "interim": False}
-
-        # Generate audio for any remaining text
-        remaining = sentence_buffer.strip()
-        if remaining:
-            try:
-                tts_result = self.text_to_speech(remaining, voice=voice)
-                if tts_result["audio"] is not None and len(tts_result["audio"]) > 0:
-                    yield {
-                        "type": "audio",
-                        "audio": tts_result["audio"],
-                        "sample_rate": tts_result["sample_rate"],
-                    }
-            except Exception:
-                pass
-
     def _sanitize_parameters(self, **kwargs):
         """Intercept our custom parameters before parent class validates them."""
         # Remove our custom parameters so parent doesn't see them
@@ -247,7 +111,6 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         kwargs.pop("max_speakers", None)
         kwargs.pop("hf_token", None)
         kwargs.pop("user_prompt", None)
-        kwargs.pop("system_prompt", None)
         kwargs.pop("diarization_backend", None)
         # TTS parameters
         kwargs.pop("return_audio", None)
@@ -269,7 +132,6 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             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: ")
-            system_prompt: Custom system prompt override (uses model's default if not provided)
             num_speakers: Exact number of speakers (if known, for diarization)
             min_speakers: Minimum number of speakers (for diarization)
             max_speakers: Maximum number of speakers (for diarization)
@@ -286,7 +148,6 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         return_audio = kwargs.pop("return_audio", False)
         tts_voice = kwargs.pop("tts_voice", DEFAULT_TTS_VOICE)
         user_prompt = kwargs.pop("user_prompt", None)
-        system_prompt = kwargs.pop("system_prompt", None)
         diarization_params = {
             "num_speakers": kwargs.pop("num_speakers", None),
             "min_speakers": kwargs.pop("min_speakers", None),
@@ -302,12 +163,6 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             original_prompt = self.model.TRANSCRIBE_PROMPT
             self.model.TRANSCRIBE_PROMPT = user_prompt
 
-        # Set custom system prompt if provided
-        original_system_prompt = None
-        if system_prompt:
-            original_system_prompt = self.model.system_prompt
-            self.model.system_prompt = system_prompt
-
         # Store audio for timestamp alignment and diarization
         if return_timestamps or return_speakers:
             self._current_audio = self._extract_audio(inputs)
@@ -369,25 +224,11 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         self._current_audio = None
         if original_prompt is not None:
             self.model.TRANSCRIBE_PROMPT = original_prompt
-        if original_system_prompt is not None:
-            self.model.system_prompt = original_system_prompt
 
         return result
 
     def _extract_audio(self, inputs) -> dict | None:
-        """Extract audio array from various input formats.
-
-        Supported input formats:
-            - str: File path to audio file
-            - bytes: Encoded audio (mp3, wav, etc.) - decoded via ffmpeg
-            - np.ndarray: Audio samples as float32 array
-            - dict with "array": Audio samples as numpy array
-            - dict with "raw": Alias for "array" (HF pipeline compat)
-            - dict with "raw_bytes": Raw PCM bytes (requires "dtype", optional "sampling_rate")
-
-        For raw PCM bytes (e.g., from pipecat), use:
-            {"raw_bytes": pcm_bytes, "dtype": "int16", "sampling_rate": 16000}
-        """
+        """Extract audio array from various input formats using HF utilities."""
         from transformers.pipelines.audio_utils import ffmpeg_read
 
         if isinstance(inputs, dict):
@@ -401,17 +242,6 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
                     "array": inputs["raw"],
                     "sampling_rate": inputs.get("sampling_rate", 16000),
                 }
-            if "raw_bytes" in inputs:
-                # Raw PCM bytes - convert to float32 array
-                dtype = inputs.get("dtype", "int16")
-                sample_rate = inputs.get("sampling_rate", 16000)
-                audio = np.frombuffer(inputs["raw_bytes"], dtype=dtype).astype(np.float32)
-                # Normalize based on dtype
-                if dtype == "int16":
-                    audio = audio / 32768.0
-                elif dtype == "int32":
-                    audio = audio / 2147483648.0
-                return {"array": audio, "sampling_rate": sample_rate}
         elif isinstance(inputs, str):
             # File path - load audio using ffmpeg (same as HF pipeline)
             with Path(inputs).open("rb") as f: