Assembled S2S model (base + AudioHead)

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

alignment.py

@@ -0,0 +1,295 @@
+"""Forced alignment for word-level timestamps using Wav2Vec2."""
+
+import numpy as np
+import torch
+
+# Offset compensation for Wav2Vec2-BASE systematic bias (in seconds)
+# Calibrated on librispeech-alignments dataset (n=25, MAE=48ms)
+START_OFFSET = 0.04  # Subtract from start times (shift earlier)
+END_OFFSET = -0.04  # Subtract from end times (shift later)
+
+
+def _get_device() -> str:
+    """Get best available device for non-transformers models."""
+    if torch.cuda.is_available():
+        return "cuda"
+    if torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
+class ForcedAligner:
+    """Lazy-loaded forced aligner for word-level timestamps using torchaudio wav2vec2.
+
+    Uses Viterbi trellis algorithm for optimal alignment path finding.
+    """
+
+    _bundle = None
+    _model = None
+    _labels = None
+    _dictionary = None
+
+    @classmethod
+    def get_instance(cls, device: str = "cuda"):
+        """Get or create the forced alignment model (singleton).
+
+        Args:
+            device: Device to run model on ("cuda" or "cpu")
+
+        Returns:
+            Tuple of (model, labels, dictionary)
+        """
+        if cls._model is None:
+            import torchaudio
+
+            cls._bundle = torchaudio.pipelines.WAV2VEC2_ASR_BASE_960H
+            cls._model = cls._bundle.get_model().to(device)
+            cls._model.eval()
+            cls._labels = cls._bundle.get_labels()
+            cls._dictionary = {c: i for i, c in enumerate(cls._labels)}
+        return cls._model, cls._labels, cls._dictionary
+
+    @staticmethod
+    def _get_trellis(emission: torch.Tensor, tokens: list[int], blank_id: int = 0) -> torch.Tensor:
+        """Build trellis for forced alignment using forward algorithm.
+
+        The trellis[t, j] represents the log probability of the best path that
+        aligns the first j tokens to the first t frames.
+
+        Args:
+            emission: Log-softmax emission matrix of shape (num_frames, num_classes)
+            tokens: List of target token indices
+            blank_id: Index of the blank/CTC token (default 0)
+
+        Returns:
+            Trellis matrix of shape (num_frames + 1, num_tokens + 1)
+        """
+        num_frames = emission.size(0)
+        num_tokens = len(tokens)
+
+        trellis = torch.full((num_frames + 1, num_tokens + 1), -float("inf"))
+        trellis[0, 0] = 0
+
+        # Force alignment to use all tokens by preventing staying in blank
+        # at the end when there are still tokens to emit
+        if num_tokens > 1:
+            trellis[-num_tokens + 1 :, 0] = float("inf")
+
+        for t in range(num_frames):
+            for j in range(num_tokens + 1):
+                # Stay: emit blank and stay at j tokens
+                stay = trellis[t, j] + emission[t, blank_id]
+
+                # Move: emit token j and advance to j+1 tokens
+                move = trellis[t, j - 1] + emission[t, tokens[j - 1]] if j > 0 else -float("inf")
+
+                trellis[t + 1, j] = max(stay, move)  # Viterbi: take best path
+
+        return trellis
+
+    @staticmethod
+    def _backtrack(
+        trellis: torch.Tensor, emission: torch.Tensor, tokens: list[int], blank_id: int = 0
+    ) -> list[tuple[int, float, float, float]]:
+        """Backtrack through trellis to find optimal forced monotonic alignment.
+
+        Guarantees:
+        - All tokens are emitted exactly once
+        - Strictly monotonic: each token's frames come after previous token's
+        - No frame skipping or token teleporting
+
+        Returns list of (token_id, start_frame, end_frame, peak_frame) for each token.
+        The peak_frame is the frame with highest emission probability for that token.
+        """
+        num_frames = emission.size(0)
+        num_tokens = len(tokens)
+
+        if num_tokens == 0:
+            return []
+
+        # Find the best ending point (should be at num_tokens)
+        # But verify trellis reached a valid state
+        if trellis[num_frames, num_tokens] == -float("inf"):
+            # Alignment failed - fall back to uniform distribution
+            frames_per_token = num_frames / num_tokens
+            return [
+                (
+                    tokens[i],
+                    i * frames_per_token,
+                    (i + 1) * frames_per_token,
+                    (i + 0.5) * frames_per_token,
+                )
+                for i in range(num_tokens)
+            ]
+
+        # Backtrack: find where each token transition occurred
+        # Store (frame, emission_score) for each token
+        token_frames: list[list[tuple[int, float]]] = [[] for _ in range(num_tokens)]
+
+        t = num_frames
+        j = num_tokens
+
+        while t > 0 and j > 0:
+            # Check: did we transition from j-1 to j at frame t-1?
+            stay_score = trellis[t - 1, j] + emission[t - 1, blank_id]
+            move_score = trellis[t - 1, j - 1] + emission[t - 1, tokens[j - 1]]
+
+            if move_score >= stay_score:
+                # Token j-1 was emitted at frame t-1
+                # Store frame and its emission probability
+                emit_prob = emission[t - 1, tokens[j - 1]].exp().item()
+                token_frames[j - 1].insert(0, (t - 1, emit_prob))
+                j -= 1
+            # Always decrement time (monotonic)
+            t -= 1
+
+        # Handle any remaining tokens at the start (edge case)
+        while j > 0:
+            token_frames[j - 1].insert(0, (0, 0.0))
+            j -= 1
+
+        # Convert to spans with peak frame
+        token_spans: list[tuple[int, float, float, float]] = []
+        for token_idx, frames_with_scores in enumerate(token_frames):
+            if not frames_with_scores:
+                # Token never emitted - assign minimal span after previous
+                if token_spans:
+                    prev_end = token_spans[-1][2]
+                    frames_with_scores = [(int(prev_end), 0.0)]
+                else:
+                    frames_with_scores = [(0, 0.0)]
+
+            token_id = tokens[token_idx]
+            frames = [f for f, _ in frames_with_scores]
+            start_frame = float(min(frames))
+            end_frame = float(max(frames)) + 1.0
+
+            # Find peak frame (highest emission probability)
+            peak_frame, _ = max(frames_with_scores, key=lambda x: x[1])
+
+            token_spans.append((token_id, start_frame, end_frame, float(peak_frame)))
+
+        return token_spans
+
+    @classmethod
+    def align(
+        cls,
+        audio: np.ndarray,
+        text: str,
+        sample_rate: int = 16000,
+    ) -> list[dict]:
+        """Align transcript to audio and return word-level timestamps.
+
+        Uses Viterbi trellis algorithm for optimal forced alignment.
+
+        Args:
+            audio: Audio waveform as numpy array
+            text: Transcript text to align
+            sample_rate: Audio sample rate (default 16000)
+
+        Returns:
+            List of dicts with 'word', 'start', 'end' keys
+        """
+        import torchaudio
+
+        device = _get_device()
+        model, _labels, dictionary = cls.get_instance(device)
+        assert cls._bundle is not None and dictionary is not None  # Initialized by get_instance
+
+        # Convert audio to tensor (copy to ensure array is writable)
+        if isinstance(audio, np.ndarray):
+            waveform = torch.from_numpy(audio.copy()).float()
+        else:
+            waveform = audio.clone().float()
+
+        # Ensure 2D (channels, time)
+        if waveform.dim() == 1:
+            waveform = waveform.unsqueeze(0)
+
+        # Resample if needed (wav2vec2 expects 16kHz)
+        if sample_rate != cls._bundle.sample_rate:
+            waveform = torchaudio.functional.resample(
+                waveform, sample_rate, cls._bundle.sample_rate
+            )
+
+        waveform = waveform.to(device)
+
+        # Get emissions from model
+        with torch.inference_mode():
+            emissions, _ = model(waveform)
+            emissions = torch.log_softmax(emissions, dim=-1)
+
+        emission = emissions[0].cpu()
+
+        # Normalize text: uppercase, keep only valid characters
+        transcript = text.upper()
+
+        # Build tokens from transcript (including word separators)
+        tokens = []
+        for char in transcript:
+            if char in dictionary:
+                tokens.append(dictionary[char])
+            elif char == " ":
+                tokens.append(dictionary.get("|", dictionary.get(" ", 0)))
+
+        if not tokens:
+            return []
+
+        # Build Viterbi trellis and backtrack for optimal path
+        trellis = cls._get_trellis(emission, tokens, blank_id=0)
+        alignment_path = cls._backtrack(trellis, emission, tokens, blank_id=0)
+
+        # Convert frame indices to time (model stride is 320 samples at 16kHz = 20ms)
+        frame_duration = 320 / cls._bundle.sample_rate
+
+        # Apply separate offset compensation for start/end (Wav2Vec2 systematic bias)
+        start_offset = START_OFFSET
+        end_offset = END_OFFSET
+
+        # Group aligned tokens into words based on pipe separator
+        # Use peak emission frame for more accurate word boundaries
+        words = text.split()
+        word_timestamps = []
+        first_char_peak = None
+        last_char_peak = None
+        word_idx = 0
+        separator_id = dictionary.get("|", dictionary.get(" ", 0))
+
+        for token_id, _start_frame, _end_frame, peak_frame in alignment_path:
+            if token_id == separator_id:  # Word separator
+                if (
+                    first_char_peak is not None
+                    and last_char_peak is not None
+                    and word_idx < len(words)
+                ):
+                    # Use peak frames for word boundaries
+                    start_time = max(0.0, first_char_peak * frame_duration - start_offset)
+                    end_time = max(0.0, (last_char_peak + 1) * frame_duration - end_offset)
+                    word_timestamps.append(
+                        {
+                            "word": words[word_idx],
+                            "start": start_time,
+                            "end": end_time,
+                        }
+                    )
+                    word_idx += 1
+                first_char_peak = None
+                last_char_peak = None
+            else:
+                if first_char_peak is None:
+                    first_char_peak = peak_frame
+                last_char_peak = peak_frame
+
+        # Don't forget the last word
+        if first_char_peak is not None and last_char_peak is not None and word_idx < len(words):
+            start_time = max(0.0, first_char_peak * frame_duration - start_offset)
+            end_time = max(0.0, (last_char_peak + 1) * frame_duration - end_offset)
+            word_timestamps.append(
+                {
+                    "word": words[word_idx],
+                    "start": start_time,
+                    "end": end_time,
+                }
+            )
+
+        return word_timestamps

asr_config.py

@@ -0,0 +1,166 @@
+from typing import Optional
+
+import transformers
+
+
+class ASRConfig(transformers.PretrainedConfig):
+    """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",
+        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: Optional[list] = None,
+        audio_sample_rate: int = 16000,
+        # Projector settings
+        projector_type: str = "mlp",
+        projector_pool_stride: int = 4,
+        projector_hidden_dim: Optional[int] = None,
+        # 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,
+        freeze_projector: bool = False,
+        label_smoothing: float = 0.0,
+        # Audio Head settings (trainable AR decoder + NeuCodec)
+        use_audio_head: bool = False,
+        freeze_audio_head: bool = False,
+        max_audio_tokens: int = 500,
+        decoder_dim: int = 512,
+        decoder_layers: int = 6,
+        decoder_heads: int = 8,
+        neucodec_model_id: str = "neuphonic/neucodec",
+        **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
+        self.encoder_conv_layers = encoder_conv_layers or [(1, 3, 1), (1, 3, 2)]
+        self.audio_sample_rate = audio_sample_rate
+
+        # Projector settings
+        self.projector_type = projector_type
+        self.projector_pool_stride = projector_pool_stride
+        self.projector_hidden_dim = projector_hidden_dim
+
+        # 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
+        self.freeze_projector = freeze_projector
+        self.label_smoothing = label_smoothing
+
+        # Audio Head settings (trainable AR decoder + NeuCodec)
+        self.use_audio_head = use_audio_head
+        self.freeze_audio_head = freeze_audio_head
+        self.max_audio_tokens = max_audio_tokens
+        self.decoder_dim = decoder_dim
+        self.decoder_layers = decoder_layers
+        self.decoder_heads = decoder_heads
+        self.neucodec_model_id = neucodec_model_id
+
+        # 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, trust_remote_code=True
+            )
+            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(
+                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
+            )
+            self.text_config.dtype = model_dtype
+        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)
+
+        super().__init__(**kwargs)
+
+        # Pipeline configuration
+        self.encoder = self.audio_config
+        self.auto_map = {
+            "AutoConfig": "asr_config.ASRConfig",
+            "AutoModel": "asr_modeling.ASRModel",
+            "AutoModelForSpeechSeq2Seq": "asr_modeling.ASRModel",
+            "AutoProcessor": "asr_processing.ASRProcessor",
+        }
+        self.custom_pipelines = {
+            "automatic-speech-recognition": {
+                "impl": "asr_pipeline.ASRPipeline",
+                "pt": ["AutoModelForSpeechSeq2Seq"],
+                "tf": [],
+                "type": "audio",
+            }
+        }
+        self.architectures = ["ASRModel"]
+        self.pipeline_tag = "automatic-speech-recognition"
+
+
+transformers.AutoConfig.register("asr_model", ASRConfig)

asr_modeling.py

@@ -0,0 +1,817 @@
+import json
+from pathlib import Path
+from typing import Optional, Union
+
+import torch
+import torch.nn as nn
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    PreTrainedModel,
+)
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+try:
+    from .asr_config import ASRConfig
+    from .projectors import PROJECTOR_CLASSES
+except ImportError:
+    from asr_config import ASRConfig  # type: ignore[no-redef]
+    from projectors import PROJECTOR_CLASSES  # type: ignore[no-redef]
+
+
+from torchaudio.transforms import SpecAugment
+
+
+class ASRModel(PreTrainedModel, GenerationMixin):
+    """Audio-to-text model combining an audio encoder, projector, and language model."""
+
+    config_class = ASRConfig
+    base_model_prefix = "model"
+    main_input_name = "input_features"
+    _supports_flash_attn_2 = True
+    supports_gradient_checkpointing = True
+    _is_loading_from_pretrained: bool = False
+    _pretrained_model_path: Optional[str] = None
+
+    TRANSCRIBE_PROMPT = ""
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, *args, **kwargs) -> "ASRModel":
+        """Load model from pretrained, handling device placement correctly."""
+        from safetensors.torch import load_file
+        from transformers.utils.hub import cached_file
+
+        config = kwargs.pop("config", None)
+        if config is None:
+            config = ASRConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+
+        # Set flag to avoid device_map="auto" in sub-model loaders
+        cls._is_loading_from_pretrained = True
+        cls._pretrained_model_path = pretrained_model_name_or_path
+
+        try:
+            model = cls(config, **kwargs)
+
+            # Load projector weights from safetensors
+            subfolder = kwargs.get("subfolder")
+            revision = kwargs.get("revision")
+            cache_kwargs = {}
+            if subfolder:
+                cache_kwargs["subfolder"] = subfolder
+            if revision:
+                cache_kwargs["revision"] = revision
+
+            model_file = cached_file(
+                pretrained_model_name_or_path,
+                "model.safetensors",
+                _raise_exceptions_for_missing_entries=False,
+                **cache_kwargs,
+            )
+
+            if model_file is not None:
+                state_dict = load_file(model_file)
+                model.load_state_dict(state_dict, strict=False)
+
+            return model
+        finally:
+            cls._is_loading_from_pretrained = False
+            cls._pretrained_model_path = None
+
+    def __init__(self, config: ASRConfig, **kwargs) -> None:
+        super().__init__(config)
+
+        self.system_prompt = config.system_prompt
+        target_dtype = getattr(torch, config.model_dtype)
+
+        # Audio encoder (frozen)
+        self.audio_tower = self._load_audio_encoder(config, target_dtype)
+
+        # Language model (frozen)
+        self.language_model = self._load_language_model(config, target_dtype)
+
+        # Initialize tokenizer and special tokens
+        self._init_tokenizer(config)
+
+        # Set up generation config with greedy decoding defaults
+        self.generation_config = self.language_model.generation_config
+        self.generation_config.max_new_tokens = config.max_new_tokens
+        self.generation_config.min_new_tokens = config.min_new_tokens
+        self.generation_config.num_beams = config.num_beams
+        self.generation_config.do_sample = config.do_sample
+        # Set sampling params from config (None means use model defaults)
+        self.generation_config.temperature = config.temperature
+        self.generation_config.top_p = config.top_p
+        self.generation_config.top_k = config.top_k
+        self.generation_config.use_cache = config.use_cache
+        self.generation_config.length_penalty = config.length_penalty
+        self.generation_config.repetition_penalty = config.repetition_penalty
+        self.generation_config.no_repeat_ngram_size = config.no_repeat_ngram_size
+        # Set EOS tokens, filtering out any that don't exist in the tokenizer
+        eos_candidates = [
+            self.tokenizer.convert_tokens_to_ids("<|im_end|>"),
+            self.tokenizer.convert_tokens_to_ids("<|endoftext|>"),
+        ]
+        self.generation_config.eos_token_id = [t for t in eos_candidates if t is not None]
+        self.generation_config.pad_token_id = self.tokenizer.pad_token_id
+
+        # Feature extractor for audio preprocessing
+        self.feature_extractor = self._create_feature_extractor(config)
+
+        # Audio projector (trainable unless freeze_projector is set)
+        self.projector = self._create_projector(config, target_dtype)
+
+        # Learned padding embedding for audio tokens (used when projector output is short)
+        # Using a learned embedding instead of zeros keeps values in the embedding distribution
+        self.audio_pad_embedding = nn.Parameter(torch.randn(1, config.llm_dim) * 0.02)
+
+        # Freeze projector if specified
+        if getattr(config, "freeze_projector", False):
+            self.projector.requires_grad_(False)
+
+        # SpecAugment for data augmentation during training
+        if getattr(config, "use_specaugment", False):
+            self.spec_augment = SpecAugment(
+                n_time_masks=config.num_time_masks,
+                time_mask_param=config.time_mask_length,
+                n_freq_masks=config.num_freq_masks,
+                freq_mask_param=config.freq_mask_length,
+            )
+        else:
+            self.spec_augment = None
+
+        # Audio head for S2S (trainable AR decoder + NeuCodec)
+        if getattr(config, "use_audio_head", False):
+            from .audio_head import AudioHead, AudioHeadConfig
+
+            device = next(self.language_model.parameters()).device
+
+            audio_head_config = AudioHeadConfig(
+                decoder_dim=config.decoder_dim,
+                decoder_layers=config.decoder_layers,
+                decoder_heads=config.decoder_heads,
+                text_vocab_size=len(self.tokenizer),
+                max_audio_tokens=config.max_audio_tokens,
+                neucodec_model_id=getattr(config, "neucodec_model_id", "neuphonic/neucodec"),
+                temperature=getattr(config, "audio_head_temperature", 1.0),
+                top_k=getattr(config, "audio_head_top_k", 50),
+            )
+            self.audio_head = AudioHead(audio_head_config).to(
+                device=device, dtype=target_dtype
+            )
+
+            if getattr(config, "freeze_audio_head", False):
+                self.audio_head.requires_grad_(False)
+        else:
+            self.audio_head = None
+
+        # Silero VAD for interruption detection (Freeze-Omni style)
+        # Loaded lazily on first use to avoid startup cost
+        self._vad_model = None
+        self._vad_utils = None
+
+        # For model parallelism
+        self._no_split_modules = getattr(self.language_model, "_no_split_modules", [])
+
+    def _tie_weights(self):
+        """No-op: AudioHead manages its own embeddings."""
+        pass
+
+    def _create_feature_extractor(self, config: ASRConfig):
+        """Create the appropriate feature extractor for the audio encoder."""
+        from transformers import AutoFeatureExtractor
+
+        feature_extractor = AutoFeatureExtractor.from_pretrained(config.audio_model_id)
+        # Disable padding by default - use actual audio length
+        feature_extractor.padding = False
+        return feature_extractor
+
+    @classmethod
+    def _load_audio_encoder(cls, config: ASRConfig, dtype: torch.dtype) -> nn.Module:
+        """Load and freeze the audio encoder."""
+        encoder_kwargs = {
+            "attn_implementation": config.attn_implementation,
+            "low_cpu_mem_usage": True,
+            "torch_dtype": dtype,
+        }
+
+        if "whisper" in config.audio_model_id.lower():
+            from transformers import WhisperModel
+
+            full_model = WhisperModel.from_pretrained(config.audio_model_id, **encoder_kwargs)
+            encoder = full_model.encoder
+            del full_model
+        elif "glm" in config.audio_model_id.lower():
+            # GLM-ASR models use audio_tower as the encoder
+            # Requires transformers >= 5.x or installed from source
+            from transformers import AutoModelForSeq2SeqLM
+
+            full_model = AutoModelForSeq2SeqLM.from_pretrained(
+                config.audio_model_id, trust_remote_code=True, **encoder_kwargs
+            )
+            # GLM stores encoder at audio_tower (GlmAsrEncoder)
+            encoder = full_model.audio_tower
+            # Clear references to free VRAM from the LLM decoder
+            full_model.language_model = None
+            full_model.multi_modal_projector = None
+            del full_model
+        else:
+            encoder = AutoModel.from_pretrained(config.audio_model_id, **encoder_kwargs)
+
+        encoder.requires_grad_(False)
+        encoder.eval()
+        return encoder
+
+    @classmethod
+    def _load_language_model(cls, config: ASRConfig, dtype: torch.dtype) -> PreTrainedModel:
+        """Load and freeze the language model."""
+        decoder_kwargs = {
+            "attn_implementation": config.attn_implementation,
+            "trust_remote_code": True,
+            "low_cpu_mem_usage": True,
+            "dtype": dtype,
+        }
+
+        decoder = AutoModelForCausalLM.from_pretrained(config.text_model_id, **decoder_kwargs)
+        decoder.config.use_cache = getattr(config, "use_cache", True)
+        decoder.requires_grad_(False)
+        decoder.eval()
+        return decoder
+
+    def _create_projector(self, config: ASRConfig, dtype: torch.dtype) -> nn.Module:
+        """Create the trainable audio projector."""
+        # Auto-detect dimensions if not specified
+        if config.encoder_dim is None:
+            enc_cfg = self.audio_tower.config
+            config.encoder_dim = getattr(enc_cfg, "hidden_size", None) or getattr(
+                enc_cfg, "d_model", None
+            )
+            if config.encoder_dim is None:
+                raise ValueError("Could not auto-detect encoder_dim. Please specify in config.")
+
+        if config.llm_dim is None:
+            dec_cfg = self.language_model.config
+            config.llm_dim = getattr(dec_cfg, "hidden_size", None) or getattr(
+                dec_cfg, "d_model", None
+            )
+            if config.llm_dim is None:
+                raise ValueError("Could not auto-detect llm_dim. Please specify in config.")
+
+        # Select projector type based on config
+        projector_type = getattr(config, "projector_type", "mlp")
+        projector_class = PROJECTOR_CLASSES.get(projector_type)
+        if projector_class is None:
+            raise ValueError(
+                f"Unknown projector_type: {projector_type}. "
+                f"Valid options: {list(PROJECTOR_CLASSES.keys())}"
+            )
+        projector = projector_class(config)
+
+        # Move projector to same device as language model (important when using quantization)
+        device = next(self.language_model.parameters()).device
+        return projector.to(device=device, dtype=dtype)
+
+    def _init_tokenizer(self, config: ASRConfig):
+        """Initialize tokenizer with audio token."""
+        self.tokenizer = AutoTokenizer.from_pretrained(config.text_model_id, trust_remote_code=True)
+
+        # Set pad token
+        if (
+            self.tokenizer.pad_token is None
+            or self.tokenizer.pad_token_id == self.tokenizer.eos_token_id
+        ) and "<|finetune_right_pad_id|>" in self.tokenizer.get_vocab():
+            self.tokenizer.pad_token = "<|finetune_right_pad_id|>"
+
+        # Add audio token
+        existing_special = getattr(self.tokenizer, "additional_special_tokens", None) or []
+        if "<audio>" not in existing_special:
+            self.tokenizer.add_special_tokens(
+                {"additional_special_tokens": existing_special + ["<audio>"]}
+            )
+            self.language_model.resize_token_embeddings(
+                len(self.tokenizer), mean_resizing=False, pad_to_multiple_of=64
+            )
+
+        self.audio_token_id = self.tokenizer.convert_tokens_to_ids("<audio>")
+        self.tokenizer.padding_side = "right"
+
+        # Sync token IDs to configs
+        for cfg in [self.config.text_config, self.language_model.config, self.generation_config]:
+            if cfg is not None:
+                cfg.pad_token_id = self.tokenizer.pad_token_id
+                cfg.eos_token_id = self.tokenizer.eos_token_id
+                cfg.bos_token_id = self.tokenizer.bos_token_id
+
+    def _set_gradient_checkpointing(self, enable: bool = True, gradient_checkpointing_func=None):
+        """Enable/disable gradient checkpointing for the language model."""
+        # The LLM still stores activations during forward for backprop to projector
+        # Gradient checkpointing trades compute for memory by recomputing activations
+        if hasattr(self.language_model, "_set_gradient_checkpointing"):
+            self.language_model._set_gradient_checkpointing(enable, gradient_checkpointing_func)
+        elif hasattr(self.language_model, "gradient_checkpointing_enable") and enable:
+            self.language_model.gradient_checkpointing_enable(
+                gradient_checkpointing_kwargs={"use_reentrant": False}
+            )
+        elif hasattr(self.language_model, "gradient_checkpointing_disable") and not enable:
+            self.language_model.gradient_checkpointing_disable()
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.language_model.set_input_embeddings(value)
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.language_model.get_output_embeddings()
+
+    def set_output_embeddings(self, value: nn.Module) -> None:
+        self.language_model.set_output_embeddings(value)
+
+    def get_processor(self):
+        """Get the processor for this model."""
+        try:
+            from .asr_processing import ASRProcessor
+        except ImportError:
+            from asr_processing import ASRProcessor  # type: ignore[no-redef]
+
+        return ASRProcessor(
+            feature_extractor=self.feature_extractor,
+            tokenizer=self.tokenizer,
+            projector=self.projector,
+            encoder_conv_layers=self.config.encoder_conv_layers,
+        )
+
+    # =========================================================================
+    # Silero VAD for Interruption Detection (Freeze-Omni style)
+    # =========================================================================
+
+    def load_vad(self, force_reload: bool = False) -> None:
+        """Load Silero VAD model for interruption detection.
+
+        Silero VAD is a lightweight (~2MB) voice activity detector that runs
+        in real-time. Used as the first layer of interruption detection.
+
+        Args:
+            force_reload: Force reload even if already loaded
+        """
+        if self._vad_model is not None and not force_reload:
+            return
+
+        model, utils = torch.hub.load(
+            repo_or_dir="snakers4/silero-vad",
+            model="silero_vad",
+            force_reload=force_reload,
+            trust_repo=True,
+        )
+
+        self._vad_model = model
+        self._vad_utils = utils
+
+        # Freeze VAD model
+        self._vad_model.eval()
+        for param in self._vad_model.parameters():
+            param.requires_grad = False
+
+    def detect_speech(
+        self,
+        audio_chunk: torch.Tensor,
+        sample_rate: int = 16000,
+        threshold: float = 0.5,
+    ) -> tuple[bool, float]:
+        """Detect speech in an audio chunk using Silero VAD.
+
+        Args:
+            audio_chunk: Audio waveform [samples] or [1, samples] at sample_rate
+            sample_rate: Audio sample rate (default 16kHz)
+            threshold: Speech probability threshold (default 0.5)
+
+        Returns:
+            Tuple of (is_speech, probability)
+        """
+        if self._vad_model is None:
+            self.load_vad()
+
+        # Ensure 1D tensor
+        if audio_chunk.dim() > 1:
+            audio_chunk = audio_chunk.squeeze()
+
+        # VAD expects specific sample rates (8000 or 16000)
+        if sample_rate not in (8000, 16000):
+            import torchaudio.functional as audio_functional
+
+            audio_chunk = audio_functional.resample(audio_chunk, sample_rate, 16000)
+            sample_rate = 16000
+
+        # Run VAD
+        with torch.no_grad():
+            speech_prob = self._vad_model(audio_chunk, sample_rate).item()
+
+        return speech_prob > threshold, speech_prob
+
+    def reset_vad_state(self) -> None:
+        """Reset VAD internal state between utterances."""
+        if self._vad_model is not None:
+            self._vad_model.reset_states()
+
+    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
+
+    def _compute_encoder_output_lengths(
+        self,
+        audio_attention_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute per-sample encoder output lengths using conv layer formulas.
+
+        Args:
+            audio_attention_mask: Mask indicating real vs padded mel frames (batch, mel_len)
+
+        Returns:
+            Tensor of encoder output lengths per sample (batch,)
+        """
+        # Get mel frame lengths from attention mask
+        lengths = audio_attention_mask.sum(dim=-1)
+
+        # Apply conv layer formulas: output = (input + 2*pad - (kernel-1) - 1) // stride + 1
+        for padding, kernel_size, stride in self.config.encoder_conv_layers:
+            lengths = (lengths + 2 * padding - (kernel_size - 1) - 1) // stride + 1
+
+        return lengths
+
+    def _encode_audio(
+        self,
+        audio_features: torch.Tensor,
+        audio_attention_mask: torch.Tensor,
+        expected_token_counts: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Encode audio and project to LLM embedding space.
+
+        Args:
+            audio_features: Mel spectrogram features (batch, n_mels, mel_len)
+            audio_attention_mask: Mask indicating real vs padded mel frames (batch, mel_len)
+            expected_token_counts: Expected number of audio tokens per sample from input_ids.
+                If provided, output will match these counts exactly (padding/truncating as needed).
+
+        Returns:
+            Flattened audio embeddings of shape (total_audio_tokens, hidden_dim).
+        """
+        with torch.no_grad():
+            encoder_out = self.audio_tower(input_features=audio_features)
+            hidden_states = encoder_out.last_hidden_state
+
+        # Project to LLM space
+        audio_embeds = self.projector(hidden_states)
+
+        # Use expected token counts if provided (from input_ids), otherwise compute from audio
+        if expected_token_counts is not None:
+            token_counts = expected_token_counts
+        else:
+            # Compute per-sample encoder output lengths using conv formulas
+            encoder_lengths = self._compute_encoder_output_lengths(audio_attention_mask)
+            token_counts = torch.tensor(
+                [
+                    self.projector.get_output_length(int(length.item()))
+                    for length in encoder_lengths
+                ],
+                device=audio_embeds.device,
+            )
+
+        # Extract embeddings matching expected token counts per sample
+        batch_size = audio_embeds.shape[0]
+
+        result_embeds = []
+        for i in range(batch_size):
+            count = int(token_counts[i].item())
+            sample_embeds = audio_embeds[i, :count, :]  # Take first 'count' embeddings
+            # Pad with learned embedding if we don't have enough embeddings
+            if sample_embeds.shape[0] < count:
+                pad_count = count - sample_embeds.shape[0]
+                padding = self.audio_pad_embedding.expand(pad_count, -1).to(
+                    device=audio_embeds.device, dtype=audio_embeds.dtype
+                )
+                sample_embeds = torch.cat([sample_embeds, padding], dim=0)
+            result_embeds.append(sample_embeds)
+
+        return torch.cat(result_embeds, dim=0)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        input_features: Optional[torch.Tensor] = None,
+        audio_attention_mask: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_values: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        """Forward pass for training and inference."""
+        # Get text embeddings if not provided
+        if inputs_embeds is None:
+            inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
+
+        if input_features is not None and input_ids is not None:
+            # Apply SpecAugment during training if enabled
+            if self.training and self.spec_augment is not None:
+                input_features = self.spec_augment(input_features)
+
+            # Count expected audio tokens from input_ids (ground truth from collator)
+            audio_token_counts = (input_ids == self.audio_token_id).sum(dim=-1)
+
+            # Encode audio -> flattened (total_audio_tokens, hidden_dim)
+            audio_embeds = self._encode_audio(
+                input_features, audio_attention_mask, audio_token_counts
+            )
+
+            # Replace <audio> token placeholders with audio embeddings using masked_scatter
+            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),
+            )
+
+        # Remove TRL-specific keys that shouldn't go to the LLM
+        kwargs.pop("prompts", None)
+        kwargs.pop("prompt_attention_mask", None)
+
+        # Run through language model (let it compute loss if labels provided)
+        outputs = self.language_model(
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            labels=labels,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        return outputs
+
+    def prepare_inputs_for_generation(self, *args, **kwargs):
+        """Prepare inputs for generation, handling audio features for cached decoding."""
+        input_features = kwargs.pop("input_features", None)
+        cache_position = kwargs.get("cache_position")
+
+        model_inputs = self.language_model.prepare_inputs_for_generation(*args, **kwargs)
+
+        # Only pass audio features on the first generation step (cache_position[0] == 0)
+        if cache_position is not None and cache_position[0] == 0 and input_features is not None:
+            model_inputs["input_features"] = input_features
+
+        return model_inputs
+
+    def _get_num_audio_tokens(
+        self,
+        audio_attention_mask: torch.Tensor,
+    ) -> int:
+        """Calculate number of audio tokens based on actual audio length.
+
+        Uses attention mask to get real audio length, then computes:
+        mel_frames -> encoder_frames (via conv formulas) -> projector output tokens
+        """
+        encoder_lengths = self._compute_encoder_output_lengths(audio_attention_mask)
+        # Use max length for batch (all samples should have same token count for generation)
+        encoder_output_len = int(encoder_lengths.max().item())
+        return int(self.projector.get_output_length(encoder_output_len))
+
+    def _build_audio_prompt(
+        self,
+        audio_attention_mask: torch.Tensor,
+        batch_size: int,
+        device: torch.device,
+        system_prompt: Optional[str] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Build input_ids and attention_mask for audio-conditioned generation.
+
+        Args:
+            audio_attention_mask: Mask for real vs padded mel frames
+            batch_size: Batch size for expanding single prompts
+            device: Device to place tensors on
+            system_prompt: Optional system prompt override
+
+        Returns:
+            Tuple of (input_ids, attention_mask) tensors
+        """
+        num_audio_tokens = self._get_num_audio_tokens(audio_attention_mask)
+        audio_placeholder = "<audio>" * num_audio_tokens
+
+        system_prompt = system_prompt or self.system_prompt
+
+        messages: list[dict[str, str]] = []
+        if system_prompt:
+            messages.append({"role": "system", "content": 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)
+
+        return input_ids, torch.ones_like(input_ids)
+
+    def _inject_audio_embeddings(
+        self,
+        input_ids: torch.Tensor,
+        audio_embeds: torch.Tensor,
+    ) -> torch.Tensor:
+        """Replace audio token placeholders with actual audio embeddings.
+
+        Args:
+            input_ids: Token IDs containing <audio> placeholder tokens
+            audio_embeds: Encoded audio embeddings to inject
+
+        Returns:
+            Input embeddings with audio tokens replaced by audio embeddings
+        """
+        inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
+        audio_token_mask = (input_ids == self.audio_token_id).unsqueeze(-1)
+        return inputs_embeds.masked_scatter(
+            audio_token_mask.to(inputs_embeds.device),
+            audio_embeds.to(inputs_embeds.device, dtype=inputs_embeds.dtype),
+        )
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        input_features: Optional[torch.Tensor] = None,
+        audio_attention_mask: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        system_prompt: Optional[str] = None,
+        **generate_kwargs,
+    ) -> torch.Tensor:
+        """Generate transcription from audio input.
+
+        Can be called in two ways:
+        1. With input_ids containing <audio> tokens (from processor)
+        2. With just audio, and we build the prompt internally
+        """
+        if input_features is None:
+            raise ValueError("input_features required for generation")
+        if audio_attention_mask is None:
+            raise ValueError("audio_attention_mask required for generation")
+
+        device = input_features.device
+        batch_size = input_features.shape[0]
+
+        # Encode audio -> flattened embeddings
+        audio_embeds = self._encode_audio(input_features, audio_attention_mask)
+
+        # If input_ids not provided, build prompt with correct number of audio tokens
+        if input_ids is None:
+            input_ids, attention_mask = self._build_audio_prompt(
+                audio_attention_mask, batch_size, device, system_prompt
+            )
+
+        # Replace audio token placeholders with audio embeddings
+        inputs_embeds = self._inject_audio_embeddings(input_ids, audio_embeds)
+
+        # Generate using language model
+        # Pass both input_ids and inputs_embeds so repetition_penalty works correctly
+        # (it needs input_ids to track which tokens have been used)
+        output = self.language_model.generate(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            generation_config=self.generation_config,
+            **generate_kwargs,
+        )
+
+        # When using inputs_embeds with input_ids, generate returns full sequence
+        # Strip the input tokens to return only generated tokens
+        sequences = output if isinstance(output, torch.Tensor) else output.sequences
+        input_len = input_ids.shape[1]
+        return sequences[:, input_len:]
+
+    def _process_audio(
+        self,
+        audio,
+        sampling_rate: int = 16000,
+    ) -> dict[str, torch.Tensor]:
+        """Process raw audio waveform to model inputs."""
+        # Convert to numpy if tensor
+        if isinstance(audio, torch.Tensor):
+            audio = audio.cpu().numpy()
+
+        # Get mel features from feature extractor
+        inputs = self.feature_extractor(
+            audio,
+            sampling_rate=sampling_rate,
+            return_attention_mask=True,
+            return_tensors="pt",
+        )
+
+        device = next(self.language_model.parameters()).device
+        return {
+            "input_features": inputs["input_features"].to(device),
+            "attention_mask": inputs["attention_mask"].to(device),
+        }
+
+    def save_pretrained(self, save_directory: Union[str, Path], **kwargs) -> None:
+        """Save model, tokenizer, and processor."""
+        import shutil
+        from pathlib import Path as PathlibPath
+
+        save_dir = PathlibPath(save_directory)
+        save_dir.mkdir(parents=True, exist_ok=True)
+
+        # Update config with actual vocab size
+        self.config.vocab_size = self.language_model.config.vocab_size
+        self.config.text_config.vocab_size = self.language_model.config.vocab_size
+
+        if hasattr(self.audio_tower.config, "num_mel_bins"):
+            self.config.audio_config.num_mel_bins = self.audio_tower.config.num_mel_bins
+
+        # Save config
+        self.config.save_pretrained(save_dir)
+
+        # Save state dict directly to avoid HuggingFace's tied weights handling
+        # which conflicts with our shared AudioHead embedding
+        state_dict = self.state_dict()
+        safe_serialization = kwargs.get("safe_serialization", True)
+
+        if safe_serialization:
+            from safetensors.torch import save_file
+
+            save_file(state_dict, save_dir / "model.safetensors")
+        else:
+            import torch
+
+            torch.save(state_dict, save_dir / "pytorch_model.bin")
+
+        # Save tokenizer and feature extractor
+        self.tokenizer.save_pretrained(save_dir)
+        self.feature_extractor.save_pretrained(save_dir)
+
+        # Add processor auto_map to preprocessor_config.json
+        config_path = save_dir / "preprocessor_config.json"
+        if config_path.exists():
+            with config_path.open() as f:
+                processor_config = json.load(f)
+        else:
+            processor_config = {}
+
+        processor_config.update(
+            {
+                "processor_class": "ASRProcessor",
+                "auto_map": {"AutoProcessor": "asr_processing.ASRProcessor"},
+            }
+        )
+
+        with config_path.open("w") as f:
+            json.dump(processor_config, f, indent=2)
+
+        # Copy source files for auto-loading
+        src_dir = PathlibPath(__file__).parent
+        for asr_file in src_dir.glob("asr_*.py"):
+            shutil.copy(asr_file, save_dir / asr_file.name)
+        # Copy projectors module
+        shutil.copy(src_dir / "projectors.py", save_dir / "projectors.py")
+        # Copy alignment module
+        shutil.copy(src_dir / "alignment.py", save_dir / "alignment.py")
+        # Copy diarization module
+        shutil.copy(src_dir / "diarization.py", save_dir / "diarization.py")
+        # Copy audio head for S2S
+        audio_head_path = src_dir / "audio_head.py"
+        if audio_head_path.exists():
+            shutil.copy(audio_head_path, save_dir / "audio_head.py")
+        # Copy full duplex session for S2S
+        full_duplex_path = src_dir / "full_duplex.py"
+        if full_duplex_path.exists():
+            shutil.copy(full_duplex_path, save_dir / "full_duplex.py")
+
+    def push_to_hub(self, repo_id: str, **kwargs) -> str:
+        """Push model to HuggingFace Hub."""
+        self.config.pretrained_model_path = repo_id
+        return super().push_to_hub(repo_id, **kwargs)
+
+    def create_or_update_model_card(self, output_dir: Union[str, Path]) -> None:
+        """No-op for model card creation - we use MODEL_CARD.md in repo instead."""
+        pass
+
+
+# Register with transformers Auto classes
+AutoConfig.register("asr_model", ASRConfig)
+AutoModel.register(ASRConfig, ASRModel)

asr_pipeline.py

@@ -0,0 +1,370 @@
+"""ASR pipeline for audio-to-text transcription with optional timestamps and diarization."""
+
+import re
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import transformers
+
+try:
+    from .alignment import ForcedAligner
+    from .asr_modeling import ASRModel
+    from .diarization import LocalSpeakerDiarizer
+except ImportError:
+    from alignment import ForcedAligner  # type: ignore[no-redef]
+    from asr_modeling import ASRModel  # type: ignore[no-redef]
+    from diarization import LocalSpeakerDiarizer  # type: ignore[no-redef]
+
+# Re-export for backwards compatibility
+__all__ = ["ForcedAligner", "LocalSpeakerDiarizer", "ASRPipeline", "strip_thinking"]
+
+
+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):
+    """ASR Pipeline for audio-to-text transcription."""
+
+    model: ASRModel
+
+    def __init__(self, model: ASRModel, **kwargs):
+        """Initialize ASR pipeline.
+
+        Args:
+            model: ASRModel instance for transcription
+            **kwargs: Additional arguments (feature_extractor, tokenizer, device)
+        """
+        feature_extractor = kwargs.pop("feature_extractor", None)
+        tokenizer = kwargs.pop("tokenizer", model.tokenizer)
+
+        if feature_extractor is None:
+            feature_extractor = model.get_processor().feature_extractor
+
+        super().__init__(
+            model=model, feature_extractor=feature_extractor, tokenizer=tokenizer, **kwargs
+        )
+        self._current_audio = None
+
+    def _sanitize_parameters(self, **kwargs):
+        """Intercept our custom parameters before parent class validates them."""
+        # Remove our custom parameters so parent doesn't see them
+        kwargs.pop("return_timestamps", None)
+        kwargs.pop("return_speakers", None)
+        kwargs.pop("num_speakers", None)
+        kwargs.pop("min_speakers", None)
+        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)
+        return super()._sanitize_parameters(**kwargs)
+
+    def __call__(
+        self,
+        inputs,
+        **kwargs,
+    ):
+        """Transcribe audio with optional word-level timestamps and speaker diarization.
+
+        Args:
+            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
+            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)
+            **kwargs: Additional arguments passed to the pipeline
+
+        Returns:
+            Dict with 'text' key, 'words' key if return_timestamps=True,
+            speaker labels on words if return_speakers=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)
+        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),
+            "max_speakers": kwargs.pop("max_speakers", None),
+        }
+
+        if return_speakers:
+            return_timestamps = True
+
+        # Set custom user prompt if provided
+        original_prompt = None
+        if user_prompt:
+            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)
+
+        # Run standard transcription
+        result = super().__call__(inputs, **kwargs)
+
+        # Add timestamps if requested
+        if return_timestamps and self._current_audio is not None:
+            text = result.get("text", "")
+            if text:
+                try:
+                    words = ForcedAligner.align(
+                        self._current_audio["array"],
+                        text,
+                        sample_rate=self._current_audio.get("sampling_rate", 16000),
+                    )
+                    result["words"] = words
+                except Exception as e:
+                    result["words"] = []
+                    result["timestamp_error"] = str(e)
+            else:
+                result["words"] = []
+
+        # Add speaker diarization if requested
+        if return_speakers and self._current_audio is not None:
+            try:
+                # Run diarization
+                speaker_segments = LocalSpeakerDiarizer.diarize(
+                    self._current_audio["array"],
+                    sample_rate=self._current_audio.get("sampling_rate", 16000),
+                    **{k: v for k, v in diarization_params.items() if v is not None},
+                )
+                result["speaker_segments"] = speaker_segments
+
+                # Assign speakers to words
+                if result.get("words"):
+                    result["words"] = LocalSpeakerDiarizer.assign_speakers_to_words(
+                        result["words"],
+                        speaker_segments,
+                    )
+            except Exception as e:
+                result["speaker_segments"] = []
+                result["diarization_error"] = str(e)
+
+        # Clean up
+        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, use:
+            {"raw_bytes": pcm_bytes, "dtype": "int16", "sampling_rate": 16000}
+        """
+        from transformers.pipelines.audio_utils import ffmpeg_read
+
+        if isinstance(inputs, dict):
+            if "array" in inputs:
+                return {
+                    "array": inputs["array"],
+                    "sampling_rate": inputs.get("sampling_rate", 16000),
+                }
+            if "raw" in inputs:
+                return {
+                    "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:
+                audio = ffmpeg_read(f.read(), sampling_rate=16000)
+            return {"array": audio, "sampling_rate": 16000}
+        elif isinstance(inputs, bytes):
+            audio = ffmpeg_read(inputs, sampling_rate=16000)
+            return {"array": audio, "sampling_rate": 16000}
+        elif isinstance(inputs, np.ndarray):
+            return {"array": inputs, "sampling_rate": 16000}
+
+        return None
+
+    def preprocess(self, inputs, **preprocess_params):
+        """Preprocess audio inputs for the model.
+
+        Args:
+            inputs: Audio input (dict with array, file path, etc.)
+            **preprocess_params: Additional preprocessing parameters
+
+        Yields:
+            Model input dicts with input_features and attention_mask
+        """
+        # Handle dict with "array" key (from datasets)
+        if isinstance(inputs, dict) and "array" in inputs:
+            inputs = {
+                "raw": inputs["array"],
+                "sampling_rate": inputs.get("sampling_rate", self.feature_extractor.sampling_rate),
+            }
+
+        for item in super().preprocess(inputs, **preprocess_params):
+            if "is_last" not in item:
+                item["is_last"] = True
+            yield item
+
+    def _forward(self, model_inputs, **generate_kwargs) -> dict[str, Any]:
+        """Run model forward pass to generate transcription.
+
+        Args:
+            model_inputs: Dict with input_features and attention_mask
+            **generate_kwargs: Generation parameters
+
+        Returns:
+            Dict with generated token IDs
+        """
+        # Extract audio features and is_last flag
+        is_last = model_inputs.pop("is_last", True) if isinstance(model_inputs, dict) else True
+
+        input_features = model_inputs["input_features"].to(self.model.device)
+        audio_attention_mask = model_inputs["attention_mask"].to(self.model.device)
+
+        generated_ids = self.model.generate(
+            input_features=input_features,
+            audio_attention_mask=audio_attention_mask,
+            **generate_kwargs,
+        )
+
+        return {"tokens": generated_ids, "is_last": is_last}
+
+    def postprocess(self, model_outputs, **kwargs) -> dict[str, str]:
+        """Convert model output tokens to text.
+
+        Args:
+            model_outputs: Dict with 'tokens' key containing generated IDs
+            **kwargs: Additional postprocessing parameters
+
+        Returns:
+            Dict with 'text' key containing transcription
+        """
+        # Handle list of outputs (from chunking)
+        if isinstance(model_outputs, list):
+            model_outputs = model_outputs[0] if model_outputs else {}
+
+        tokens = model_outputs.get("tokens")
+        if tokens is None:
+            return super().postprocess(model_outputs, **kwargs)
+
+        if torch.is_tensor(tokens):
+            tokens = tokens.cpu()
+            if tokens.dim() > 1:
+                tokens = tokens[0]
+
+        # Filter out eos tokens that the tokenizer doesn't recognize as special
+        # (generation_config.eos_token_id may differ from tokenizer.eos_token_id)
+        if hasattr(self, "model") and hasattr(self.model, "generation_config"):
+            eos_ids = self.model.generation_config.eos_token_id
+            if eos_ids is not None:
+                eos_set = set(eos_ids) if isinstance(eos_ids, list) else {eos_ids}
+                tokens = [t for t in tokens.tolist() if t not in eos_set]
+
+        text = self.tokenizer.decode(tokens, skip_special_tokens=True).strip()
+        # Strip <think>...</think> tags (Qwen3 doesn't respect /no_think prompt)
+        text = strip_thinking(text)
+        # Truncate repetitions at end of text
+        text = _truncate_repetitions(text)
+        return {"text": text}
+
+
+def _truncate_repetitions(text: str, min_repeats: int = 3) -> str:
+    """Truncate repeated words/phrases/characters at end of text.
+
+    Detects patterns like:
+    - Repeated words: "the the the the" -> "the"
+    - Repeated phrases: "i am sorry i am sorry i am sorry" -> "i am sorry"
+    - Repeated characters: "444444" -> "4"
+
+    Args:
+        text: Input text to process
+        min_repeats: Minimum repetitions to trigger truncation (default 3)
+
+    Returns:
+        Text with trailing repetitions removed
+    """
+    if not text:
+        return text
+
+    # 1. Truncate repeated characters at end (e.g., "444444" -> "4")
+    char_pattern = re.compile(r"(.)\1{" + str(min_repeats - 1) + r",}$")
+    text = char_pattern.sub(r"\1", text)
+
+    # 2. Truncate repeated words at end (e.g., "the the the" -> "the")
+    word_pattern = re.compile(
+        r"\b(\w+)(?:\s+\1){" + str(min_repeats - 1) + r",}\s*$", re.IGNORECASE
+    )
+    while word_pattern.search(text):
+        text = word_pattern.sub(r"\1", text)
+
+    # 3. Truncate repeated phrases (2-20 words) at end
+    # e.g., "i am sorry i am sorry i am sorry" -> "i am sorry"
+    words = text.split()
+    if len(words) >= min_repeats * 2:
+        # Try phrase lengths from 2 to 20 words
+        for phrase_len in range(2, min(21, len(words) // min_repeats + 1)):
+            # Check if the last phrase_len words repeat
+            phrase = " ".join(words[-phrase_len:])
+            # Build pattern to match repeated phrases at end
+            phrase_escaped = re.escape(phrase)
+            phrase_pattern = re.compile(
+                r"(^|.*?\s)("
+                + phrase_escaped
+                + r")(?:\s+"
+                + phrase_escaped
+                + r"){"
+                + str(min_repeats - 1)
+                + r",}\s*$",
+                re.IGNORECASE,
+            )
+            match = phrase_pattern.match(text)
+            if match:
+                # Keep prefix + one instance of the phrase
+                text = (match.group(1) + match.group(2)).strip()
+                words = text.split()
+                break
+
+    return text

asr_processing.py

@@ -0,0 +1,133 @@
+from typing import Optional, Union
+
+import torch
+import transformers
+from transformers import ProcessorMixin
+
+try:
+    from .asr_config import ASRConfig
+except ImportError:
+    from asr_config import ASRConfig  # type: ignore[no-redef]
+
+
+class ASRProcessor(ProcessorMixin):
+    """Processor for Whisper-based ASR models."""
+
+    attributes = ["feature_extractor", "tokenizer"]
+    feature_extractor_class = "AutoFeatureExtractor"
+    tokenizer_class = "AutoTokenizer"
+    AUDIO_TOKEN = "<audio>"
+    TRANSCRIBE_PROMPT = ""
+    # Default conv layers for Whisper/GLM-ASR: [(pad, kernel, stride), ...]
+    DEFAULT_ENCODER_CONV_LAYERS = [(1, 3, 1), (1, 3, 2)]
+
+    def __init__(
+        self,
+        feature_extractor,
+        tokenizer,
+        projector=None,
+        encoder_conv_layers: Optional[list] = None,
+    ):
+        """Initialize the ASR processor.
+
+        Args:
+            feature_extractor: Audio feature extractor (WhisperFeatureExtractor)
+            tokenizer: Text tokenizer for the language model
+            projector: Audio projector module (for computing output lengths)
+            encoder_conv_layers: Conv layer specs [(pad, kernel, stride), ...]
+        """
+        self.feature_extractor = feature_extractor
+        self.tokenizer = tokenizer
+        self.audio_token_id = tokenizer.convert_tokens_to_ids(self.AUDIO_TOKEN)
+        self.projector = projector
+        self.encoder_conv_layers = encoder_conv_layers or self.DEFAULT_ENCODER_CONV_LAYERS
+
+    def _compute_encoder_output_length(self, mel_length: int) -> int:
+        """Compute encoder output length using conv layer formulas."""
+        length = mel_length
+        for padding, kernel_size, stride in self.encoder_conv_layers:
+            length = (length + 2 * padding - (kernel_size - 1) - 1) // stride + 1
+        return length
+
+    def __call__(
+        self,
+        audio: Optional[Union[list, "torch.Tensor"]] = None,
+        text: Optional[str] = None,
+        system_prompt: Optional[str] = None,
+        return_tensors: str = "pt",
+        **kwargs,
+    ) -> dict:
+        """Process audio and text inputs for inference.
+
+        Args:
+            audio: Raw audio waveform(s)
+            text: Target transcription (optional, for training - but use DataCollator instead)
+            system_prompt: Optional system prompt
+            return_tensors: Return format ("pt" for PyTorch)
+
+        Returns:
+            Dict with input_features, input_ids, attention_mask
+        """
+        result = {}
+
+        # Process audio
+        if audio is not None:
+            audio_inputs = self.feature_extractor(
+                audio,
+                sampling_rate=getattr(self.feature_extractor, "sampling_rate", 16000),
+                return_attention_mask=True,
+                return_tensors=return_tensors,
+                **kwargs,
+            )
+            result["input_features"] = audio_inputs["input_features"]
+            result["audio_attention_mask"] = audio_inputs["attention_mask"]
+
+            # Use actual audio length (from attention mask) for token count
+            real_mel_len = int(audio_inputs["attention_mask"].sum(dim=-1).max().item())
+            encoder_output_len = self._compute_encoder_output_length(real_mel_len)
+            num_audio_tokens = self.projector.get_output_length(encoder_output_len)
+        else:
+            num_audio_tokens = 0
+
+        # Build prompt with audio token placeholders (instruction-free)
+        if num_audio_tokens > 0:
+            user_content = self.AUDIO_TOKEN * num_audio_tokens
+            if self.TRANSCRIBE_PROMPT:
+                user_content += " " + self.TRANSCRIBE_PROMPT
+        else:
+            user_content = self.TRANSCRIBE_PROMPT or ""
+
+        messages = []
+        if system_prompt:
+            messages.append({"role": "system", "content": system_prompt})
+        messages.append({"role": "user", "content": user_content})
+        if text is not None:
+            messages.append({"role": "assistant", "content": text})
+
+        # Tokenize
+        tokenized = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_generation_prompt=(text is None),
+            return_tensors=return_tensors,
+            enable_thinking=False,  # Disable Qwen3 thinking mode for ASR
+        )
+
+        # Handle both tensor and BatchEncoding returns
+        if isinstance(tokenized, torch.Tensor):
+            input_ids = tokenized
+        else:
+            # BatchEncoding or dict-like object
+            input_ids = tokenized.get("input_ids", tokenized.input_ids)
+
+        if input_ids.dim() == 1:
+            input_ids = input_ids.unsqueeze(0)
+
+        result["input_ids"] = input_ids
+        result["attention_mask"] = torch.ones_like(input_ids)
+
+        return result
+
+
+ASRProcessor.register_for_auto_class()
+transformers.AutoProcessor.register(ASRConfig, ASRProcessor)

audio_head.py

@@ -0,0 +1,357 @@
+"""Audio head for speech-to-speech using a trainable AR decoder + NeuCodec.
+
+Generates audio from text tokens via a trainable LlamaModel decoder:
+  Text tokens -> Embedding -> LlamaModel -> head -> NeuCodec FSQ codes -> audio
+
+NeuCodec uses a single FSQ codebook (levels=[4]*8, vocab=65536) at 50 tokens/sec,
+outputting 24kHz audio. No multi-codebook handling needed.
+
+Training: S2SDataCollator prepares codec_input_ids/codec_labels (both 2D: [batch, seq_len]).
+AudioHead predicts FSQ codes via a single head with teacher forcing.
+
+Inference: Autoregressive generation with KV cache, feeding back predicted codes.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Iterator, Optional
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F  # noqa: N812
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+logger = logging.getLogger(__name__)
+
+# NeuCodec FSQ constants (levels=[4]*8, 1 quantizer -> 4^8 = 65536 codes)
+NEUCODEC_VOCAB_SIZE = 65536
+NEUCODEC_SAMPLE_RATE = 24000
+
+# Special tokens (above vocab range)
+BOS_TOKEN = NEUCODEC_VOCAB_SIZE
+EOS_TOKEN = NEUCODEC_VOCAB_SIZE + 1
+PAD_TOKEN = NEUCODEC_VOCAB_SIZE + 2
+TOTAL_VOCAB = NEUCODEC_VOCAB_SIZE + 3  # 65539
+
+
+class AudioHeadConfig(PretrainedConfig):
+    """Configuration class for the AudioHead model."""
+
+    model_type = "audio_head"
+
+    def __init__(
+        self,
+        decoder_dim: int = 512,
+        decoder_layers: int = 6,
+        decoder_heads: int = 8,
+        text_vocab_size: int = 32000,
+        max_audio_tokens: int = 500,
+        neucodec_model_id: str = "neuphonic/neucodec",
+        temperature: float = 1.0,
+        top_k: int = 50,
+        **kwargs,
+    ):
+        self.decoder_dim = decoder_dim
+        self.decoder_layers = decoder_layers
+        self.decoder_heads = decoder_heads
+        self.text_vocab_size = text_vocab_size
+        self.max_audio_tokens = max_audio_tokens
+        self.neucodec_model_id = neucodec_model_id
+        self.temperature = temperature
+        self.top_k = top_k
+        super().__init__(**kwargs)
+
+
+@dataclass
+class AudioHeadOutput(ModelOutput):
+    """Output of AudioHead forward pass.
+
+    Attributes:
+        loss: Cross-entropy loss when codec_labels are provided.
+        codes: Generated codec codes when in inference mode [batch, gen_len].
+    """
+
+    loss: Optional[torch.Tensor] = None
+    codes: Optional[torch.Tensor] = None
+
+
+class AudioHead(PreTrainedModel):
+    """Trainable AR decoder that predicts NeuCodec FSQ codes.
+
+    NeuCodec uses a single FSQ codebook (4^8 = 65536 codes) at 50 tokens/sec.
+    No multi-codebook handling needed — just a flat sequence of codes.
+    """
+
+    config_class = AudioHeadConfig
+
+    def __init__(self, config: AudioHeadConfig):
+        super().__init__(config)
+        self.text_vocab_size = config.text_vocab_size
+        self.decoder_dim = config.decoder_dim
+        self.max_tokens = config.max_audio_tokens
+        self.vocab_size = NEUCODEC_VOCAB_SIZE
+
+        # Embed text tokens to decoder dim
+        self.text_embedding = nn.Embedding(config.text_vocab_size, config.decoder_dim)
+
+        # Codec token embedding (FSQ codes + special tokens)
+        self.token_embedding = nn.Embedding(TOTAL_VOCAB, config.decoder_dim)
+
+        # Small LlamaModel as decoder backbone (from config, NOT pretrained)
+        from transformers import LlamaConfig, LlamaModel
+
+        llama_config = LlamaConfig(
+            hidden_size=config.decoder_dim,
+            intermediate_size=config.decoder_dim * 4,
+            num_hidden_layers=config.decoder_layers,
+            num_attention_heads=config.decoder_heads,
+            vocab_size=TOTAL_VOCAB,
+            max_position_embeddings=4096,
+        )
+        self.decoder = LlamaModel(llama_config)
+        # We handle embeddings ourselves, remove the unused one to save memory
+        self.decoder.embed_tokens = None
+
+        # Sampling parameters for inference
+        self.temperature = config.temperature
+        self.top_k = config.top_k
+
+        # NeuCodec model (loaded lazily, frozen, inference only)
+        self.neucodec_model = None
+
+        # Initialize weights
+        self.post_init()
+
+    def forward(
+        self,
+        text_token_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        codec_labels: Optional[torch.Tensor] = None,
+        codec_input_ids: Optional[torch.Tensor] = None,
+        codec_attention_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> AudioHeadOutput:
+        """Forward pass for training or inference.
+
+        Args:
+            text_token_ids: Text token IDs [batch, seq_len]
+            attention_mask: Text attention mask [batch, seq_len] (1=real, 0=padding)
+            codec_labels: Target codes [batch, audio_len] (-100 for ignore)
+            codec_input_ids: Teacher-forced input [batch, audio_len]
+            codec_attention_mask: Codec attention mask [batch, audio_len]
+
+        Returns:
+            AudioHeadOutput with loss (training) or codes (inference).
+        """
+        # Embed text tokens (clamp to valid range)
+        if (text_token_ids >= self.text_vocab_size).any() or (text_token_ids < 0).any():
+            logger.warning(
+                "text_token_ids out of range [0, %d): min=%d max=%d. Clamping.",
+                self.text_vocab_size, text_token_ids.min().item(), text_token_ids.max().item(),
+            )
+            text_token_ids = text_token_ids.clamp(0, self.text_vocab_size - 1)
+        prefix = self.text_embedding(text_token_ids)  # [batch, text_len, decoder_dim]
+        batch_size, text_len, _ = prefix.shape
+
+        if codec_labels is not None:
+            # Teacher forcing: codec_input_ids is [batch, audio_len]
+            cb_input = codec_input_ids
+            if (cb_input >= TOTAL_VOCAB).any() or (cb_input < 0).any():
+                logger.warning(
+                    "codec_input_ids out of range [0, %d): min=%d max=%d. Clamping.",
+                    TOTAL_VOCAB, cb_input.min().item(), cb_input.max().item(),
+                )
+                cb_input = cb_input.clamp(0, TOTAL_VOCAB - 1)
+            token_emb = self.token_embedding(cb_input)  # [batch, audio_len, dim]
+
+            audio_len = token_emb.shape[1]
+
+            # Concatenate prefix + codec tokens
+            hidden = torch.cat([prefix, token_emb], dim=1)  # [batch, text+audio, dim]
+
+            # Build combined attention mask
+            if attention_mask is not None:
+                prefix_mask = attention_mask
+            else:
+                prefix_mask = torch.ones(
+                    batch_size, text_len, device=hidden.device, dtype=torch.long
+                )
+
+            if codec_attention_mask is not None:
+                audio_mask = codec_attention_mask
+            else:
+                audio_mask = torch.ones(
+                    batch_size, audio_len, device=hidden.device, dtype=torch.long
+                )
+
+            combined_mask = torch.cat([prefix_mask, audio_mask], dim=1)
+
+            # Build causal mask for codec positions while prefix attends bidirectionally
+            total_len = text_len + audio_len
+            causal_mask = torch.triu(
+                torch.full((total_len, total_len), float("-inf"), device=hidden.device),
+                diagonal=1,
+            )
+            causal_mask[:text_len, :text_len] = 0.0
+            causal_mask = causal_mask.unsqueeze(0).unsqueeze(0).expand(batch_size, 1, -1, -1)
+
+            padding_mask = (1 - combined_mask).bool()
+            padding_mask_expanded = padding_mask.unsqueeze(1).unsqueeze(2).expand_as(causal_mask)
+            causal_mask = causal_mask.masked_fill(padding_mask_expanded, float("-inf"))
+
+            position_ids = (
+                torch.arange(total_len, device=hidden.device).unsqueeze(0).expand(batch_size, -1)
+            )
+
+            # Run through LlamaModel
+            outputs = self.decoder(
+                inputs_embeds=hidden,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+            )
+
+            # Extract audio positions only
+            audio_hidden = outputs.last_hidden_state[:, text_len:]  # [batch, audio_len, dim]
+
+            # Predict codes and compute loss
+            labels = codec_labels.clone()  # [batch, audio_len]
+            valid_mask = labels != -100
+            labels[valid_mask] = labels[valid_mask].clamp(0, TOTAL_VOCAB - 1)
+
+            logits = F.linear(audio_hidden, self.token_embedding.weight)  # [batch, audio_len, total_vocab]
+            loss = F.cross_entropy(
+                logits.reshape(-1, TOTAL_VOCAB),
+                labels.reshape(-1),
+                ignore_index=-100,
+            )
+            return AudioHeadOutput(loss=loss)
+
+        # Inference: autoregressive generation
+        codes = self._generate(prefix, attention_mask)
+        return AudioHeadOutput(codes=codes)
+
+    def _generate(
+        self, prefix: torch.Tensor, prefix_mask: Optional[torch.Tensor]
+    ) -> torch.Tensor:
+        """AR generation: predict codes one timestep at a time with KV cache."""
+        batch_size, text_len, _ = prefix.shape
+        device = prefix.device
+
+        all_codes = []
+
+        # Build initial input: prefix + BOS embedding
+        bos_token = torch.full((batch_size, 1), BOS_TOKEN, dtype=torch.long, device=device)
+        bos_emb = self.token_embedding(bos_token)  # [batch, 1, dim]
+        hidden = torch.cat([prefix, bos_emb], dim=1)  # [batch, text_len+1, dim]
+
+        # Position IDs for initial forward
+        position_ids = torch.arange(text_len + 1, device=device).unsqueeze(0).expand(batch_size, -1)
+
+        # Initial forward pass (no KV cache yet)
+        outputs = self.decoder(
+            inputs_embeds=hidden,
+            position_ids=position_ids,
+            use_cache=True,
+        )
+        past_key_values = outputs.past_key_values
+        last_hidden = outputs.last_hidden_state[:, -1:]  # [batch, 1, dim]
+
+        for step in range(self.max_tokens):
+            # Predict code token
+            logits = F.linear(last_hidden.squeeze(1), self.token_embedding.weight)  # [batch, vocab]
+
+            # Apply temperature and top-k sampling
+            if self.temperature > 0 and self.top_k > 0:
+                logits = logits / self.temperature
+                # Zero out logits below top-k threshold
+                top_k_vals, _ = logits.topk(self.top_k, dim=-1)
+                logits[logits < top_k_vals[:, -1:]] = float("-inf")
+                probs = F.softmax(logits, dim=-1)
+                token = torch.multinomial(probs, num_samples=1).squeeze(-1)  # [batch]
+            else:
+                token = logits.argmax(dim=-1)  # [batch]
+
+            # Check for EOS
+            if (token == EOS_TOKEN).all():
+                break
+
+            all_codes.append(token)
+
+            # Feed back prediction for next step
+            next_emb = self.token_embedding(token.unsqueeze(1))  # [batch, 1, dim]
+
+            next_pos = torch.full(
+                (batch_size, 1), text_len + 1 + step + 1, dtype=torch.long, device=device
+            )
+
+            # Forward with KV cache
+            outputs = self.decoder(
+                inputs_embeds=next_emb,
+                position_ids=next_pos,
+                past_key_values=past_key_values,
+                use_cache=True,
+            )
+            past_key_values = outputs.past_key_values
+            last_hidden = outputs.last_hidden_state  # [batch, 1, dim]
+
+        if all_codes:
+            # [batch, gen_len]
+            codes = torch.stack(all_codes, dim=1)
+        else:
+            codes = torch.empty(batch_size, 0, dtype=torch.long, device=device)
+
+        return codes
+
+    def _load_neucodec(self):
+        """Load frozen NeuCodec model for audio decoding."""
+        from neucodec import NeuCodec
+
+        self.neucodec_model = NeuCodec.from_pretrained(self.config.neucodec_model_id)
+        self.neucodec_model.eval()
+        self.neucodec_model.requires_grad_(False)
+        logger.info("Loaded frozen NeuCodec model for audio decoding")
+
+    def decode_to_audio(self, codes: torch.Tensor) -> list[torch.Tensor]:
+        """Decode NeuCodec FSQ tokens to audio waveforms.
+
+        Args:
+            codes: Codec tokens [batch, seq_len]
+
+        Returns:
+            List of audio waveform tensors (one per batch item)
+        """
+        if self.neucodec_model is None:
+            self._load_neucodec()
+        assert self.neucodec_model is not None
+
+        # NeuCodec decode_code expects [batch, 1, seq_len]
+        codes_3d = codes.unsqueeze(1).to(self.neucodec_model.device)
+
+        with torch.no_grad():
+            audio_values = self.neucodec_model.decode_code(codes_3d)  # [batch, 1, samples]
+
+        return [audio_values[i, 0] for i in range(audio_values.shape[0])]
+
+    def generate_streaming(
+        self,
+        text_token_ids: torch.Tensor,
+        chunk_samples: int = 24000,
+    ) -> Iterator[torch.Tensor]:
+        """Generate audio and yield waveform chunks for streaming playback.
+
+        Args:
+            text_token_ids: Text token IDs [batch, seq_len]
+            chunk_samples: Audio samples per chunk (default 1s at 24kHz)
+
+        Yields:
+            Audio waveform chunks [samples]
+        """
+        output = self(text_token_ids)
+        codes = output.codes
+        audios = self.decode_to_audio(codes)
+
+        for audio in audios:
+            for start in range(0, audio.shape[-1], chunk_samples):
+                end = min(start + chunk_samples, audio.shape[-1])
+                yield audio[..., start:end]

chat_template.jinja

@@ -0,0 +1,94 @@
+{# ───── defaults ───── #}
+{%- if enable_thinking is not defined -%}
+{%- set enable_thinking = true -%}
+{%- endif -%}
+
+{# ───── reasoning mode ───── #}
+{%- if enable_thinking -%}
+  {%- set reasoning_mode = "/think" -%}
+{%- else -%}
+  {%- set reasoning_mode = "/no_think" -%}
+{%- endif -%}
+
+{# ───── header (system message) ───── #}
+{{- "<|im_start|>system\n" -}}
+
+{%- if messages[0].role == "system" -%}
+  {%- set system_message = messages[0].content -%}
+  {%- if "/no_think" in system_message -%}
+    {%- set reasoning_mode = "/no_think" -%}
+  {%- elif "/think" in system_message -%}
+    {%- set reasoning_mode = "/think" -%}
+  {%- endif -%}
+  {%- set custom_instructions = system_message.replace("/no_think", "").replace("/think", "").rstrip() -%}
+{%- endif -%}
+
+{%- if "/system_override" in system_message -%}
+  {{- custom_instructions.replace("/system_override", "").rstrip() -}}
+  {{- "<|im_end|>\n" -}}
+{%- else -%}
+  {{- "## Metadata\n\n" -}}
+  {{- "Knowledge Cutoff Date: June 2025\n" -}}
+  {%- set today = strftime_now("%d %B %Y") -%}
+  {{- "Today Date: " ~ today ~ "\n" -}}
+  {{- "Reasoning Mode: " + reasoning_mode + "\n\n" -}}
+  
+  {{- "## Custom Instructions\n\n" -}}
+  {%- if custom_instructions -%}
+    {{- custom_instructions + "\n\n" -}}
+  {%- elif reasoning_mode == "/think" -%}
+    {{- "You are a helpful AI assistant named SmolLM, trained by Hugging Face. Your role as an assistant involves thoroughly exploring questions through a systematic thinking process before providing the final precise and accurate solutions. This requires engaging in a comprehensive cycle of analysis, summarizing, exploration, reassessment, reflection, backtracking, and iteration to develop well-considered thinking process. Please structure your response into two main sections: Thought and Solution using the specified format: <think> Thought section </think> Solution section. In the Thought section, detail your reasoning process in steps. Each step should include detailed considerations such as analysing questions, summarizing relevant findings, brainstorming new ideas, verifying the accuracy of the current steps, refining any errors, and revisiting previous steps. In the Solution section, based on various attempts, explorations, and reflections from the Thought section, systematically present the final solution that you deem correct. The Solution section should be logical, accurate, and concise and detail necessary steps needed to reach the conclusion.\n\n" -}}
+  {%- else -%}
+    {{- "You are a helpful AI assistant named SmolLM, trained by Hugging Face.\n\n" -}}
+  {%- endif -%}
+
+  {%- if xml_tools or python_tools or tools -%}
+    {{- "### Tools\n\n" -}}
+    {%- if xml_tools or tools -%}
+      {%- if tools -%}
+        {%- set xml_tools = tools -%}
+      {%- endif -%}
+      {%- set ns = namespace(xml_tool_string="You may call one or more functions to assist with the user query.\nYou are provided with function signatures within <tools></tools> XML tags:\n\n<tools>\n") -%}
+      {%- for tool in xml_tools[:] -%} {# The slicing makes sure that xml_tools is a list #}
+        {%- set ns.xml_tool_string = ns.xml_tool_string ~ (tool | string) ~ "\n" -%}
+      {%- endfor -%}
+      {%- set xml_tool_string = ns.xml_tool_string + "</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call>" -%}
+      {{- xml_tool_string -}}
+    {%- endif -%}
+    {%- if python_tools -%}
+      {%- set ns = namespace(python_tool_string="When you send a message containing Python code between '<code>' and '</code>' tags, it will be executed in a stateful Jupyter notebook environment, and you will then be given the output to continued reasoning in an agentic loop.\n\nYou can use the following tools in your python code like regular functions:\n<tools>\n") -%}
+      {%- for tool in python_tools[:] -%} {# The slicing makes sure that python_tools is a list #}
+        {%- set ns.python_tool_string = ns.python_tool_string ~ (tool | string) ~ "\n" -%}
+      {%- endfor -%}
+      {%- set python_tool_string = ns.python_tool_string + "</tools>\n\nThe state persists between code executions: so variables that you define in one step are still available thereafter." -%}
+      {{- python_tool_string -}}
+    {%- endif -%}
+    {{- "\n\n" -}}
+    {{- "<|im_end|>\n" -}}
+  {%- endif -%}
+{%- endif -%}
+{# ───── main loop ───── #}
+{%- for message in messages -%}
+    {%- set content = message.content if message.content is string else "" -%}
+    {%- if message.role == "user" -%}
+        {{ "<|im_start|>" + message.role + "\n"  + content + "<|im_end|>\n" }}
+    {%- elif message.role == "assistant" -%}
+        {% generation %}
+        {%- if reasoning_mode == "/think" -%}
+            {{ "<|im_start|>assistant\n" + content.lstrip("\n") + "<|im_end|>\n" }}
+        {%- else -%}
+            {{ "<|im_start|>assistant\n" + "<think>\n\n</think>\n" + content.lstrip("\n") + "<|im_end|>\n" }}
+        {%- endif -%}
+        {% endgeneration %}
+    {%- elif message.role == "tool" -%}
+    {{ "<|im_start|>" + "user\n"  + content + "<|im_end|>\n" }}
+    {%- endif -%}
+{%- endfor -%}
+{# ───── generation prompt ───── #}
+{%- if add_generation_prompt -%}
+    {%- if reasoning_mode == "/think" -%}
+        {{ "<|im_start|>assistant\n" }}
+    {%- else -%}
+        {{ "<|im_start|>assistant\n" + "<think>\n\n</think>\n"  }}
+    {%- endif -%}
+{%- endif -%}

config.json

@@ -0,0 +1,409 @@
+{
+  "architectures": [
+    "ASRModel"
+  ],
+  "attn_implementation": null,
+  "audio_config": {
+    "_name_or_path": "zai-org/GLM-ASR-Nano-2512",
+    "architectures": [
+      "GlmAsrForConditionalGeneration"
+    ],
+    "audio_config": {
+      "_name_or_path": "",
+      "architectures": null,
+      "attention_dropout": 0.0,
+      "chunk_size_feed_forward": 0,
+      "dtype": null,
+      "head_dim": 64,
+      "hidden_act": "gelu",
+      "hidden_size": 1280,
+      "id2label": {
+        "0": "LABEL_0",
+        "1": "LABEL_1"
+      },
+      "initializer_range": 0.02,
+      "intermediate_size": 5120,
+      "is_encoder_decoder": false,
+      "label2id": {
+        "LABEL_0": 0,
+        "LABEL_1": 1
+      },
+      "max_position_embeddings": 1500,
+      "model_type": "glmasr_encoder",
+      "num_attention_heads": 20,
+      "num_hidden_layers": 32,
+      "num_key_value_heads": 20,
+      "num_mel_bins": 128,
+      "output_attentions": false,
+      "output_hidden_states": false,
+      "partial_rotary_factor": 0.5,
+      "problem_type": null,
+      "return_dict": true,
+      "rope_parameters": {
+        "partial_rotary_factor": 0.5,
+        "rope_theta": 10000.0,
+        "rope_type": "default"
+      }
+    },
+    "audio_token_id": 59260,
+    "dtype": "bfloat16",
+    "hidden_size": 2048,
+    "model_type": "glmasr",
+    "num_mel_bins": 128,
+    "projector_hidden_act": "gelu",
+    "text_config": {
+      "_name_or_path": "",
+      "architectures": null,
+      "attention_bias": false,
+      "attention_dropout": 0.0,
+      "bos_token_id": 1,
+      "chunk_size_feed_forward": 0,
+      "dtype": null,
+      "eos_token_id": [
+        59246,
+        59253,
+        59255
+      ],
+      "head_dim": 128,
+      "hidden_act": "silu",
+      "hidden_size": 2048,
+      "id2label": {
+        "0": "LABEL_0",
+        "1": "LABEL_1"
+      },
+      "initializer_range": 0.02,
+      "intermediate_size": 6144,
+      "is_encoder_decoder": false,
+      "label2id": {
+        "LABEL_0": 0,
+        "LABEL_1": 1
+      },
+      "max_position_embeddings": 8192,
+      "mlp_bias": false,
+      "model_type": "llama",
+      "num_attention_heads": 16,
+      "num_hidden_layers": 28,
+      "num_key_value_heads": 4,
+      "output_attentions": false,
+      "output_hidden_states": false,
+      "pad_token_id": null,
+      "pretraining_tp": 1,
+      "problem_type": null,
+      "return_dict": true,
+      "rms_norm_eps": 1e-05,
+      "rope_parameters": {
+        "rope_theta": 10000.0,
+        "rope_type": "default"
+      },
+      "tie_word_embeddings": false,
+      "use_cache": true,
+      "vocab_size": 59264
+    },
+    "vocab_size": 59264
+  },
+  "audio_model_id": "zai-org/GLM-ASR-Nano-2512",
+  "audio_sample_rate": 16000,
+  "auto_map": {
+    "AutoConfig": "asr_config.ASRConfig",
+    "AutoModel": "asr_modeling.ASRModel",
+    "AutoModelForSpeechSeq2Seq": "asr_modeling.ASRModel",
+    "AutoProcessor": "asr_processing.ASRProcessor"
+  },
+  "custom_pipelines": {
+    "automatic-speech-recognition": {
+      "impl": "asr_pipeline.ASRPipeline",
+      "pt": [
+        "AutoModelForSpeechSeq2Seq"
+      ],
+      "tf": [],
+      "type": "audio"
+    }
+  },
+  "decoder_dim": 256,
+  "decoder_heads": 4,
+  "decoder_layers": 4,
+  "do_sample": false,
+  "downsample_rate": 5,
+  "dtype": "bfloat16",
+  "enable_thinking": false,
+  "encoder": {
+    "_name_or_path": "zai-org/GLM-ASR-Nano-2512",
+    "architectures": [
+      "GlmAsrForConditionalGeneration"
+    ],
+    "audio_config": {
+      "_name_or_path": "",
+      "architectures": null,
+      "attention_dropout": 0.0,
+      "chunk_size_feed_forward": 0,
+      "dtype": null,
+      "head_dim": 64,
+      "hidden_act": "gelu",
+      "hidden_size": 1280,
+      "id2label": {
+        "0": "LABEL_0",
+        "1": "LABEL_1"
+      },
+      "initializer_range": 0.02,
+      "intermediate_size": 5120,
+      "is_encoder_decoder": false,
+      "label2id": {
+        "LABEL_0": 0,
+        "LABEL_1": 1
+      },
+      "max_position_embeddings": 1500,
+      "model_type": "glmasr_encoder",
+      "num_attention_heads": 20,
+      "num_hidden_layers": 32,
+      "num_key_value_heads": 20,
+      "num_mel_bins": 128,
+      "output_attentions": false,
+      "output_hidden_states": false,
+      "partial_rotary_factor": 0.5,
+      "problem_type": null,
+      "return_dict": true,
+      "rope_parameters": {
+        "partial_rotary_factor": 0.5,
+        "rope_theta": 10000.0,
+        "rope_type": "default"
+      }
+    },
+    "audio_token_id": 59260,
+    "dtype": "bfloat16",
+    "hidden_size": 2048,
+    "model_type": "glmasr",
+    "num_mel_bins": 128,
+    "projector_hidden_act": "gelu",
+    "text_config": {
+      "_name_or_path": "",
+      "architectures": null,
+      "attention_bias": false,
+      "attention_dropout": 0.0,
+      "bos_token_id": 1,
+      "chunk_size_feed_forward": 0,
+      "dtype": null,
+      "eos_token_id": [
+        59246,
+        59253,
+        59255
+      ],
+      "head_dim": 128,
+      "hidden_act": "silu",
+      "hidden_size": 2048,
+      "id2label": {
+        "0": "LABEL_0",
+        "1": "LABEL_1"
+      },
+      "initializer_range": 0.02,
+      "intermediate_size": 6144,
+      "is_encoder_decoder": false,
+      "label2id": {
+        "LABEL_0": 0,
+        "LABEL_1": 1
+      },
+      "max_position_embeddings": 8192,
+      "mlp_bias": false,
+      "model_type": "llama",
+      "num_attention_heads": 16,
+      "num_hidden_layers": 28,
+      "num_key_value_heads": 4,
+      "output_attentions": false,
+      "output_hidden_states": false,
+      "pad_token_id": null,
+      "pretraining_tp": 1,
+      "problem_type": null,
+      "return_dict": true,
+      "rms_norm_eps": 1e-05,
+      "rope_parameters": {
+        "rope_theta": 10000.0,
+        "rope_type": "default"
+      },
+      "tie_word_embeddings": false,
+      "use_cache": true,
+      "vocab_size": 59264
+    },
+    "vocab_size": 59264
+  },
+  "encoder_conv_layers": [
+    [
+      1,
+      3,
+      1
+    ],
+    [
+      1,
+      3,
+      2
+    ]
+  ],
+  "encoder_dim": 1280,
+  "freeze_audio_head": false,
+  "freeze_projector": false,
+  "freq_mask_length": 27,
+  "label_smoothing": 0.0,
+  "length_penalty": 1.0,
+  "llm_dim": 2048,
+  "lora_alpha": 32,
+  "lora_dropout": 0.0,
+  "lora_rank": 8,
+  "lora_target_modules": [
+    "q_proj",
+    "k_proj",
+    "v_proj",
+    "o_proj",
+    "gate_proj",
+    "up_proj",
+    "down_proj"
+  ],
+  "max_audio_tokens": 500,
+  "max_new_tokens": 128,
+  "min_new_tokens": 0,
+  "model_dtype": "bfloat16",
+  "model_type": "asr_model",
+  "neucodec_model_id": "neuphonic/neucodec",
+  "no_repeat_ngram_size": 0,
+  "num_beams": 1,
+  "num_experts": 4,
+  "num_experts_per_tok": 2,
+  "num_freq_masks": 2,
+  "num_time_masks": 2,
+  "pipeline_tag": "automatic-speech-recognition",
+  "pretrained_model_path": "mazesmazes/tiny-audio-s2s-full",
+  "projector_dropout": 0.0,
+  "projector_hidden_dim": 1024,
+  "projector_init_std": 0.02,
+  "projector_num_layers": 2,
+  "projector_pool_stride": 4,
+  "projector_type": "mlp",
+  "qformer_hidden_size": null,
+  "qformer_intermediate_size": null,
+  "qformer_num_heads": 16,
+  "qformer_num_layers": 2,
+  "qformer_window_size": 15,
+  "repetition_penalty": 1.1,
+  "router_aux_loss_coef": 0.01,
+  "system_prompt": "",
+  "temperature": 1.0,
+  "text_config": {
+    "_name_or_path": "HuggingFaceTB/SmolLM3-3B",
+    "architectures": [
+      "SmolLM3ForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": null,
+    "dtype": "bfloat16",
+    "eos_token_id": 128012,
+    "hidden_act": "silu",
+    "hidden_size": 2048,
+    "initializer_range": 0.02,
+    "intermediate_size": 11008,
+    "layer_types": [
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 65536,
+    "max_window_layers": 28,
+    "mlp_bias": false,
+    "model_type": "smollm3",
+    "no_rope_layer_interval": 4,
+    "no_rope_layers": [
+      1,
+      1,
+      1,
+      0,
+      1,
+      1,
+      1,
+      0,
+      1,
+      1,
+      1,
+      0,
+      1,
+      1,
+      1,
+      0,
+      1,
+      1,
+      1,
+      0,
+      1,
+      1,
+      1,
+      0,
+      1,
+      1,
+      1,
+      0,
+      1,
+      1,
+      1,
+      0,
+      1,
+      1,
+      1,
+      0
+    ],
+    "num_attention_heads": 16,
+    "num_hidden_layers": 36,
+    "num_key_value_heads": 4,
+    "pad_token_id": 128004,
+    "pretraining_tp": 2,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "rope_theta": 5000000.0,
+      "rope_type": "default"
+    },
+    "sliding_window": null,
+    "tie_word_embeddings": true,
+    "use_cache": false,
+    "use_sliding_window": false,
+    "vocab_size": 128320
+  },
+  "text_model_id": "HuggingFaceTB/SmolLM3-3B",
+  "text_vocab_size": 128257,
+  "time_mask_length": 100,
+  "top_k": 50,
+  "top_p": null,
+  "transformers_version": "5.0.0",
+  "use_audio_head": true,
+  "use_cache": false,
+  "use_lora": false,
+  "use_specaugment": true,
+  "vocab_size": 128320
+}

diarization.py

@@ -0,0 +1,706 @@
+"""Speaker diarization using TEN-VAD + ECAPA-TDNN + spectral clustering.
+
+Spectral clustering implementation adapted from FunASR/3D-Speaker:
+https://github.com/alibaba-damo-academy/FunASR
+MIT License (https://opensource.org/licenses/MIT)
+"""
+
+import warnings
+
+import numpy as np
+import scipy
+import sklearn.metrics.pairwise
+import torch
+from sklearn.cluster._kmeans import k_means
+from sklearn.preprocessing import normalize
+
+
+def _get_device() -> torch.device:
+    """Get best available device for inference."""
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+class SpectralCluster:
+    """Spectral clustering using unnormalized Laplacian of affinity matrix.
+
+    Adapted from FunASR/3D-Speaker and SpeechBrain implementations.
+    Uses eigenvalue gap to automatically determine number of speakers.
+    """
+
+    def __init__(self, min_num_spks: int = 1, max_num_spks: int = 15, pval: float = 0.06):
+        self.min_num_spks = min_num_spks
+        self.max_num_spks = max_num_spks
+        self.pval = pval
+
+    def __call__(self, embeddings: np.ndarray, oracle_num: int | None = None) -> np.ndarray:
+        """Run spectral clustering on embeddings.
+
+        Args:
+            embeddings: Speaker embeddings of shape [N, D]
+            oracle_num: Optional known number of speakers
+
+        Returns:
+            Cluster labels of shape [N]
+        """
+        # Similarity matrix computation
+        sim_mat = self.get_sim_mat(embeddings)
+
+        # Refining similarity matrix with pval
+        prunned_sim_mat = self.p_pruning(sim_mat)
+
+        # Symmetrization
+        sym_prund_sim_mat = 0.5 * (prunned_sim_mat + prunned_sim_mat.T)
+
+        # Laplacian calculation
+        laplacian = self.get_laplacian(sym_prund_sim_mat)
+
+        # Get Spectral Embeddings
+        emb, num_of_spk = self.get_spec_embs(laplacian, oracle_num)
+
+        # Perform clustering
+        return self.cluster_embs(emb, num_of_spk)
+
+    def get_sim_mat(self, embeddings: np.ndarray) -> np.ndarray:
+        """Compute cosine similarity matrix."""
+        return sklearn.metrics.pairwise.cosine_similarity(embeddings, embeddings)
+
+    def p_pruning(self, affinity: np.ndarray) -> np.ndarray:
+        """Prune low similarity values in affinity matrix (keep top pval fraction)."""
+        n = affinity.shape[0]
+        pval = max(self.pval, 6.0 / n)
+        k_keep = max(1, int(pval * n))
+
+        # Vectorized: find top-k indices per row and zero out the rest
+        top_k_idx = np.argpartition(affinity, -k_keep, axis=1)[:, -k_keep:]
+        mask = np.zeros_like(affinity, dtype=bool)
+        np.put_along_axis(mask, top_k_idx, True, axis=1)
+        affinity[~mask] = 0
+        return affinity
+
+    def get_laplacian(self, sim_mat: np.ndarray) -> np.ndarray:
+        """Compute unnormalized Laplacian matrix."""
+        from scipy.sparse.csgraph import laplacian
+
+        np.fill_diagonal(sim_mat, 0)
+        return laplacian(sim_mat, normed=False)
+
+    def get_spec_embs(
+        self, laplacian: np.ndarray, k_oracle: int | None = None
+    ) -> tuple[np.ndarray, int]:
+        """Extract spectral embeddings from Laplacian.
+
+        Uses the eigengap heuristic to estimate the number of clusters:
+        The number of clusters k is chosen where the gap between consecutive
+        eigenvalues is largest, indicating a transition from "cluster" eigenvalues
+        (near 0) to "noise" eigenvalues.
+        """
+        lambdas, eig_vecs = scipy.linalg.eigh(laplacian)
+
+        num_of_spk = k_oracle if k_oracle is not None else self._estimate_num_speakers(lambdas)
+
+        emb = eig_vecs[:, :num_of_spk]
+        return emb, num_of_spk
+
+    def _estimate_num_speakers(self, lambdas: np.ndarray) -> int:
+        """Estimate number of speakers using refined eigengap heuristic.
+
+        For spectral clustering, we look for the largest gap in eigenvalues.
+        The eigenvalues corresponding to clusters are close to 0, and there
+        should be a significant jump to the remaining eigenvalues.
+        """
+        # Consider eigenvalues from index 1 to max_num_spks (skip first, it's always ~0)
+        # We need gaps between positions, so look at indices 1 to max_num_spks+1
+        max_idx = min(self.max_num_spks + 1, len(lambdas))
+        relevant_lambdas = lambdas[1:max_idx]  # Skip first eigenvalue
+
+        if len(relevant_lambdas) < 2:
+            return self.min_num_spks
+
+        # Compute absolute gaps (not ratios - ratios are unstable near 0)
+        gaps = np.diff(relevant_lambdas)
+
+        # Find the largest gap - the index gives us (k-1) since we skipped first
+        # Add 1 to convert from gap index to number of speakers
+        # Add 1 again because we skipped the first eigenvalue
+        max_gap_idx = int(np.argmax(gaps))
+        num_of_spk = max_gap_idx + 2  # +1 for gap->count, +1 for skipped eigenvalue
+
+        # Clamp between min and max
+        return max(self.min_num_spks, min(num_of_spk, self.max_num_spks))
+
+    def cluster_embs(self, emb: np.ndarray, k: int) -> np.ndarray:
+        """Cluster spectral embeddings using k-means."""
+        _, labels, _ = k_means(emb, k, n_init=10)
+        return labels
+
+
+class SpeakerClusterer:
+    """Speaker clustering backend using spectral clustering with speaker merging.
+
+    Features:
+    - Spectral clustering with eigenvalue gap for auto speaker count detection
+    - P-pruning for affinity matrix refinement
+    - Post-clustering speaker merging by cosine similarity
+    """
+
+    def __init__(
+        self,
+        min_num_spks: int = 2,
+        max_num_spks: int = 10,
+        merge_thr: float = 0.90,  # Moderate merging
+    ):
+        self.min_num_spks = min_num_spks
+        self.max_num_spks = max_num_spks
+        self.merge_thr = merge_thr
+        self._spectral_cluster: SpectralCluster | None = None
+
+    def _get_spectral_cluster(self) -> SpectralCluster:
+        """Lazy-load spectral clusterer."""
+        if self._spectral_cluster is None:
+            self._spectral_cluster = SpectralCluster(
+                min_num_spks=self.min_num_spks,
+                max_num_spks=self.max_num_spks,
+            )
+        return self._spectral_cluster
+
+    def __call__(self, embeddings: np.ndarray, num_speakers: int | None = None) -> np.ndarray:
+        """Cluster speaker embeddings and return labels.
+
+        Args:
+            embeddings: Speaker embeddings of shape [N, D]
+            num_speakers: Optional oracle number of speakers
+
+        Returns:
+            Cluster labels of shape [N]
+        """
+        import warnings
+
+        if len(embeddings.shape) != 2:
+            raise ValueError(f"Expected 2D array, got shape {embeddings.shape}")
+
+        # Handle edge cases
+        if embeddings.shape[0] == 0:
+            return np.array([], dtype=int)
+        if embeddings.shape[0] == 1:
+            return np.array([0], dtype=int)
+        if embeddings.shape[0] < 6:
+            return np.zeros(embeddings.shape[0], dtype=int)
+
+        # Normalize embeddings and replace NaN/inf
+        embeddings = np.nan_to_num(embeddings, nan=0.0, posinf=0.0, neginf=0.0)
+        embeddings = normalize(embeddings)
+
+        # Run spectral clustering (suppress numerical warnings)
+        spectral = self._get_spectral_cluster()
+
+        # Update min/max for oracle case
+        if num_speakers is not None:
+            spectral.min_num_spks = num_speakers
+            spectral.max_num_spks = num_speakers
+
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=RuntimeWarning)
+            labels = spectral(embeddings, oracle_num=num_speakers)
+
+        # Reset min/max
+        if num_speakers is not None:
+            spectral.min_num_spks = self.min_num_spks
+            spectral.max_num_spks = self.max_num_spks
+
+        # Merge similar speakers if no oracle
+        if num_speakers is None:
+            labels = self._merge_by_cos(labels, embeddings, self.merge_thr)
+
+        # Re-index labels sequentially
+        _, labels = np.unique(labels, return_inverse=True)
+
+        return labels
+
+    def _merge_by_cos(self, labels: np.ndarray, embs: np.ndarray, cos_thr: float) -> np.ndarray:
+        """Merge similar speakers by cosine similarity of centroids."""
+        from scipy.cluster.hierarchy import fcluster, linkage
+        from scipy.spatial.distance import pdist
+
+        unique_labels = np.unique(labels)
+        if len(unique_labels) <= 1:
+            return labels
+
+        # Compute normalized speaker centroids
+        centroids = np.array([embs[labels == lbl].mean(0) for lbl in unique_labels])
+        centroids = normalize(centroids)
+
+        # Hierarchical clustering with cosine distance
+        distances = pdist(centroids, metric="cosine")
+        linkage_matrix = linkage(distances, method="average")
+        merged_labels = fcluster(linkage_matrix, t=1.0 - cos_thr, criterion="distance") - 1
+
+        # Map original labels to merged labels
+        label_map = dict(zip(unique_labels, merged_labels))
+        return np.array([label_map[lbl] for lbl in labels])
+
+
+class LocalSpeakerDiarizer:
+    """Local speaker diarization using TEN-VAD + ECAPA-TDNN + spectral clustering.
+
+    Pipeline:
+    1. TEN-VAD detects speech segments
+    2. Sliding window (1.0s, 75% overlap) for uniform embedding extraction
+    3. ECAPA-TDNN extracts speaker embeddings per window
+    4. Spectral clustering with eigenvalue gap for auto speaker detection
+    5. Frame-level consensus voting for segment reconstruction
+    6. Post-processing merges short segments to reduce flicker
+
+    Tunable Parameters (class attributes):
+    - WINDOW_SIZE: Embedding extraction window size in seconds
+    - STEP_SIZE: Sliding window step size (overlap = WINDOW_SIZE - STEP_SIZE)
+    - VAD_THRESHOLD: Speech detection threshold (lower = more sensitive)
+    - VAD_MIN_DURATION: Minimum speech segment duration
+    - VAD_MAX_GAP: Maximum gap to bridge between segments
+    - VAD_PAD_ONSET/OFFSET: Padding added to speech segments
+    - VOTING_RATE: Frame resolution for consensus voting
+    - MIN_SEGMENT_DURATION: Minimum final segment duration
+    - SAME_SPEAKER_GAP: Maximum gap to merge same-speaker segments
+    - TAIL_COVERAGE_RATIO: Minimum tail coverage to add extra window
+    """
+
+    _ten_vad_model = None
+    _ecapa_model = None
+    _device = None
+
+    # ==================== TUNABLE PARAMETERS ====================
+
+    # Sliding window for embedding extraction
+    WINDOW_SIZE = 0.75  # seconds - shorter window for finer resolution
+    STEP_SIZE = 0.15  # seconds (80% overlap for more votes)
+    TAIL_COVERAGE_RATIO = 0.1  # Add extra window if tail > this ratio of window
+
+    # VAD hysteresis parameters
+    VAD_THRESHOLD = 0.25  # Balanced threshold
+    VAD_MIN_DURATION = 0.05  # Minimum speech segment duration (seconds)
+    VAD_MAX_GAP = 0.50  # Bridge gaps shorter than this (seconds)
+    VAD_PAD_ONSET = 0.05  # Padding at segment start (seconds)
+    VAD_PAD_OFFSET = 0.05  # Padding at segment end (seconds)
+
+    # Frame-level voting
+    VOTING_RATE = 0.01  # 10ms resolution for consensus voting
+
+    # Post-processing
+    MIN_SEGMENT_DURATION = 0.15  # Minimum final segment duration (seconds)
+    SHORT_SEGMENT_GAP = 0.1  # Gap threshold for merging short segments
+    SAME_SPEAKER_GAP = 0.5  # Gap threshold for merging same-speaker segments
+
+    # ===========================================================
+
+    @classmethod
+    def _get_ten_vad_model(cls):
+        """Lazy-load TEN-VAD model (singleton)."""
+        if cls._ten_vad_model is None:
+            from ten_vad import TenVad
+
+            cls._ten_vad_model = TenVad(hop_size=256, threshold=cls.VAD_THRESHOLD)
+        return cls._ten_vad_model
+
+    @classmethod
+    def _get_device(cls) -> torch.device:
+        """Get the best available device."""
+        if cls._device is None:
+            cls._device = _get_device()
+        return cls._device
+
+    @classmethod
+    def _get_ecapa_model(cls):
+        """Lazy-load ECAPA-TDNN speaker embedding model (singleton)."""
+        if cls._ecapa_model is None:
+            # Suppress torchaudio deprecation warning from SpeechBrain
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", message="torchaudio._backend")
+                from speechbrain.inference.speaker import EncoderClassifier
+
+                device = cls._get_device()
+                cls._ecapa_model = EncoderClassifier.from_hparams(
+                    source="speechbrain/spkrec-ecapa-voxceleb",
+                    run_opts={"device": str(device)},
+                )
+
+        return cls._ecapa_model
+
+    @classmethod
+    def diarize(
+        cls,
+        audio: np.ndarray | str,
+        sample_rate: int = 16000,
+        num_speakers: int | None = None,
+        min_speakers: int = 2,
+        max_speakers: int = 10,
+        **_kwargs,
+    ) -> list[dict]:
+        """Run speaker diarization on audio.
+
+        Args:
+            audio: Audio waveform as numpy array or path to audio file
+            sample_rate: Audio sample rate (default 16000)
+            num_speakers: Exact number of speakers (if known)
+            min_speakers: Minimum number of speakers
+            max_speakers: Maximum number of speakers
+
+        Returns:
+            List of dicts with 'speaker', 'start', 'end' keys
+        """
+        # Handle file path input
+        if isinstance(audio, str):
+            import librosa
+
+            audio, sample_rate = librosa.load(audio, sr=16000)
+
+        # Ensure correct sample rate
+        if sample_rate != 16000:
+            import librosa
+
+            audio = librosa.resample(audio, orig_sr=sample_rate, target_sr=16000)
+            sample_rate = 16000
+
+        audio = audio.astype(np.float32)
+        total_duration = len(audio) / sample_rate
+
+        # Step 1: VAD (returns segments and raw frame-level decisions)
+        segments, vad_frames = cls._get_speech_segments(audio, sample_rate)
+        if not segments:
+            return []
+
+        # Step 2: Extract embeddings
+        embeddings, window_segments = cls._extract_embeddings(audio, segments, sample_rate)
+        if len(embeddings) == 0:
+            return []
+
+        # Step 3: Cluster
+        clusterer = SpeakerClusterer(min_num_spks=min_speakers, max_num_spks=max_speakers)
+        labels = clusterer(embeddings, num_speakers)
+
+        # Step 4: Post-process with consensus voting (VAD-aware)
+        return cls._postprocess_segments(window_segments, labels, total_duration, vad_frames)
+
+    @classmethod
+    def _get_speech_segments(
+        cls, audio_array: np.ndarray, sample_rate: int = 16000
+    ) -> tuple[list[dict], list[bool]]:
+        """Get speech segments using TEN-VAD.
+
+        Returns:
+            Tuple of (segments list, vad_frames list of per-frame speech decisions)
+        """
+        vad_model = cls._get_ten_vad_model()
+
+        # Convert to int16 as required by TEN-VAD
+        # Clip to prevent integer overflow
+        if audio_array.dtype != np.int16:
+            audio_int16 = (np.clip(audio_array, -1.0, 1.0) * 32767).astype(np.int16)
+        else:
+            audio_int16 = audio_array
+
+        # Process frame by frame
+        hop_size = 256
+        frame_duration = hop_size / sample_rate
+        speech_frames: list[bool] = []
+
+        for i in range(0, len(audio_int16) - hop_size, hop_size):
+            frame = audio_int16[i : i + hop_size]
+            _, is_speech = vad_model.process(frame)
+            speech_frames.append(is_speech)
+
+        # Convert frame-level decisions to segments
+        segments = []
+        in_speech = False
+        start_idx = 0
+
+        for i, is_speech in enumerate(speech_frames):
+            if is_speech and not in_speech:
+                start_idx = i
+                in_speech = True
+            elif not is_speech and in_speech:
+                start_time = start_idx * frame_duration
+                end_time = i * frame_duration
+                segments.append(
+                    {
+                        "start": start_time,
+                        "end": end_time,
+                        "start_sample": int(start_time * sample_rate),
+                        "end_sample": int(end_time * sample_rate),
+                    }
+                )
+                in_speech = False
+
+        # Handle trailing speech
+        if in_speech:
+            start_time = start_idx * frame_duration
+            end_time = len(speech_frames) * frame_duration
+            segments.append(
+                {
+                    "start": start_time,
+                    "end": end_time,
+                    "start_sample": int(start_time * sample_rate),
+                    "end_sample": int(end_time * sample_rate),
+                }
+            )
+
+        return cls._apply_vad_hysteresis(segments, sample_rate), speech_frames
+
+    @classmethod
+    def _apply_vad_hysteresis(cls, segments: list[dict], sample_rate: int = 16000) -> list[dict]:
+        """Apply hysteresis-like post-processing to VAD segments."""
+        if not segments:
+            return segments
+
+        segments = sorted(segments, key=lambda x: x["start"])
+
+        # Fill short gaps
+        merged = [segments[0].copy()]
+        for seg in segments[1:]:
+            gap = seg["start"] - merged[-1]["end"]
+            if gap <= cls.VAD_MAX_GAP:
+                merged[-1]["end"] = seg["end"]
+                merged[-1]["end_sample"] = seg["end_sample"]
+            else:
+                merged.append(seg.copy())
+
+        # Remove short segments
+        filtered = [seg for seg in merged if (seg["end"] - seg["start"]) >= cls.VAD_MIN_DURATION]
+
+        # Dilate segments (add padding)
+        for seg in filtered:
+            seg["start"] = max(0.0, seg["start"] - cls.VAD_PAD_ONSET)
+            seg["end"] = seg["end"] + cls.VAD_PAD_OFFSET
+            seg["start_sample"] = int(seg["start"] * sample_rate)
+            seg["end_sample"] = int(seg["end"] * sample_rate)
+
+        return filtered
+
+    @classmethod
+    def _extract_embeddings(
+        cls, audio_array: np.ndarray, segments: list[dict], sample_rate: int
+    ) -> tuple[np.ndarray, list[dict]]:
+        """Extract speaker embeddings using sliding windows."""
+        speaker_model = cls._get_ecapa_model()
+
+        window_samples = int(cls.WINDOW_SIZE * sample_rate)
+        step_samples = int(cls.STEP_SIZE * sample_rate)
+
+        embeddings = []
+        window_segments = []
+
+        with torch.no_grad():
+            for seg in segments:
+                seg_start = seg["start_sample"]
+                seg_end = seg["end_sample"]
+                seg_len = seg_end - seg_start
+
+                # Generate window positions
+                if seg_len <= window_samples:
+                    starts = [seg_start]
+                    ends = [seg_end]
+                else:
+                    starts = list(range(seg_start, seg_end - window_samples + 1, step_samples))
+                    ends = [s + window_samples for s in starts]
+
+                    # Cover tail if > TAIL_COVERAGE_RATIO of window remains
+                    if ends and ends[-1] < seg_end:
+                        remainder = seg_end - ends[-1]
+                        if remainder > (window_samples * cls.TAIL_COVERAGE_RATIO):
+                            starts.append(seg_end - window_samples)
+                            ends.append(seg_end)
+
+                for c_start, c_end in zip(starts, ends):
+                    chunk = audio_array[c_start:c_end]
+
+                    # Pad short chunks with reflection
+                    if len(chunk) < window_samples:
+                        pad_width = window_samples - len(chunk)
+                        chunk = np.pad(chunk, (0, pad_width), mode="reflect")
+
+                    # Extract embedding using SpeechBrain's encode_batch
+                    chunk_tensor = torch.from_numpy(chunk).float().unsqueeze(0)
+                    embedding = (
+                        speaker_model.encode_batch(chunk_tensor).squeeze(0).squeeze(0).cpu().numpy()
+                    )
+
+                    # Validate embedding
+                    if np.isfinite(embedding).all() and np.linalg.norm(embedding) > 1e-8:
+                        embeddings.append(embedding)
+                        window_segments.append(
+                            {
+                                "start": c_start / sample_rate,
+                                "end": c_end / sample_rate,
+                            }
+                        )
+
+        # Normalize all embeddings at once
+        if embeddings:
+            return normalize(np.array(embeddings)), window_segments
+        return np.array([]), []
+
+    @classmethod
+    def _resample_vad(cls, vad_frames: list[bool], num_frames: int) -> np.ndarray:
+        """Resample VAD frame decisions to match voting grid resolution.
+
+        VAD operates at 256 samples / 16000 Hz = 16ms per frame.
+        Voting operates at VOTING_RATE (default 10ms) per frame.
+        This maps VAD decisions to the finer voting grid.
+        """
+        if not vad_frames:
+            return np.zeros(num_frames, dtype=bool)
+
+        vad_rate = 256 / 16000  # 16ms per VAD frame
+        vad_arr = np.array(vad_frames)
+
+        # Vectorized: compute VAD frame indices for each voting frame
+        voting_times = np.arange(num_frames) * cls.VOTING_RATE
+        vad_indices = np.clip((voting_times / vad_rate).astype(int), 0, len(vad_arr) - 1)
+        return vad_arr[vad_indices]
+
+    @classmethod
+    def _postprocess_segments(
+        cls,
+        window_segments: list[dict],
+        labels: np.ndarray,
+        total_duration: float,
+        vad_frames: list[bool],
+    ) -> list[dict]:
+        """Post-process using frame-level consensus voting with VAD-aware silence."""
+        if not window_segments or len(labels) == 0:
+            return []
+
+        # Correct labels to be contiguous
+        unique_labels = np.unique(labels)
+        label_map = {old: new for new, old in enumerate(unique_labels)}
+        clean_labels = np.array([label_map[lbl] for lbl in labels])
+        num_speakers = len(unique_labels)
+
+        if num_speakers == 0:
+            return []
+
+        # Create voting grid
+        num_frames = int(np.ceil(total_duration / cls.VOTING_RATE)) + 1
+        votes = np.zeros((num_frames, num_speakers), dtype=np.float32)
+
+        # Accumulate votes
+        for win, label in zip(window_segments, clean_labels):
+            start_frame = int(win["start"] / cls.VOTING_RATE)
+            end_frame = int(win["end"] / cls.VOTING_RATE)
+            end_frame = min(end_frame, num_frames)
+            if start_frame < end_frame:
+                votes[start_frame:end_frame, label] += 1.0
+
+        # Determine winner per frame
+        frame_speakers = np.argmax(votes, axis=1)
+        max_votes = np.max(votes, axis=1)
+
+        # Resample VAD to voting grid resolution for silence-aware voting
+        vad_resampled = cls._resample_vad(vad_frames, num_frames)
+
+        # Convert frames to segments
+        final_segments = []
+        current_speaker = -1
+        seg_start = 0.0
+
+        for f in range(num_frames):
+            speaker = int(frame_speakers[f])
+            score = max_votes[f]
+
+            # Force silence if VAD says no speech OR no votes
+            if score == 0 or not vad_resampled[f]:
+                speaker = -1
+
+            if speaker != current_speaker:
+                if current_speaker != -1:
+                    final_segments.append(
+                        {
+                            "speaker": f"SPEAKER_{current_speaker}",
+                            "start": seg_start,
+                            "end": f * cls.VOTING_RATE,
+                        }
+                    )
+                current_speaker = speaker
+                seg_start = f * cls.VOTING_RATE
+
+        # Close last segment
+        if current_speaker != -1:
+            final_segments.append(
+                {
+                    "speaker": f"SPEAKER_{current_speaker}",
+                    "start": seg_start,
+                    "end": num_frames * cls.VOTING_RATE,
+                }
+            )
+
+        return cls._merge_short_segments(final_segments)
+
+    @classmethod
+    def _merge_short_segments(cls, segments: list[dict]) -> list[dict]:
+        """Merge short segments to reduce flicker."""
+        if not segments:
+            return []
+
+        clean: list[dict] = []
+        for seg in segments:
+            dur = seg["end"] - seg["start"]
+            if dur < cls.MIN_SEGMENT_DURATION:
+                if (
+                    clean
+                    and clean[-1]["speaker"] == seg["speaker"]
+                    and seg["start"] - clean[-1]["end"] < cls.SHORT_SEGMENT_GAP
+                ):
+                    clean[-1]["end"] = seg["end"]
+                continue
+
+            if (
+                clean
+                and clean[-1]["speaker"] == seg["speaker"]
+                and seg["start"] - clean[-1]["end"] < cls.SAME_SPEAKER_GAP
+            ):
+                clean[-1]["end"] = seg["end"]
+            else:
+                clean.append(seg)
+
+        return clean
+
+    @classmethod
+    def assign_speakers_to_words(
+        cls,
+        words: list[dict],
+        speaker_segments: list[dict],
+    ) -> list[dict]:
+        """Assign speaker labels to words based on timestamp overlap.
+
+        Args:
+            words: List of word dicts with 'word', 'start', 'end' keys
+            speaker_segments: List of speaker dicts with 'speaker', 'start', 'end' keys
+
+        Returns:
+            Words list with 'speaker' key added to each word
+        """
+        for word in words:
+            word_mid = (word["start"] + word["end"]) / 2
+
+            # Find the speaker segment that contains this word's midpoint
+            best_speaker = None
+            for seg in speaker_segments:
+                if seg["start"] <= word_mid <= seg["end"]:
+                    best_speaker = seg["speaker"]
+                    break
+
+            # If no exact match, find closest segment
+            if best_speaker is None and speaker_segments:
+                min_dist = float("inf")
+                for seg in speaker_segments:
+                    seg_mid = (seg["start"] + seg["end"]) / 2
+                    dist = abs(word_mid - seg_mid)
+                    if dist < min_dist:
+                        min_dist = dist
+                        best_speaker = seg["speaker"]
+
+            word["speaker"] = best_speaker
+
+        return words

full_duplex.py

@@ -0,0 +1,475 @@
+"""Full-duplex audio session for speech-to-speech.
+
+Implements Freeze-Omni style full-duplex conversation where the model
+can listen and speak simultaneously, with support for user interruption.
+
+Architecture:
+- Dual queue system: PCMQueue (input) + AudioQueue (output)
+- Multi-threaded: Listen thread + Generate thread run concurrently
+- State machine: listen -> speak -> (interrupt) -> listen
+- VAD-based turn detection using model's built-in Silero VAD
+
+Usage (sync):
+    session = FullDuplexSession(model)
+    session.start()
+
+    while has_audio:
+        session.push_audio(audio_chunk)
+        output = session.pop_audio()
+        if output is not None:
+            speaker.play(output)
+
+    session.stop()
+
+Usage (async/web):
+    session = FullDuplexSession(
+        model,
+        on_state_change=lambda s: send_status(s),
+        on_text=lambda t: send_text(t),
+        on_audio=lambda a: send_audio(a),
+    )
+    session.start()
+
+    # In your receive loop:
+    session.push_audio(audio_chunk)
+"""
+
+import logging
+import queue
+import threading
+import time
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import TYPE_CHECKING, Callable, Optional
+
+import numpy as np
+import torch
+
+if TYPE_CHECKING:
+    from .asr_modeling import ASRModel
+
+logger = logging.getLogger(__name__)
+
+
+class ConversationState(Enum):
+    """State machine for full-duplex conversation."""
+
+    IDLE = "idle"
+    LISTENING = "listening"
+    PROCESSING = "processing"
+    SPEAKING = "speaking"
+
+
+@dataclass
+class FullDuplexConfig:
+    """Configuration for full-duplex session."""
+
+    # Audio settings
+    sample_rate: int = 16000
+    chunk_size: int = 512  # Samples per chunk (32ms at 16kHz)
+    output_sample_rate: int = 44100  # DAC output rate
+
+    # VAD settings
+    vad_threshold: float = 0.5
+    silence_duration_ms: float = 700  # Silence to end turn
+    min_speech_duration_ms: float = 100  # Minimum speech to trigger
+
+    # Generation settings
+    audio_chunk_size: int = 4  # Tokens per audio chunk
+
+    # Timing
+    poll_interval: float = 0.01
+
+
+class PCMQueue:
+    """Thread-safe queue for streaming PCM audio input."""
+
+    def __init__(self):
+        self.buffer = np.array([], dtype=np.float32)
+        self.lock = threading.Lock()
+
+    def put(self, audio: np.ndarray) -> None:
+        with self.lock:
+            self.buffer = np.concatenate([self.buffer, audio.astype(np.float32)])
+
+    def get(self, length: int) -> Optional[np.ndarray]:
+        with self.lock:
+            if len(self.buffer) < length:
+                return None
+            result = self.buffer[:length]
+            self.buffer = self.buffer[length:]
+            return result
+
+    def clear(self) -> None:
+        with self.lock:
+            self.buffer = np.array([], dtype=np.float32)
+
+    def __len__(self) -> int:
+        with self.lock:
+            return len(self.buffer)
+
+
+class AudioQueue:
+    """Thread-safe queue for output audio chunks."""
+
+    def __init__(self):
+        self._queue: queue.Queue = queue.Queue()
+
+    def put(self, audio: torch.Tensor) -> None:
+        self._queue.put(audio)
+
+    def get(self) -> Optional[torch.Tensor]:
+        try:
+            return self._queue.get_nowait()
+        except queue.Empty:
+            return None
+
+    def clear(self) -> None:
+        while not self._queue.empty():
+            try:
+                self._queue.get_nowait()
+            except queue.Empty:
+                break
+
+    def is_empty(self) -> bool:
+        return self._queue.empty()
+
+
+@dataclass
+class _SessionState:
+    """Internal state for full-duplex session."""
+
+    state: ConversationState = ConversationState.IDLE
+    speech_buffer: list = field(default_factory=list)
+    speech_start_time: float = 0.0
+    last_speech_time: float = 0.0
+    silence_frames: int = 0
+    stop_generate: bool = False
+    is_generating: bool = False
+    generated_text: str = ""
+
+
+class FullDuplexSession:
+    """Full-duplex speech-to-speech session (Freeze-Omni style).
+
+    Manages simultaneous listening and speaking with VAD-based turn detection.
+    Designed to be easy to integrate with both sync and async (web) applications.
+
+    Args:
+        model: ASRModel with audio_head configured
+        config: FullDuplexConfig for session parameters
+        on_state_change: Callback when state changes (state: ConversationState)
+        on_text: Callback when text is generated (text: str, interim: bool)
+        on_audio: Callback when audio chunk is ready (audio: torch.Tensor)
+            If provided, audio is sent here instead of output_queue
+        on_interrupted: Callback when generation is interrupted
+    """
+
+    def __init__(
+        self,
+        model: "ASRModel",
+        config: Optional[FullDuplexConfig] = None,
+        on_state_change: Optional[Callable[[ConversationState], None]] = None,
+        on_text: Optional[Callable[[str, bool], None]] = None,
+        on_audio: Optional[Callable[[torch.Tensor], None]] = None,
+        on_interrupted: Optional[Callable[[], None]] = None,
+    ):
+        self.model = model
+        self.config = config or FullDuplexConfig()
+
+        # Callbacks
+        self.on_state_change = on_state_change
+        self.on_text = on_text
+        self.on_audio = on_audio
+        self.on_interrupted = on_interrupted
+
+        # Queues
+        self.input_queue = PCMQueue()
+        self.output_queue = AudioQueue()
+
+        # State
+        self._state = _SessionState()
+        self._running = False
+        self._state_lock = threading.Lock()
+
+        # Threads
+        self._listen_thread: Optional[threading.Thread] = None
+        self._generate_thread: Optional[threading.Thread] = None
+
+        # Precompute timing thresholds
+        ms_per_chunk = self.config.chunk_size / self.config.sample_rate * 1000
+        self._silence_threshold = int(self.config.silence_duration_ms / ms_per_chunk)
+        self._min_speech_chunks = int(self.config.min_speech_duration_ms / ms_per_chunk)
+
+        # Ensure VAD is loaded
+        self.model.load_vad()
+
+    @property
+    def state(self) -> ConversationState:
+        with self._state_lock:
+            return self._state.state
+
+    def _set_state(self, value: ConversationState) -> None:
+        with self._state_lock:
+            old_state = self._state.state
+            self._state.state = value
+        if old_state != value:
+            logger.debug(f"State: {old_state.value} -> {value.value}")
+            if self.on_state_change:
+                try:
+                    self.on_state_change(value)
+                except Exception as e:
+                    logger.error(f"on_state_change callback error: {e}")
+
+    @property
+    def is_generating(self) -> bool:
+        with self._state_lock:
+            return self._state.is_generating
+
+    @property
+    def generated_text(self) -> str:
+        with self._state_lock:
+            return self._state.generated_text
+
+    def start(self) -> None:
+        """Start the full-duplex session."""
+        if self._running:
+            return
+
+        self._running = True
+        self._set_state(ConversationState.LISTENING)
+
+        self._listen_thread = threading.Thread(target=self._listen_loop, daemon=True)
+        self._listen_thread.start()
+
+        logger.info("Full-duplex session started")
+
+    def stop(self) -> None:
+        """Stop the full-duplex session."""
+        self._running = False
+
+        with self._state_lock:
+            self._state.stop_generate = True
+
+        if self._listen_thread:
+            self._listen_thread.join(timeout=2.0)
+        if self._generate_thread:
+            self._generate_thread.join(timeout=2.0)
+
+        self.input_queue.clear()
+        self.output_queue.clear()
+        self._set_state(ConversationState.IDLE)
+
+        logger.info("Full-duplex session stopped")
+
+    def push_audio(self, audio: np.ndarray) -> None:
+        """Push audio samples to the input queue.
+
+        Args:
+            audio: Audio samples as numpy array (float32 normalized or int16)
+        """
+        if audio.dtype == np.int16:
+            audio = audio.astype(np.float32) / 32768.0
+        self.input_queue.put(audio)
+
+    def pop_audio(self) -> Optional[torch.Tensor]:
+        """Pop generated audio from the output queue.
+
+        Only used if on_audio callback is not set.
+
+        Returns:
+            Audio tensor [samples] or None
+        """
+        return self.output_queue.get()
+
+    def interrupt(self) -> None:
+        """Interrupt current generation and return to listening."""
+        with self._state_lock:
+            self._state.stop_generate = True
+
+        # Wait for generation to stop
+        timeout = 2.0
+        start = time.time()
+        while self._state.is_generating and (time.time() - start) < timeout:
+            time.sleep(self.config.poll_interval)
+
+        # Clear output queue
+        self.output_queue.clear()
+
+        # Reset state
+        with self._state_lock:
+            self._state.stop_generate = False
+            self._state.generated_text = ""
+            self._state.speech_buffer.clear()
+            self._state.silence_frames = 0
+
+        self._set_state(ConversationState.LISTENING)
+        self.model.reset_vad_state()
+
+        if self.on_interrupted:
+            try:
+                self.on_interrupted()
+            except Exception as e:
+                logger.error(f"on_interrupted callback error: {e}")
+
+        logger.debug("Generation interrupted")
+
+    def _emit_audio(self, audio: torch.Tensor) -> None:
+        """Send audio to callback or queue."""
+        if self.on_audio:
+            try:
+                self.on_audio(audio)
+            except Exception as e:
+                logger.error(f"on_audio callback error: {e}")
+        else:
+            self.output_queue.put(audio)
+
+    def _emit_text(self, text: str, interim: bool = False) -> None:
+        """Send text to callback."""
+        if self.on_text:
+            try:
+                self.on_text(text, interim)
+            except Exception as e:
+                logger.error(f"on_text callback error: {e}")
+
+    def _listen_loop(self) -> None:
+        """Main listening loop - processes audio and detects speech."""
+        is_speaking = False
+
+        while self._running:
+            audio = self.input_queue.get(self.config.chunk_size)
+            if audio is None:
+                time.sleep(self.config.poll_interval)
+                continue
+
+            # Run VAD
+            audio_tensor = torch.from_numpy(audio)
+            is_speech, prob = self.model.detect_speech(
+                audio_tensor,
+                self.config.sample_rate,
+                self.config.vad_threshold,
+            )
+
+            current_time = time.time()
+
+            # Check for interruption during generation
+            if self._state.is_generating and is_speech:
+                logger.debug(f"Interruption detected (prob={prob:.2f})")
+                self.interrupt()
+                # Start new utterance with this chunk
+                is_speaking = True
+                with self._state_lock:
+                    self._state.speech_buffer = [audio]
+                    self._state.speech_start_time = current_time
+                    self._state.last_speech_time = current_time
+                    self._state.silence_frames = 0
+                continue
+
+            # Normal VAD state machine
+            if is_speech:
+                if not is_speaking:
+                    is_speaking = True
+                    with self._state_lock:
+                        self._state.speech_buffer = []
+                        self._state.speech_start_time = current_time
+                with self._state_lock:
+                    self._state.speech_buffer.append(audio)
+                    self._state.last_speech_time = current_time
+                    self._state.silence_frames = 0
+
+            elif is_speaking:
+                with self._state_lock:
+                    self._state.speech_buffer.append(audio)
+                    self._state.silence_frames += 1
+
+                    if self._state.silence_frames >= self._silence_threshold:
+                        is_speaking = False
+
+                        # Check minimum speech duration
+                        if len(self._state.speech_buffer) >= self._min_speech_chunks:
+                            speech_audio = np.concatenate(self._state.speech_buffer)
+                            self._state.speech_buffer = []
+                            self._state.silence_frames = 0
+
+                            # Start generation
+                            self._generate_thread = threading.Thread(
+                                target=self._generate_loop,
+                                args=(speech_audio,),
+                                daemon=True,
+                            )
+                            self._generate_thread.start()
+                        else:
+                            self._state.speech_buffer = []
+                            self._state.silence_frames = 0
+
+    def _generate_loop(self, speech_audio: np.ndarray) -> None:
+        """Generation loop - produces text and audio response."""
+        with self._state_lock:
+            self._state.is_generating = True
+            self._state.generated_text = ""
+            self._state.stop_generate = False
+
+        try:
+            self._set_state(ConversationState.PROCESSING)
+
+            # Process input audio
+            device = next(self.model.language_model.parameters()).device
+            inputs = self.model._process_audio(speech_audio, self.config.sample_rate)
+            input_features = inputs["input_features"]
+            audio_attention_mask = inputs["attention_mask"]
+
+            # Encode
+            audio_embeds = self.model._encode_audio(input_features, audio_attention_mask)
+            input_ids, attention_mask = self.model._build_audio_prompt(
+                audio_attention_mask, 1, device
+            )
+            inputs_embeds = self.model._inject_audio_embeddings(input_ids, audio_embeds)
+
+            # Check for interruption
+            if self._state.stop_generate:
+                return
+
+            # Generate text
+            with torch.no_grad():
+                output = self.model.language_model.generate(
+                    input_ids=input_ids,
+                    inputs_embeds=inputs_embeds,
+                    attention_mask=attention_mask,
+                    generation_config=self.model.generation_config,
+                )
+
+            if self._state.stop_generate:
+                return
+
+            # Extract text
+            text_ids = output[:, input_ids.shape[1] :]
+            text = self.model.tokenizer.decode(text_ids[0], skip_special_tokens=True)
+
+            with self._state_lock:
+                self._state.generated_text = text
+
+            self._emit_text(text, interim=False)
+
+            if self._state.stop_generate:
+                return
+
+            # Generate audio
+            if self.model.audio_head is not None:
+                self._set_state(ConversationState.SPEAKING)
+
+                for audio_chunk in self.model.audio_head.generate_streaming(
+                    text_token_ids=text_ids,
+                ):
+                    if self._state.stop_generate:
+                        return
+                    self._emit_audio(audio_chunk)
+
+            self._set_state(ConversationState.LISTENING)
+
+        except Exception as e:
+            logger.error(f"Generation error: {e}")
+            self._set_state(ConversationState.LISTENING)
+
+        finally:
+            with self._state_lock:
+                self._state.is_generating = False

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4df14201e66792d6b8ccefd124852fdc5d47f013a6276f9973540d63956097ad
+size 122303840

preprocessor_config.json

@@ -0,0 +1,19 @@
+{
+  "chunk_length": 30,
+  "dither": 0.0,
+  "feature_extractor_type": "WhisperFeatureExtractor",
+  "feature_size": 128,
+  "hop_length": 160,
+  "n_fft": 400,
+  "n_samples": 480000,
+  "nb_max_frames": 3000,
+  "padding": false,
+  "padding_side": "right",
+  "padding_value": 0.0,
+  "return_attention_mask": false,
+  "sampling_rate": 16000,
+  "processor_class": "ASRProcessor",
+  "auto_map": {
+    "AutoProcessor": "asr_processing.ASRProcessor"
+  }
+}

projectors.py

@@ -0,0 +1,64 @@
+"""Audio projector module for bridging encoder and decoder embeddings.
+
+MLPAudioProjector: Simple 2-layer MLP with frame stacking downsampling.
+"""
+
+import torch
+import torch.nn as nn
+from transformers.models.llama.modeling_llama import LlamaRMSNorm
+
+
+class MLPAudioProjector(nn.Module):
+    """2-layer MLP projector with frame-stacking downsampling (matches GLM-ASR)."""
+
+    def __init__(self, config):
+        """Initialize MLP projector.
+
+        Args:
+            config: ASRConfig with encoder_dim, llm_dim, projector_pool_stride
+        """
+        super().__init__()
+
+        encoder_dim = getattr(config, "encoder_dim", 768)
+        llm_dim = getattr(config, "llm_dim", 2048)
+        self.k = getattr(config, "projector_pool_stride", 4)
+
+        # Frame stacking: concat k adjacent frames then project
+        in_dim = encoder_dim * self.k
+        # Hidden dim defaults to llm_dim, can be overridden via config
+        hidden_dim = getattr(config, "projector_hidden_dim", None) or llm_dim
+        self.linear_1 = nn.Linear(in_dim, hidden_dim, bias=False)
+        self.norm = LlamaRMSNorm(hidden_dim, eps=1e-6)
+        self.act = nn.GELU()
+        self.linear_2 = nn.Linear(hidden_dim, llm_dim, bias=False)
+
+    def get_output_length(self, input_length: int) -> int:
+        """Calculate output sequence length given input length (matches GLM-ASR)."""
+        # GLM-ASR formula: (L - merge_factor) // merge_factor + 1
+        return (input_length - self.k) // self.k + 1
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Project audio features to LLM embedding space.
+
+        Args:
+            x: Audio encoder output of shape [batch, seq_len, encoder_dim]
+
+        Returns:
+            Projected features of shape [batch, (seq_len - k) // k + 1, llm_dim]
+        """
+        batch, seq, dim = x.shape
+        # Truncate to match GLM-ASR: use (seq - k) // k + 1 frames
+        # This drops trailing frames that don't fill a complete k-frame window
+        out_len = (seq - self.k) // self.k + 1
+        x = x[:, : out_len * self.k, :]  # Truncate to exact multiple
+        x = x.reshape(batch, out_len, dim * self.k)
+
+        x = self.linear_1(x)
+        x = self.norm(x)
+        x = self.act(x)
+        return self.linear_2(x)
+
+
+PROJECTOR_CLASSES = {
+    "mlp": MLPAudioProjector,
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d4aeaf198f783cbf58d8cd59812baac429ffe49147bf9648f6618de20b8d4a4c
+size 17209003

tokenizer_config.json

@@ -0,0 +1,19 @@
+{
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|im_end|>",
+  "extra_special_tokens": [
+    "<audio>"
+  ],
+  "fast": false,
+  "is_local": false,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 131072,
+  "model_specific_special_tokens": {},
+  "pad_token": "<|finetune_right_pad_id|>",
+  "tokenizer_class": "TokenizersBackend"
+}