Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,8 @@ 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
+sample/01_count_bark/sample01_02.wav filter=lfs diff=lfs merge=lfs -text
+sample/02_translate/sample02_02.wav filter=lfs diff=lfs merge=lfs -text
+sample/02_translate/sample02_04.wav filter=lfs diff=lfs merge=lfs -text
+sample/03_cough_music/sample03_01.wav filter=lfs diff=lfs merge=lfs -text
+sample/03_cough_music/sample03_03.wav filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,28 @@
 ---
 title: Audio Interaction Model
-emoji: 🏢
-colorFrom: purple
-colorTo: purple
+emoji: 🎧
+colorFrom: blue
+colorTo: red
 sdk: gradio
 sdk_version: 6.19.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Streaming audio-language model that listens and responds
+python_version: "3.12"
+startup_duration_timeout: 30m
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Audio Interaction Model
+
+This Space demos the **Audio Interaction Model** from the paper
+[Audio Interaction Model](https://arxiv.org/abs/2606.05121).
+
+The model is an always-on streaming audio language model that listens to
+audio and decides for itself when to speak. Upload an audio clip or record
+from microphone, and the model will listen, understand, and respond with text.
+
+## Model
+
+- **Weights**: [zhifeixie/AudioInteraction](https://huggingface.co/zhifeixie/AudioInteraction)
+- **Architecture**: 3B parameter GPT based on Qwen2.5-Omni
+- **Paper**: [arXiv:2606.05121](https://arxiv.org/abs/2606.05121)
+- **Code**: [GitHub](https://github.com/xzf-thu/Audio-Interaction)

app.py

@@ -0,0 +1,391 @@
+"""Audio Interaction Model — Gradio demo on ZeroGPU.
+
+Upload one or more audio clips (or record from microphone) and the model
+will listen, decide whether to respond, and generate a text answer.
+
+Based on the paper: "Audio Interaction Model" (arXiv:2606.05121)
+Model weights: https://huggingface.co/zhifeixie/AudioInteraction
+"""
+
+import os
+
+# Avoid allocator fragmentation under memory spikes
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+# Prevent numba from initializing CUDA before spaces can patch torch
+os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")
+
+# Install openai-whisper with --no-deps to avoid it downgrading torch.
+# Its real runtime deps (tiktoken, more-itertools, tqdm, scipy) are in requirements.txt.
+import subprocess, sys
+try:
+    import whisper
+except ImportError:
+    subprocess.run(
+        [sys.executable, "-m", "pip", "install", "--no-deps", "openai-whisper"],
+        check=True,
+    )
+
+import spaces  # MUST come before any CUDA-touching import
+import torch
+
+# Patch torch.load to disable weights_only (the audio encoder checkpoint
+# was saved with older torch that pickles numpy/object globals).
+_orig_torch_load = torch.load
+torch.load = lambda *a, **k: _orig_torch_load(*a, **{**k, "weights_only": k.get("weights_only", False)})
+
+import json
+import sys as _sys
+import time
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+import gradio as gr
+import numpy as np
+import lightning as L
+from huggingface_hub import snapshot_download
+from transformers import AutoConfig
+from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder
+from safetensors.torch import load_file
+
+# Add the app root to PYTHONPATH so `src.*` and `utils` imports resolve
+_sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+from src.audiointeraction.dataset.TOKENS import (
+    ASSISTANT, AUDIO_BEGIN, ENGLISH, KEEP_SILENCE, ONLINE, PAD,
+    SYSTEM, TEXT_BEGIN, TEXT_END,
+    HAPPY, SAD, ANGRY, SURPRISE, NORMAL, URGENT,
+)
+from src.audiointeraction.generate.base import (
+    AUDIO_TOKENS_PER_CHUNK,
+    sample,
+    encode_audio_chunks,
+    _encode_audio_samples,
+)
+from src.audiointeraction.model import GPT
+from src.audiointeraction.config import Config
+from src.audiointeraction.tokenizer import Tokenizer
+from src.audiointeraction.utils import get_default_supported_precision
+
+MODEL_REPO = "zhifeixie/AudioInteraction"
+EMOTION_EMOJI = {
+    HAPPY: "😊", SAD: "😢", ANGRY: "😠", SURPRISE: "😲",
+    NORMAL: "😐", URGENT: "⚠️",
+}
+
+SYSTEM_PROMPT = (
+    "You are a helpful assistant. When there is no user text, if the audio contains a question, "
+    "please answer it. If it is a sound effect, determine based on the sound whether help is needed."
+)
+
+
+# ---------------------------------------------------------------------------
+# Model loading
+# ---------------------------------------------------------------------------
+
+def _download_checkpoints() -> str:
+    """Download model weights from HuggingFace Hub and return local path."""
+    local_dir = "/home/user/checkpoints"
+    os.makedirs(local_dir, exist_ok=True)
+    snapshot_download(
+        repo_id=MODEL_REPO,
+        repo_type="model",
+        local_dir=local_dir,
+        resume_download=True,
+    )
+    return local_dir
+
+
+CHECKPOINT_DIR = _download_checkpoints()
+
+ckpt = Path(CHECKPOINT_DIR)
+MODEL_CONFIG_DIR = str(ckpt)
+TRAINED_CHECKPOINT = str(ckpt)
+QWEN_OMNI_CKPT = str(ckpt / "qwen25OmniConfig")
+AUDIO_TOWER_CKPT = str(ckpt / "audiointeraction_ChunkwisedEncoder.pth")
+
+
+def _load_model():
+    """Load the GPT model from sharded safetensors."""
+    config = Config.from_file(Path(MODEL_CONFIG_DIR) / "model_config.yaml")
+    # Force SDPA attention (flash_attn not available / sm_120 incompatible with FA3)
+    config.use_flash_attention = False
+
+    fabric = L.Fabric(
+        devices=1, num_nodes=1, strategy="auto",
+        precision=get_default_supported_precision(training=False),
+        loggers="tensorboard",
+    )
+
+    with fabric.init_module(empty_init=False):
+        model = GPT(config)
+    model = fabric.setup(model)
+
+    index_path = Path(TRAINED_CHECKPOINT) / "model.safetensors.index.json"
+    with open(index_path) as f:
+        index = json.load(f)
+    shard_files = sorted(set(index["weight_map"].values()))
+    state_dict = {}
+    for shard in shard_files:
+        state_dict.update(load_file(str(Path(TRAINED_CHECKPOINT) / shard), device="cpu"))
+    missing, unexpected = model.load_state_dict(state_dict, strict=True)
+    if missing or unexpected:
+        print(f"[load_model] missing={missing[:3]}... unexpected={unexpected[:3]}...")
+
+    return fabric, model, config
+
+
+def _load_audio_encoder(device):
+    """Load the Qwen2.5-Omni audio encoder."""
+    cfg = AutoConfig.from_pretrained(QWEN_OMNI_CKPT)
+    audio_cfg = cfg.thinker_config.audio_config
+    encoder = Qwen2_5OmniAudioEncoder._from_config(audio_cfg)
+    state_dict = torch.load(AUDIO_TOWER_CKPT, map_location=device)
+    encoder.load_state_dict(state_dict)
+    encoder.to(device).requires_grad_(False).eval()
+    return encoder
+
+
+print("Loading model and audio encoder...")
+fabric, model, model_config = _load_model()
+model.to("cuda")
+model.eval()
+
+audio_encoder = _load_audio_encoder("cuda")
+tokenizer = Tokenizer(MODEL_CONFIG_DIR)
+
+system_ids = tokenizer.encode(SYSTEM_PROMPT).cpu().tolist()
+prefix_ids = torch.LongTensor(
+    [ONLINE, ENGLISH, SYSTEM, TEXT_BEGIN] + system_ids + [TEXT_END]
+).to(model.device)
+
+with fabric.init_tensor():
+    model.set_kv_cache(batch_size=1)
+model.eval()
+
+print("Model and audio encoder loaded successfully!")
+
+
+# ---------------------------------------------------------------------------
+# Inference
+# ---------------------------------------------------------------------------
+
+@spaces.GPU(duration=120)
+def interact_with_audio(
+    audio_file: str,
+    text_instruction: str = "",
+) -> str:
+    """Listen to an audio clip and generate a text response.
+
+    Upload an audio file (wav/mp3/m4a/flac/ogg) and optionally add a text
+    instruction. The model will listen to the audio and decide whether to
+    respond — it may answer a question, transcribe speech, translate,
+    describe a sound, or stay silent if the audio doesn't need a response.
+    """
+    if audio_file is None:
+        return "Please provide an audio file."
+
+    device = model.device
+
+    # Reset KV cache for a fresh session
+    model.clear_kv_cache()
+    with fabric.init_tensor():
+        model.set_kv_cache(batch_size=1)
+
+    token = prefix_ids
+    input_pos = torch.arange(0, prefix_ids.size(0), device=device, dtype=torch.int64)
+    prompt_size = prefix_ids.size(0)
+
+    # Encode the audio file into chunks
+    audio_chunks = encode_audio_chunks(audio_file, audio_encoder, device)
+
+    # If there's a text instruction, prepend it as text tokens
+    if text_instruction.strip():
+        # Insert text instruction after the system prompt
+        instruction_ids = tokenizer.encode(text_instruction.strip()).cpu().tolist()
+        text_tokens = [TEXT_BEGIN] + instruction_ids + [TEXT_END]
+        text_tensor = torch.LongTensor(text_tokens).to(device)
+        token = torch.cat([token, text_tensor])
+
+    turns: List[List[int]] = []
+    listening = True
+    audio_idx = -1
+    current_turn: List[int] = []
+    text_started = False
+    emotion_prefix = ""
+
+    max_steps = 4096
+
+    for _ in range(max_steps - input_pos.numel()):
+        if listening:
+            audio_idx += 1
+            if audio_idx >= len(audio_chunks):
+                break
+            # Append listening block: [AUDIO_BEGIN, PAD*N, ASSISTANT]
+            new_tokens = torch.LongTensor(
+                [AUDIO_BEGIN] + [PAD] * AUDIO_TOKENS_PER_CHUNK + [ASSISTANT]
+            ).to(device)
+            new_positions = input_pos[-1] + torch.arange(1, len(new_tokens) + 1, device=device)
+            token = torch.cat([token, new_tokens])
+            input_pos = torch.cat([input_pos, new_positions])
+
+            logits = model(
+                token.view(1, -1), None, 1,
+                audio_chunks[audio_idx].to(device),
+                input_pos,
+                input_pos_maxp1=None,
+                audio_tokens_per_chunk=AUDIO_TOKENS_PER_CHUNK,
+            )
+        else:
+            logits = model(
+                token.view(1, -1), None, 1,
+                None,
+                input_pos,
+                input_pos_maxp1=None,
+                audio_tokens_per_chunk=AUDIO_TOKENS_PER_CHUNK,
+            )
+
+        token = sample(logits, temperature=0.0, top_p=0.0).to(torch.int64)
+        int_token = token.item()
+        # Advance input_pos by one
+        new_pos = input_pos[-1].unsqueeze(0).add_(1)
+        input_pos = torch.cat([input_pos, new_pos])
+
+        if listening:
+            if int_token == TEXT_BEGIN:
+                listening = False
+                current_turn = [int_token]
+                text_started = False
+                emotion_prefix = ""
+            elif int_token == KEEP_SILENCE:
+                turns.append([int_token])
+            else:
+                # Unexpected token
+                break
+        else:
+            current_turn.append(int_token)
+            if int_token == TEXT_END:
+                turns.append(current_turn)
+                current_turn = []
+                text_started = False
+                listening = True
+            else:
+                n = len(current_turn)
+                if n == 2:
+                    text_started = True
+                    if int_token in EMOTION_EMOJI:
+                        emotion_prefix = EMOTION_EMOJI[int_token] + " "
+                    else:
+                        piece = tokenizer.decode(torch.tensor([int_token]))
+                        if piece:
+                            # Start output
+                            pass
+                elif n >= 3:
+                    pass
+
+    # Decode all turns into text
+    responses = []
+    for turn in turns:
+        if turn[0] == KEEP_SILENCE:
+            continue
+        # Extract text tokens: skip TEXT_BEGIN and optional emotion tag
+        text_tokens_list = turn[1:]  # skip TEXT_BEGIN
+        if text_tokens_list and text_tokens_list[-1] == TEXT_END:
+            text_tokens_list = text_tokens_list[:-1]
+        # Check if first token is an emotion tag
+        emotion = ""
+        if text_tokens_list and text_tokens_list[0] in EMOTION_EMOJI:
+            emotion = EMOTION_EMOJI[text_tokens_list[0]] + " "
+            text_tokens_list = text_tokens_list[1:]
+
+        if text_tokens_list:
+            decoded = tokenizer.decode(torch.tensor(text_tokens_list))
+            if decoded:
+                responses.append(emotion + decoded)
+
+    model.clear_kv_cache()
+
+    if not responses:
+        return "🔇 The model listened to the audio and chose to stay silent. (This is expected for non-question sounds — the model decides for itself when to speak.)"
+
+    return "\n\n".join(responses)
+
+
+# ---------------------------------------------------------------------------
+# Gradio UI
+# ---------------------------------------------------------------------------
+
+DESCRIPTION = """
+# Audio Interaction Model
+
+This is a demo of the **Audio Interaction Model** from the paper
+[Audio Interaction Model](https://arxiv.org/abs/2606.05121).
+
+The model is an always-on streaming audio language model that listens to
+audio and **decides for itself when to speak**. Upload an audio clip
+(or record from microphone) and the model will:
+
+- 🎧 **Listen** to the audio and understand its content
+- 🧠 **Decide** whether a response is needed
+- 💬 **Respond** with text (transcription, answer, translation, description, etc.)
+
+The model may also choose to stay silent (⟨Silent⟩) if the audio doesn't
+warrant a response — this is a core feature of the Audio Interaction Model,
+not a bug!
+"""
+
+with gr.Blocks(title="Audio Interaction Model") as demo:
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Row():
+        with gr.Column():
+            audio_input = gr.Audio(
+                label="Audio Input",
+                type="filepath",
+                sources=["upload", "microphone"],
+            )
+            text_instruction = gr.Textbox(
+                label="Text Instruction (optional)",
+                placeholder="e.g. 'Translate this to English' or leave empty",
+                lines=2,
+            )
+            submit_btn = gr.Button("Interact", variant="primary")
+
+        with gr.Column():
+            output_text = gr.Textbox(
+                label="Model Response",
+                lines=10,
+                interactive=False,
+            )
+
+    submit_btn.click(
+        fn=interact_with_audio,
+        inputs=[audio_input, text_instruction],
+        outputs=output_text,
+    )
+
+    gr.Examples(
+        examples=[
+            ["sample/01_count_bark/sample01_01.mp3", ""],
+            ["sample/01_count_bark/sample01_02.wav", ""],
+            ["sample/02_translate/sample02_01.mp3", "Translate to English"],
+            ["sample/03_cough_music/sample03_01.wav", ""],
+        ],
+        inputs=[audio_input, text_instruction],
+        outputs=output_text,
+        fn=interact_with_audio,
+        cache_examples=True,
+        cache_mode="lazy",
+    )
+
+    gr.Markdown("""
+    ---
+    **Model**: [zhifeixie/AudioInteraction](https://huggingface.co/zhifeixie/AudioInteraction)
+    | **Paper**: [arXiv:2606.05121](https://arxiv.org/abs/2606.05121)
+    | **Code**: [GitHub](https://github.com/xzf-thu/Audio-Interaction)
+
+    The model is a 3B parameter streaming audio language model based on
+    Qwen2.5-Omni. It processes audio in 400ms chunks and generates text
+    responses when it decides intervention is needed.
+    """)
+
+demo.launch(mcp_server=True)

requirements.txt

@@ -0,0 +1,16 @@
+lightning
+litgpt==0.5.8
+transformers==4.57.6
+safetensors
+tokenizers
+sentencepiece
+torchaudio
+librosa
+numpy
+pyyaml
+einops
+tiktoken
+more-itertools
+tqdm
+scipy
+numba

sample/01_count_bark/sample01_02.wav

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

sample/01_count_bark/sequence.json

@@ -0,0 +1,8 @@
+[
+  "sample01_01.mp3",
+  "sample01_02.wav",
+  "sample01_02.wav",
+  "sample01_02.wav",
+  "sample01_02.wav",
+  "sample01_02.wav"
+]

sample/02_translate/sample02_02.wav

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

sample/02_translate/sample02_04.wav

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

sample/02_translate/sequence.json

@@ -0,0 +1,6 @@
+[
+  "sample02_01.mp3",
+  "sample02_02.wav",
+  "sample02_03.mp3",
+  "sample02_04.wav"
+]

sample/03_cough_music/sample03_01.wav

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

sample/03_cough_music/sample03_03.wav

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

sample/03_cough_music/sequence.json

@@ -0,0 +1,5 @@
+[
+  "sample03_01.wav",
+  "sample03_02.mp3",
+  "sample03_03.wav"
+]

src/audiointeraction/TOKENS.py

@@ -0,0 +1,59 @@
+"""
+Single source of truth for special token IDs used by the audio-enhanced GPT.
+
+Both `cons_online_data.py` and `cons_offline_data.py` import from here so
+that any change to a token id propagates everywhere.
+"""
+
+# === Vocabulary layout ===
+VOCAB_SHIFT = 151600
+
+# === Role / segment markers ===
+USER = VOCAB_SHIFT + 1
+ASSISTANT = VOCAB_SHIFT + 2
+
+TEXT_BEGIN = 151644
+TEXT_END = 151643
+SYSTEM = TEXT_BEGIN  # 151644 — same id used as both turn separator and system marker
+
+# === Audio markers ===
+AUDIO_BEGIN = 151647
+AUDIO_END = 151648
+AUDIO_PAD = 151646
+
+# === Control / padding ===
+PAD = VOCAB_SHIFT + 8
+MASK = -100
+KEEP_SILENCE = VOCAB_SHIFT + 5
+
+# === Task mode (offline batch vs online streaming) ===
+ONLINE = VOCAB_SHIFT + 9
+OFFLINE = VOCAB_SHIFT + 10
+
+# === Language tag ===
+ENGLISH = VOCAB_SHIFT + 11
+
+# === Emotion tags (online only) ===
+HAPPY = VOCAB_SHIFT + 14
+SAD = VOCAB_SHIFT + 15
+ANGRY = VOCAB_SHIFT + 16
+SURPRISE = VOCAB_SHIFT + 17
+NORMAL = VOCAB_SHIFT + 18
+URGENT = VOCAB_SHIFT + 19
+
+EMOTION_TO_ID = {
+    "surprise": SURPRISE,
+    "happy":    HAPPY,
+    "normal":   NORMAL,
+    "sad":      SAD,
+    "angry":    ANGRY,
+    "urgent":   URGENT,
+}
+
+
+HAPPY_SMILE     = "😊"
+SAD_CRY         = "😢"
+ANGRY_MAD       = "😠"
+SURPRISE_WOW    = "😲"
+NORMAL_NEUTRAL  = "😐"
+URGENT_WARNING  = "⚠️"

src/audiointeraction/__init__.py

@@ -0,0 +1,18 @@
+# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
+
+import logging
+import re
+
+from src.audiointeraction.model import GPT  # needs to be imported before config
+from src.audiointeraction.config import Config
+from src.audiointeraction.tokenizer import Tokenizer
+
+# Suppress excessive warnings, see https://github.com/pytorch/pytorch/issues/111632
+pattern = re.compile(".*Profiler function .* will be ignored")
+logging.getLogger("torch._dynamo.variables.torch").addFilter(lambda record: not pattern.search(record.getMessage()))
+
+# Avoid printing state-dict profiling output at the WARNING level when saving a checkpoint
+logging.getLogger("torch.distributed.fsdp._optim_utils").disabled = True
+logging.getLogger("torch.distributed.fsdp._debug_utils").disabled = True
+
+__all__ = ["GPT", "Config", "Tokenizer"]

src/audiointeraction/args.py

@@ -0,0 +1,50 @@
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class TrainArgs:
+    """Training-related arguments."""
+
+    save_interval: Optional[int] = 1000
+    """Number of optimizer steps between checkpoint saves."""
+    log_interval: int = 1
+    """Number of iterations between log lines."""
+    global_batch_size: int = 64
+    """Total samples per optimizer step across all data-parallel ranks."""
+    micro_batch_size: int = 4
+    """Samples per data-parallel rank per forward/backward pass."""
+    lr_warmup_steps: Optional[int] = 100
+    """Number of warmup iterations (linear ramp from 0 to max_lr)."""
+    epochs: Optional[int] = None
+    """Number of epochs to train."""
+    max_steps: Optional[int] = None
+    """Hard cap on optimizer steps (overrides epochs if reached first)."""
+    max_seq_length: Optional[int] = None
+    """Truncate samples longer than this."""
+
+    def gradient_accumulation_iters(self, devices: int, num_nodes: int = 1) -> int:
+        n = self.batch_size(devices, num_nodes) // self.micro_batch_size
+        assert n > 0
+        return n
+
+    def batch_size(self, devices: int, num_nodes: int = 1) -> int:
+        n = self.global_batch_size // (devices * num_nodes)
+        assert n > 0
+        return n
+
+
+@dataclass
+class EvalArgs:
+    """Evaluation-related arguments."""
+
+    interval: int = 600
+    """Number of optimizer steps between validation runs."""
+    max_new_tokens: Optional[int] = None
+    """Generation length cap (only used by validation-time generation paths)."""
+    max_iters: int = 100
+    """Max number of validation batches per run."""
+    initial_validation: bool = False
+    """Run one validation pass before the first training step."""
+    final_validation: bool = True
+    """Run one validation pass after training completes."""

src/audiointeraction/config.py

@@ -0,0 +1,150 @@
+# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
+
+from copy import deepcopy
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Literal, Optional, Type, Union
+
+import torch
+import yaml
+from typing_extensions import Self
+from src.audiointeraction.utils import find_multiple
+
+
+@dataclass
+class Config:
+    name: str = ""
+    hf_config: dict = field(default_factory=dict)
+    # General size parameters
+    block_size: int = 4096
+    n_layer: int = 16
+    n_embd: int = 4096
+    vocab_size: int = 50254
+    padding_multiple: int = 512
+    padded_vocab_size: Optional[int] = None
+    # Transformer block (structure, normalizations)
+    norm_class_name: Literal["LayerNorm", "RMSNorm"] = "LayerNorm"
+    norm_eps: float = 1e-5
+    norm_qk: bool = False
+    post_attention_norm: bool = False
+    post_mlp_norm: bool = False
+    parallel_residual: bool = True
+    shared_attention_norm: bool = False
+    # Transformer block (self-attention)
+    n_head: int = 32
+    head_size: Optional[int] = None
+    # to use multi-head attention (MHA), set this to `n_head` (default)
+    # to use multi-query attention (MQA), set this to 1
+    # to use grouped-query attention (GQA), set this to a value in between
+    # Example with `n_head=4`
+    # ┌───┐┌───┐┌───┐┌───┐     ┌───┐    ┌───┐             ┌───┐
+    # │ v ││ v ││ v ││ v │     │ v │    │ v │             │ v │
+    # └───┘└───┘└───┘└───┘     └───┘    └───┘             └───┘
+    #   │    │    │    │         │        │                 │
+    # ┌───┐┌───┐┌───┐┌───┐     ┌───┐    ┌───┐             ┌───┐
+    # │ k ││ k ││ k ││ k │     │ k │    │ k │             │ k │
+    # └───┘└───┘└───┘└───┘     └───┘    └───┘             └───┘
+    #   │    │    │    │      ┌──┴──┐  ┌──┴──┐      ┌────┬──┴─┬────┐
+    # ┌───┐┌───┐┌───┐┌───┐  ┌───┐┌───┐┌───┐┌───┐  ┌───┐┌───┐┌───┐┌───┐
+    # │ q ││ q ││ q ││ q │  │ q ││ q ││ q ││ q │  │ q ││ q ││ q ││ q │
+    # └───┘└───┘└───┘└───┘  └───┘└───┘└───┘└───┘  └───┘└───┘└───┘└───┘
+    # ◀──────────────────▶  ◀──────────────────▶  ◀──────────────────▶
+    #         MHA                    GQA                   MQA
+    #   n_query_groups=4       n_query_groups=2      n_query_groups=1
+    #
+    # credit https://arxiv.org/pdf/2305.13245.pdf
+    n_query_groups: Optional[int] = None
+    attn_bias: bool = False
+    attention_scores_scalar: Optional[int] = None
+    sliding_window_size: Optional[int] = None
+    sliding_window_layer_placing: Optional[Literal["all", "interleaved"]] = None
+    # if `attention_logit_softcapping` is used, cannot use optimized
+    # `torch.nn.functional.scaled_dot_product_attention` (which implements
+    # Flash attention), may result in higher memory and runtime footprint.
+    attention_logit_softcapping: Optional[float] = None
+    # Rotary position embedding (RoPE)
+    rope_base: int = 10000
+    rotary_percentage: float = 0.25
+    rope_condense_ratio: int = 1
+    rope_adjustments: Optional[dict] = None
+    # Transformer block (MLP)
+    intermediate_size: Optional[int] = None
+    bias: bool = True
+    mlp_class_name: Literal["GptNeoxMLP", "LLaMAMLP", "GemmaMLP", "LLaMAMoE"] = "GptNeoxMLP"
+    gelu_approximate: str = "none"
+    n_expert: int = 0
+    n_expert_per_token: int = 0
+    # GPT before/after blocks
+    scale_embeddings: bool = False
+    lm_head_bias: bool = False
+    final_logit_softcapping: Optional[float] = None
+
+    def __post_init__(self):
+        if not self.name:
+            self.name = self.hf_config.get("name", self.name)
+
+        if self.head_size is None:
+            assert self.n_embd % self.n_head == 0
+            self.head_size = self.n_embd // self.n_head
+
+        # vocab size should be a power of 2 to be optimal on hardware. compute the closest value
+        if self.padded_vocab_size is None:
+            self.padded_vocab_size = find_multiple(self.vocab_size, self.padding_multiple)
+        else:
+            # vocab size shouldn't be larger than padded vocab size
+            self.vocab_size = min(self.vocab_size, self.padded_vocab_size)
+
+        # compute the number of query groups
+        if self.n_query_groups is not None:
+            assert self.n_head % self.n_query_groups == 0
+        else:
+            self.n_query_groups = self.n_head
+
+        # compute the intermediate size for MLP if not set
+        if self.intermediate_size is None:
+            if self.mlp_class_name == "LLaMAMLP":
+                raise ValueError(f"The config {self.name!r}, needs to set the `intermediate_size`")
+            self.intermediate_size = 4 * self.n_embd
+
+        self.rope_n_elem = int(self.rotary_percentage * self.head_size)
+
+        if self.sliding_window_size is not None:
+            self.sliding_window_layer_stride = (
+                1 if (self.sliding_window_layer_placing is None or self.sliding_window_layer_placing == "all") else 2
+            )
+
+    @classmethod
+    def from_file(cls, path: Union[str, Path], **kwargs: Any) -> Self:
+        with open(path, encoding="utf-8") as fp:
+            file_kwargs = yaml.safe_load(fp)
+            if file_kwargs is None:
+                raise ValueError(f"{path} is empty which is likely unexpected.")
+        file_kwargs.update(kwargs)
+        return cls(**file_kwargs)
+
+    @property
+    def mlp_class(self) -> Type:
+        # `self.mlp_class_name` cannot be the type to keep the config serializable
+        import src.audiointeraction.model
+        return getattr(src.audiointeraction.model, self.mlp_class_name)
+
+    @property
+    def norm_class(self) -> Type:
+        # `self.norm_class_name` cannot be the type to keep the config serializable
+
+        from functools import partial
+
+        if self.norm_class_name == "RMSNorm":
+
+            from src.audiointeraction.model import RMSNorm
+
+            return partial(RMSNorm, add_unit_offset="Gemma" in self.name)
+
+        if self.norm_class_name == "LayerNorm" and "OLMo" in self.name:
+            # this makes it equivalent to `torch.nn.functional.layer_norm`
+            # that is used by OLMo
+            # Table 5 caption in the OLMo paper shows this - https://aclanthology.org/2024.acl-long.841
+            return partial(torch.nn.LayerNorm, elementwise_affine=False)
+
+        return getattr(torch.nn, self.norm_class_name)
+

src/audiointeraction/dataset/TOKENS.py

@@ -0,0 +1,59 @@
+"""
+Single source of truth for special token IDs used by the audio-enhanced GPT.
+
+Both `cons_online_data.py` and `cons_offline_data.py` import from here so
+that any change to a token id propagates everywhere.
+"""
+
+# === Vocabulary layout ===
+VOCAB_SHIFT = 151600
+
+# === Role / segment markers ===
+USER = VOCAB_SHIFT + 1
+ASSISTANT = VOCAB_SHIFT + 2
+
+TEXT_BEGIN = 151644
+TEXT_END = 151643
+SYSTEM = TEXT_BEGIN  # 151644 — same id used as both turn separator and system marker
+
+# === Audio markers ===
+AUDIO_BEGIN = 151647
+AUDIO_END = 151648
+AUDIO_PAD = 151646
+
+# === Control / padding ===
+PAD = VOCAB_SHIFT + 8
+MASK = -100
+KEEP_SILENCE = VOCAB_SHIFT + 5
+
+# === Task mode (offline batch vs online streaming) ===
+ONLINE = VOCAB_SHIFT + 9
+OFFLINE = VOCAB_SHIFT + 10
+
+# === Language tag ===
+ENGLISH = VOCAB_SHIFT + 11
+
+# === Emotion tags (online only) ===
+HAPPY = VOCAB_SHIFT + 14
+SAD = VOCAB_SHIFT + 15
+ANGRY = VOCAB_SHIFT + 16
+SURPRISE = VOCAB_SHIFT + 17
+NORMAL = VOCAB_SHIFT + 18
+URGENT = VOCAB_SHIFT + 19
+
+EMOTION_TO_ID = {
+    "surprise": SURPRISE,
+    "happy":    HAPPY,
+    "normal":   NORMAL,
+    "sad":      SAD,
+    "angry":    ANGRY,
+    "urgent":   URGENT,
+}
+
+
+HAPPY_SMILE     = "😊"
+SAD_CRY         = "😢"
+ANGRY_MAD       = "😠"
+SURPRISE_WOW    = "😲"
+NORMAL_NEUTRAL  = "😐"
+URGENT_WARNING  = "⚠️"

src/audiointeraction/dataset/cons_offline_data.py

@@ -0,0 +1,206 @@
+"""Build offline single-turn SFT samples.
+
+Token layout: prefix + [audio_block?] + [user_block?] + assistant_block.
+Three task variants come from which blocks are present:
+  - audio + user  -> A_T_T
+  - audio         -> A_T
+  -         user  -> T_T
+
+Input  (per line): {"conversation": [{"user", "assistant", "audio_path"}, ...]} (first turn only)
+                   or flat {"user", "assistant", "audio_path"}.
+Output (per line): {"tasks": "offline", "idx", "input_ids", "labels", "audio_pos", "pt_path_dir"}.
+
+Edit the constants below, then `python cons_offline_data.py`.
+Re-running picks up where the previous run stopped (already-written idx are skipped).
+"""
+
+import json
+import os
+
+import torch
+from tqdm import tqdm
+from transformers import AutoConfig, Qwen2_5OmniForConditionalGeneration
+
+from src.audiointeraction.dataset.utils.load_audio import _load_mel
+from src.audiointeraction.dataset.TOKENS import (
+    ASSISTANT, AUDIO_BEGIN, AUDIO_END, AUDIO_PAD, ENGLISH, MASK,
+    OFFLINE, SYSTEM, TEXT_BEGIN, TEXT_END, USER,
+)
+from src.audiointeraction.generate.base import resolve_checkpoint_paths
+from src.audiointeraction.tokenizer import Tokenizer
+
+
+# ============================================================
+# Fill these in before running.
+# ============================================================
+CHECKPOINT_DIR = ""   # checkpoint root (tokenizer + qwen_2_5_omni_config + audiointeraction_ChunkwisedEncoder.pth)
+INPUT_JSONL    = ""   # raw input jsonl
+OUTPUT_JSONL   = ""   # training-ready output jsonl
+ERROR_LOG      = ""   # path to append per-sample error messages
+FEATURE_DIR    = ""   # dir to save audio feature .pt files (one subdir per sample)
+DEVICE         = "cuda"
+# ============================================================
+
+
+DEFAULT_SYSTEM_PROMPT = (
+    "You are a helpful assistant. When there is no user text, if the audio contains a question, "
+    "please answer it. If it is a sound effect, determine based on the sound whether help is needed."
+)
+
+# Cap offline audio length at 20 seconds (matches the prior behavior).
+MAX_AUDIO_SECONDS = 20
+
+
+# === Audio encoder (lazy, cached per device) ===
+_audio_encoder_cache = {}
+
+
+def _get_audio_encoder(qwen_omni_ckpt, audio_tower_ckpt, device):
+    key = (qwen_omni_ckpt, audio_tower_ckpt, str(device))
+    if key not in _audio_encoder_cache:
+        print(f"[offline] loading audio_encoder on {device} ...")
+        cfg = AutoConfig.from_pretrained(qwen_omni_ckpt)
+        enc = Qwen2_5OmniForConditionalGeneration._from_config(cfg).thinker.audio_tower
+        enc.load_state_dict(torch.load(audio_tower_ckpt, map_location=device))
+        enc.to(device).eval()
+        _audio_encoder_cache[key] = enc
+    return _audio_encoder_cache[key]
+
+
+def _extract_and_save_audio_feature(audio_path, save_dir, *,
+                                    qwen_omni_ckpt, audio_tower_ckpt, device):
+    """Encode audio, save to `<save_dir>/AudioFeat.pt`, return output_len."""
+    _, mel, len_feature, input_len, _ = _load_mel(audio_path, max_seconds=MAX_AUDIO_SECONDS)
+    encoder = _get_audio_encoder(qwen_omni_ckpt, audio_tower_ckpt, device)
+    with torch.no_grad():
+        hidden = encoder(
+            mel.to(device),
+            torch.tensor([len_feature]).to(device),
+            torch.tensor([input_len]).to(device),
+        ).last_hidden_state
+    hidden = hidden.squeeze(0).cpu()
+    output_len = hidden.shape[0]
+    os.makedirs(save_dir, exist_ok=True)
+    torch.save(hidden, os.path.join(save_dir, "AudioFeat.pt"))
+    return output_len
+
+
+# === Input parsing ===
+
+def _extract_first_turn(data_item):
+    """Return (user_text, assistant_text, audio_path). Offline is single-turn,
+    so a `conversation` list contributes only its first turn.
+    """
+    convs = data_item.get("conversation", [])
+    src = convs[0] if convs else data_item
+    return (
+        src.get("user"),
+        src.get("assistant"),
+        src.get("audio_path") or src.get("merge_path"),
+    )
+
+
+# === Sample builder ===
+
+def build_offline_sample(data_item, idx, *,
+                         tokenizer, system_ids, feature_dir,
+                         qwen_omni_ckpt, audio_tower_ckpt, device):
+    """Build one offline SFT sample. Returns a dict with input_ids/labels/audio_pos/pt_path_dir."""
+    user_text, assistant_text, audio_path = _extract_first_turn(data_item)
+    if not assistant_text or assistant_text == "None":
+        raise ValueError("missing assistant text")
+
+    has_audio = audio_path not in (None, "", "None")
+    has_user = user_text not in (None, "", "None")
+    assistant_ids = tokenizer.encode(assistant_text).cpu().tolist()
+    user_ids = tokenizer.encode(user_text).cpu().tolist() if has_user else []
+
+    if has_audio:
+        pt_path_dir = os.path.join(feature_dir, str(idx))
+        output_len = _extract_and_save_audio_feature(
+            audio_path, pt_path_dir,
+            qwen_omni_ckpt=qwen_omni_ckpt, audio_tower_ckpt=audio_tower_ckpt, device=device,
+        )
+    else:
+        pt_path_dir = None
+        output_len = 0
+
+    prefix          = [OFFLINE, ENGLISH, SYSTEM, TEXT_BEGIN] + system_ids + [TEXT_END]
+    audio_block     = [AUDIO_BEGIN] + [AUDIO_PAD] * output_len + [AUDIO_END]    if has_audio else []
+    user_block      = [USER, TEXT_BEGIN] + user_ids + [TEXT_END]                if has_user  else []
+    assistant_block = [ASSISTANT, TEXT_BEGIN] + assistant_ids + [TEXT_END]
+
+    input_ids = prefix + audio_block + user_block + assistant_block
+    labels    = [MASK] * (len(prefix) + len(audio_block) + len(user_block)) + assistant_block
+
+    # AUDIO_PAD sits right after [AUDIO_BEGIN] (which is at index len(prefix)).
+    audio_pos = [(len(prefix) + 1, len(prefix) + 1 + output_len)] if has_audio else None
+
+    return {
+        "input_ids": input_ids,
+        "labels": labels,
+        "audio_pos": audio_pos,
+        "pt_path_dir": pt_path_dir,
+    }
+
+
+# === Main driver ===
+
+def _load_processed_indices(output_path):
+    """For resume: collect idx of all valid records already in output_path."""
+    done = set()
+    if not os.path.exists(output_path):
+        return done
+    with open(output_path, "r", encoding="utf-8") as fr:
+        for line in fr:
+            try:
+                done.add(json.loads(line)["idx"])
+            except (json.JSONDecodeError, KeyError):
+                continue
+    return done
+
+
+def main():
+    for name, value in [("CHECKPOINT_DIR", CHECKPOINT_DIR), ("INPUT_JSONL", INPUT_JSONL),
+                        ("OUTPUT_JSONL", OUTPUT_JSONL), ("ERROR_LOG", ERROR_LOG),
+                        ("FEATURE_DIR", FEATURE_DIR)]:
+        if not value:
+            raise SystemExit(f"Set {name} at the top of cons_offline_data.py before running.")
+
+    tokenizer_dir, _, qwen_omni_ckpt, audio_tower_ckpt = resolve_checkpoint_paths(CHECKPOINT_DIR)
+    os.makedirs(FEATURE_DIR, exist_ok=True)
+
+    with open(INPUT_JSONL, "r", encoding="utf-8") as f:
+        data_lines = f.readlines()
+
+    processed = _load_processed_indices(OUTPUT_JSONL)
+    todo = [i for i in range(len(data_lines)) if i not in processed]
+    print(f"Total {len(data_lines)} | already done {len(processed)} | to process {len(todo)}")
+
+    tokenizer = Tokenizer(tokenizer_dir)
+    system_ids = tokenizer.encode(DEFAULT_SYSTEM_PROMPT).cpu().tolist()
+
+    with open(OUTPUT_JSONL, "a", encoding="utf-8", buffering=1) as fout, \
+         open(ERROR_LOG,   "a", encoding="utf-8", buffering=1) as ferr:
+        for idx in tqdm(todo):
+            try:
+                data_item = json.loads(data_lines[idx])
+                sample = build_offline_sample(
+                    data_item, idx,
+                    tokenizer=tokenizer, system_ids=system_ids,
+                    feature_dir=FEATURE_DIR,
+                    qwen_omni_ckpt=qwen_omni_ckpt,
+                    audio_tower_ckpt=audio_tower_ckpt,
+                    device=DEVICE,
+                )
+                fout.write(json.dumps(
+                    {"tasks": "offline", "idx": idx, **sample},
+                    ensure_ascii=False,
+                ) + "\n")
+            except Exception as e:
+                ferr.write(f"[idx {idx}] {type(e).__name__}: {e}\n")
+                ferr.write(data_lines[idx])
+
+
+if __name__ == "__main__":
+    main()

src/audiointeraction/dataset/cons_online_data.py

@@ -0,0 +1,387 @@
+"""Build online (streaming) SFT samples in 4 inspectable steps.
+
+Pipeline:
+  step1 — sample a leading-silence length per turn + a tail silence
+          (encoder-output frames). No audio touched.
+  step2 — load every turn's audio, prepend the sampled silence, concat into
+          one waveform, write `<wavs_dir>/<idx>.wav` (16 kHz mono int16).
+  step3 — run the Qwen2.5-Omni audio_tower on each wav and dump
+          `<features_dir>/<idx>/AudioFeat.pt`.
+  step4 — lay out the token-level streaming sequence (input_ids / labels /
+          audio_pos) and write a training-ready jsonl.
+
+Edit the constants below, then `python cons_online_data.py`. Each step's
+intermediate jsonl is left on disk so you can re-run any step independently
+or inspect the artifacts.
+"""
+
+import glob
+import json
+import os
+import random
+import wave
+
+import numpy as np
+import soundfile as sf
+import whisper
+from tqdm import tqdm
+
+from src.audiointeraction.dataset.utils.extract_online_feature import extract_audio_features
+from src.audiointeraction.dataset.utils.load_audio import SAMPLES_PER_FRAME, _load_mel
+from src.audiointeraction.dataset.TOKENS import (
+    ASSISTANT, AUDIO_BEGIN, EMOTION_TO_ID, ENGLISH, KEEP_SILENCE, MASK,
+    NORMAL, ONLINE, PAD, SYSTEM, TEXT_BEGIN, TEXT_END,
+)
+from src.audiointeraction.generate.base import resolve_checkpoint_paths
+from src.audiointeraction.tokenizer import Tokenizer
+
+
+# ============================================================
+# Fill these in before running.
+# ============================================================
+CHECKPOINT_DIR  = ""    # checkpoint root (tokenizer + qwen_2_5_omni_config + audiointeraction_ChunkwisedEncoder.pth)
+INPUT_JSONL     = ""    # user-supplied raw jsonl
+WORK_DIR        = ""    # holds intermediate step1/2/3 jsonls + wavs/ + features/
+OUT_TRAIN_JSONL = ""    # final training-ready jsonl
+NOISE_DIR       = ""    # dir of background-noise audio files (wav/flac/ogg/mp3); recursed
+
+MIN_NOISE_LEN = 20      # min leading/tail noise length per turn (encoder frames)
+MAX_NOISE_LEN = 60      # max leading/tail noise length per turn
+CHUNK_SIZE    = 10      # encoder-output frames per audio chunk
+SEED          = 1337
+DEVICE        = "cuda"
+# ============================================================
+
+
+DEFAULT_SYSTEM_PROMPT = (
+    "You are a helpful assistant. When there is no user text, if the audio contains a question, "
+    "please answer it. If it is a sound effect, determine based on the sound whether help is needed."
+)
+
+NO_RESPONSE_MARKERS = {"<no need to response>", "no need to response"}
+
+
+# === Step 1 — sample silence (filled from a noise library) ===
+
+NOISE_AUDIO_EXTS = (".wav", ".flac", ".ogg", ".mp3")
+
+
+def _scan_noise_dir(noise_dir):
+    """Walk noise_dir and return [(path, duration_s), ...] for every audio file."""
+    paths = []
+    for root, _, files in os.walk(noise_dir):
+        for f in files:
+            if f.lower().endswith(NOISE_AUDIO_EXTS):
+                paths.append(os.path.join(root, f))
+    index = []
+    for p in tqdm(paths, desc="scan noise"):
+        try:
+            index.append((p, sf.info(p).duration))
+        except Exception:
+            # sf.info won't read mp3; fall back to librosa.
+            try:
+                import librosa
+                index.append((p, librosa.get_duration(path=p)))
+            except Exception as e:
+                print(f"[warn] skip noise {p}: {e}")
+    return index
+
+
+def _pick_noise(noise_index, needed_s, *, margin_s=0.0, max_tries=100):
+    """Pick a random noise file. If the draw is shorter than needed, re-draw.
+    Returns (path, start_s) — a slice of `needed_s + margin_s` is guaranteed to fit."""
+    budget = needed_s + margin_s
+    for _ in range(max_tries):
+        path, dur = random.choice(noise_index)
+        if dur >= budget:
+            start_s = random.uniform(0.0, dur - budget)
+            return path, round(start_s, 3)
+    raise RuntimeError(
+        f"Couldn't find a noise file long enough for {budget:.2f}s after {max_tries} tries — "
+        f"add longer files to NOISE_DIR."
+    )
+
+
+def _extract_turns(data_item):
+    """Return list of {audio_path, assistant, emotion} dicts."""
+    convs = data_item.get("conversation", [])
+    if convs:
+        return [
+            {"audio_path": c["audio_path"],
+             "assistant":  c["assistant"],
+             "emotion":    c.get("emotion") or "normal"}
+            for c in convs
+        ]
+    if "merge_path" in data_item and "assistant" in data_item:
+        return [{
+            "audio_path": data_item["merge_path"],
+            "assistant":  data_item["assistant"],
+            "emotion":    data_item.get("emotion", "normal"),
+        }]
+    raise ValueError("missing 'conversation' or single-turn fields")
+
+
+def step1_sample_silence(input_jsonl, output_jsonl, *,
+                         noise_dir, min_noise_len, max_noise_len, chunk_size, seed):
+    """For each turn, pick a leading noise slice (file + start_s); also one for tail.
+
+    The tail length is adjusted to a chunk boundary inside step 2, so its noise
+    slice is selected with a `chunk_size`-frame margin to guarantee the eventual
+    actual length fits in the picked file.
+    """
+    random.seed(seed)
+    os.makedirs(os.path.dirname(os.path.abspath(output_jsonl)) or ".", exist_ok=True)
+    noise_index = _scan_noise_dir(noise_dir)
+    if not noise_index:
+        raise RuntimeError(f"No audio files found under {noise_dir}")
+    tail_margin_s = chunk_size * 0.04   # one chunk = 10 frames × 40 ms
+
+    with open(input_jsonl, "r", encoding="utf-8") as fin, \
+         open(output_jsonl, "w", encoding="utf-8") as fout:
+        for idx, line in enumerate(tqdm(fin.readlines(), desc="step1")):
+            try:
+                data_item = json.loads(line)
+                turns = _extract_turns(data_item)
+                for t in turns:
+                    t["leading_silence_frames"] = random.randint(min_noise_len, max_noise_len)
+                    needed_s = t["leading_silence_frames"] * 0.04
+                    t["leading_noise_path"], t["leading_noise_start_s"] = \
+                        _pick_noise(noise_index, needed_s)
+                tail_frames = random.randint(min_noise_len, max_noise_len)
+                tail_path, tail_start_s = _pick_noise(
+                    noise_index, tail_frames * 0.04, margin_s=tail_margin_s,
+                )
+                rec = {
+                    "idx": idx,
+                    "turns": turns,
+                    "tail_silence_frames": tail_frames,
+                    "tail_noise_path": tail_path,
+                    "tail_noise_start_s": tail_start_s,
+                }
+                fout.write(json.dumps(rec, ensure_ascii=False) + "\n")
+            except Exception as e:
+                print(f"[step1 idx {idx}] {type(e).__name__}: {e}")
+
+
+# === Step 2 — concatenate audio per sample ===
+
+def _load_audio_aligned(audio_path):
+    """Load audio @ 16 kHz, pad/crop to (output_len * SAMPLES_PER_FRAME) samples.
+    Returns (np.float32 array, output_len_frames)."""
+    audio, _, _, _, output_len = _load_mel(audio_path)
+    arr = np.asarray(audio, dtype=np.float32)
+    target = output_len * SAMPLES_PER_FRAME
+    if len(arr) < target:
+        arr = np.concatenate([arr, np.zeros(target - len(arr), dtype=np.float32)])
+    elif len(arr) > target:
+        max_start = (len(arr) - target) // SAMPLES_PER_FRAME
+        start = random.randint(0, max_start) * SAMPLES_PER_FRAME
+        arr = arr[start: start + target]
+    return arr, output_len
+
+
+def _load_noise_segment(noise_path, start_s, n_samples):
+    """Pull exactly n_samples from noise_path starting at start_s (16 kHz, mono float).
+    Raises if the file actually yields fewer samples than asked (step 1 chose this
+    file based on metadata duration — caller should treat it as a data error)."""
+    full = whisper.load_audio(noise_path, sr=16000)   # float32 in [-1, 1]
+    start_idx = int(round(start_s * 16000))
+    seg = full[start_idx: start_idx + n_samples]
+    if len(seg) < n_samples:
+        raise ValueError(
+            f"Noise {noise_path} only yields {len(seg)} samples from start {start_s}s, "
+            f"need {n_samples}."
+        )
+    return seg
+
+
+def _write_wav(path, float_samples):
+    """Write a 16 kHz mono int16 wav from float samples in [-1, 1]."""
+    arr = np.asarray(float_samples, dtype=np.float32)
+    arr_i16 = np.clip(arr * 32767.0, -32768, 32767).astype(np.int16)
+    with wave.open(path, "wb") as wf:
+        wf.setnchannels(1)
+        wf.setsampwidth(2)
+        wf.setframerate(16000)
+        wf.writeframes(arr_i16.tobytes())
+
+
+def step2_concat_audio(input_jsonl, output_jsonl, wavs_dir, *, chunk_size, seed):
+    """For each step-1 record, splice [noise → audio → noise → audio → ... → tail noise]
+    into one wav using the noise slices step 1 picked. Records audio_frames per turn."""
+    random.seed(seed)
+    os.makedirs(wavs_dir, exist_ok=True)
+    os.makedirs(os.path.dirname(os.path.abspath(output_jsonl)) or ".", exist_ok=True)
+    with open(input_jsonl, "r", encoding="utf-8") as fin, \
+         open(output_jsonl, "w", encoding="utf-8") as fout:
+        for line in tqdm(fin.readlines(), desc="step2"):
+            rec = json.loads(line)
+            try:
+                segments = []
+                total_frames = 0
+                for t in rec["turns"]:
+                    n = t["leading_silence_frames"] * SAMPLES_PER_FRAME
+                    segments.append(_load_noise_segment(
+                        t["leading_noise_path"], t["leading_noise_start_s"], n,
+                    ))
+                    seg, n_frames = _load_audio_aligned(t["audio_path"])
+                    segments.append(seg)
+                    t["audio_frames"] = n_frames
+                    total_frames += t["leading_silence_frames"] + n_frames
+
+                # Round tail length so the total lands on a chunk boundary.
+                tail = rec["tail_silence_frames"] - (total_frames + rec["tail_silence_frames"]) % chunk_size
+                if tail < 0:
+                    tail = (chunk_size - (total_frames % chunk_size)) % chunk_size
+                segments.append(_load_noise_segment(
+                    rec["tail_noise_path"], rec["tail_noise_start_s"], tail * SAMPLES_PER_FRAME,
+                ))
+                rec["tail_silence_frames_actual"] = tail
+
+                wav_path = os.path.join(wavs_dir, f"{rec['idx']}.wav")
+                _write_wav(wav_path, np.concatenate(segments))
+                rec["concat_wav_path"] = wav_path
+                fout.write(json.dumps(rec, ensure_ascii=False) + "\n")
+            except Exception as e:
+                print(f"[step2 idx {rec.get('idx')}] {type(e).__name__}: {e}")
+
+
+# === Step 3 — run audio_tower on each wav ===
+
+def step3_extract_features(input_jsonl, output_jsonl, features_dir, *,
+                           qwen_omni_ckpt, audio_tower_ckpt, device):
+    """For each step-2 record, encode the concatenated wav into AudioFeat.pt."""
+    os.makedirs(features_dir, exist_ok=True)
+    os.makedirs(os.path.dirname(os.path.abspath(output_jsonl)) or ".", exist_ok=True)
+    with open(input_jsonl, "r", encoding="utf-8") as fin, \
+         open(output_jsonl, "w", encoding="utf-8") as fout:
+        for line in tqdm(fin.readlines(), desc="step3"):
+            rec = json.loads(line)
+            try:
+                audio = whisper.load_audio(rec["concat_wav_path"], sr=16000)
+                pt_path_dir = os.path.join(features_dir, str(rec["idx"]))
+                extract_audio_features(
+                    audio, pt_path_dir,
+                    qwen_omni_ckpt=qwen_omni_ckpt,
+                    audio_tower_ckpt=audio_tower_ckpt,
+                    device=device,
+                )
+                rec["pt_path_dir"] = pt_path_dir
+                fout.write(json.dumps(rec, ensure_ascii=False) + "\n")
+            except Exception as e:
+                print(f"[step3 idx {rec.get('idx')}] {type(e).__name__}: {e}")
+
+
+# === Step 4 — build token-level training samples ===
+
+def _silence_chunk(chunk_size):
+    """Mid-turn or tail chunk: model should keep waiting."""
+    ids = [AUDIO_BEGIN] + [PAD] * chunk_size + [ASSISTANT, KEEP_SILENCE]
+    labels = [MASK] + [MASK] * chunk_size + [MASK, KEEP_SILENCE]
+    return ids, labels
+
+
+def _response_chunk(assistant_text, emotion, tokenizer, chunk_size):
+    """Last chunk of a turn: actual reply, or silence if marked no-response."""
+    if isinstance(assistant_text, str) and assistant_text.strip().lower() in NO_RESPONSE_MARKERS:
+        return _silence_chunk(chunk_size)
+    emotion_tok = EMOTION_TO_ID.get(emotion.lower(), NORMAL) if isinstance(emotion, str) else NORMAL
+    assistant_ids = tokenizer.encode(assistant_text).cpu().tolist()
+    ids    = [AUDIO_BEGIN] + [PAD] * chunk_size + [ASSISTANT, TEXT_BEGIN, emotion_tok] + assistant_ids + [TEXT_END]
+    labels = [MASK]        + [MASK] * chunk_size + [MASK,      TEXT_BEGIN, emotion_tok] + assistant_ids + [TEXT_END]
+    return ids, labels
+
+
+def step4_build_tokens(input_jsonl, output_jsonl, *, tokenizer_dir, chunk_size):
+    """For each step-3 record, lay out input_ids / labels / audio_pos."""
+    os.makedirs(os.path.dirname(os.path.abspath(output_jsonl)) or ".", exist_ok=True)
+    tokenizer = Tokenizer(tokenizer_dir)
+    system_ids = tokenizer.encode(DEFAULT_SYSTEM_PROMPT).cpu().tolist()
+    with open(input_jsonl, "r", encoding="utf-8") as fin, \
+         open(output_jsonl, "w", encoding="utf-8") as fout:
+        for line in tqdm(fin.readlines(), desc="step4"):
+            rec = json.loads(line)
+            try:
+                input_ids = [ONLINE, ENGLISH, SYSTEM, TEXT_BEGIN] + system_ids + [TEXT_END]
+                labels    = [MASK] * len(input_ids)
+                audio_pos = []
+
+                total_frames = 0
+                for i, t in enumerate(rec["turns"]):
+                    new_chunks = (
+                        (total_frames + t["leading_silence_frames"] + t["audio_frames"]) // chunk_size
+                        - total_frames // chunk_size
+                    )
+                    total_frames += t["leading_silence_frames"] + t["audio_frames"]
+                    # First turn gets one extra trailing chunk so the reply has room to land.
+                    if i == 0:
+                        new_chunks += 1
+                    for j in range(new_chunks):
+                        pos_start = len(input_ids) + 1
+                        audio_pos.append((pos_start, pos_start + chunk_size))
+                        if j == new_chunks - 1:
+                            chunk_ids, chunk_labels = _response_chunk(t["assistant"], t["emotion"], tokenizer, chunk_size)
+                        else:
+                            chunk_ids, chunk_labels = _silence_chunk(chunk_size)
+                        input_ids += chunk_ids
+                        labels    += chunk_labels
+
+                tail = rec["tail_silence_frames_actual"]
+                tail_chunks = (total_frames + tail) // chunk_size - total_frames // chunk_size - 1
+                for _ in range(max(0, tail_chunks)):
+                    pos_start = len(input_ids) + 1
+                    audio_pos.append((pos_start, pos_start + chunk_size))
+                    chunk_ids, chunk_labels = _silence_chunk(chunk_size)
+                    input_ids += chunk_ids
+                    labels    += chunk_labels
+
+                fout.write(json.dumps({
+                    "tasks":       "online",
+                    "idx":         rec["idx"],
+                    "input_ids":   input_ids,
+                    "labels":      labels,
+                    "audio_pos":   audio_pos,
+                    "pt_path_dir": rec["pt_path_dir"],
+                }, ensure_ascii=False) + "\n")
+            except Exception as e:
+                print(f"[step4 idx {rec.get('idx')}] {type(e).__name__}: {e}")
+
+
+# === Main driver ===
+
+def main():
+    for name, value in [("CHECKPOINT_DIR", CHECKPOINT_DIR), ("INPUT_JSONL", INPUT_JSONL),
+                        ("WORK_DIR", WORK_DIR), ("OUT_TRAIN_JSONL", OUT_TRAIN_JSONL),
+                        ("NOISE_DIR", NOISE_DIR)]:
+        if not value:
+            raise SystemExit(f"Set {name} at the top of cons_online_data.py before running.")
+
+    tokenizer_dir, _, qwen_omni_ckpt, audio_tower_ckpt = resolve_checkpoint_paths(CHECKPOINT_DIR)
+    step1_jsonl = os.path.join(WORK_DIR, "step1.jsonl")
+    step2_jsonl = os.path.join(WORK_DIR, "step2.jsonl")
+    step3_jsonl = os.path.join(WORK_DIR, "step3.jsonl")
+    wavs_dir    = os.path.join(WORK_DIR, "wavs")
+    features_dir = os.path.join(WORK_DIR, "features")
+
+    step1_sample_silence(
+        INPUT_JSONL, step1_jsonl,
+        noise_dir=NOISE_DIR,
+        min_noise_len=MIN_NOISE_LEN, max_noise_len=MAX_NOISE_LEN,
+        chunk_size=CHUNK_SIZE, seed=SEED,
+    )
+    step2_concat_audio(
+        step1_jsonl, step2_jsonl, wavs_dir,
+        chunk_size=CHUNK_SIZE, seed=SEED,
+    )
+    step3_extract_features(
+        step2_jsonl, step3_jsonl, features_dir,
+        qwen_omni_ckpt=qwen_omni_ckpt, audio_tower_ckpt=audio_tower_ckpt, device=DEVICE,
+    )
+    step4_build_tokens(
+        step3_jsonl, OUT_TRAIN_JSONL,
+        tokenizer_dir=tokenizer_dir, chunk_size=CHUNK_SIZE,
+    )
+
+
+if __name__ == "__main__":
+    main()

src/audiointeraction/generate/base.py

@@ -0,0 +1,413 @@
+"""Core streaming generation primitives for the audio-enhanced GPT.
+
+The model alternates between two states inside `streaming_generate`:
+  - LISTENING: each step consumes one encoder-output chunk of audio. The model
+    emits either KEEP_SILENCE (keep listening) or TEXT_BEGIN (start replying).
+  - SPEAKING:  autoregressive text generation until TEXT_END, then back to
+    LISTENING for the next audio chunk.
+
+Public surface:
+    AUDIO_TOKENS_PER_CHUNK
+    sample, encode_audio_chunks, streaming_generate
+"""
+SYSTEM_PROMPT = (
+    "You are a helpful assistant. When there is no user text, if the audio contains a question, "
+    "please answer it. If it is a sound effect, determine based on the sound whether help is needed."
+)
+
+
+
+import os
+import sys
+from typing import List, Optional
+
+import numpy as np
+import torch
+import whisper
+
+from src.audiointeraction.dataset.TOKENS import (
+    ASSISTANT, AUDIO_BEGIN, KEEP_SILENCE, PAD, TEXT_BEGIN, TEXT_END,
+    HAPPY, SAD, ANGRY, SURPRISE, NORMAL, URGENT,
+)
+from src.audiointeraction.model import GPT
+from src.audiointeraction.tokenizer import Tokenizer
+
+
+# Encoder-output frames per [AUDIO_BEGIN, PAD*N, ASSISTANT, ...] block.
+AUDIO_TOKENS_PER_CHUNK = 10
+
+
+# 情绪 token -> 颜文字（终端/GitHub 都能正常显示）
+EMOTION_KAOMOJI = {
+    HAPPY:    "(◕‿◕)",
+    SAD:      "(╥﹏╥)",
+    ANGRY:    "(╬ಠ益ಠ)",
+    SURPRISE: "(⊙o⊙)",
+    NORMAL:   "(・_・)",
+    URGENT:   "(°□°;)",
+}
+
+
+# === Sampling ===
+
+def _top_p_filter(logits: torch.Tensor, top_p: float) -> torch.Tensor:
+    sorted_logits, sorted_idx = torch.sort(logits, descending=False)
+    cumprobs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
+    remove = cumprobs <= (1 - top_p)
+    remove[-1:] = 0  # always keep the most probable token
+    return logits.masked_fill(remove.scatter(0, sorted_idx, remove), float("-inf"))
+
+
+def sample(logits: torch.Tensor, *, temperature=1.0, top_k=None, top_p=1.0) -> torch.Tensor:
+    """Sample one token id from the last position of `logits` ([1, T, V])."""
+    if not 0.0 <= top_p <= 1.0:
+        raise ValueError(f"top_p must be in [0, 1], got {top_p}")
+    logits = logits[0, -1]
+    if top_k is not None:
+        v, i = torch.topk(logits, min(top_k, logits.size(-1)))
+        logits = torch.full_like(logits, float("-inf")).scatter_(-1, i, v)
+    if temperature <= 0.0 and top_p <= 0.0:
+        return torch.argmax(logits, dim=-1, keepdim=True)
+    if temperature > 0.0:
+        logits = logits / temperature
+    if top_p < 1.0:
+        logits = _top_p_filter(logits, top_p)
+    probs = torch.nn.functional.softmax(logits, dim=-1)
+    return torch.multinomial(probs, num_samples=1)
+
+
+# === Audio feature extraction ===
+
+def _split_into_chunks(n: int, chunk_size: int) -> List[int]:
+    chunks = [chunk_size] * (n // chunk_size)
+    if n % chunk_size:
+        chunks.append(n % chunk_size)
+    return chunks
+
+
+def _encode_audio_samples(audio: List[float], audio_encoder: torch.nn.Module, device) -> List[torch.Tensor]:
+    """Run the audio_tower on raw 16 kHz samples and split the output into AUDIO_TOKENS_PER_CHUNK chunks."""
+    audio = list(audio)
+    # Pad to a 0.4-s boundary (6400 samples @ 16 kHz).
+    if len(audio) % 6400 != 0:
+        audio += [0] * (6400 - len(audio) % 6400)
+    mel = whisper.log_mel_spectrogram(np.array(audio, dtype=np.float32), n_mels=128)
+    len_feature = mel.shape[1]
+
+    with torch.no_grad():
+        feat = audio_encoder(
+            torch.tensor(mel).to(device),
+            torch.tensor(_split_into_chunks(len_feature, 40)).to(device),
+            torch.tensor((len_feature - 1) // 2 + 1).to(device),
+        ).last_hidden_state
+
+    # Drop any trailing partial chunk so each chunk is exactly AUDIO_TOKENS_PER_CHUNK frames.
+    keep = feat.shape[0] - feat.shape[0] % AUDIO_TOKENS_PER_CHUNK
+    return [feat[i: i + AUDIO_TOKENS_PER_CHUNK] for i in range(0, keep, AUDIO_TOKENS_PER_CHUNK)]
+
+
+def encode_audio_chunks(audio_path: str, audio_encoder: torch.nn.Module, device) -> List[torch.Tensor]:
+    """Run the audio_tower on `audio_path` and split the output into AUDIO_TOKENS_PER_CHUNK chunks."""
+    audio = whisper.load_audio(audio_path, sr=16000).tolist()
+    return _encode_audio_samples(audio, audio_encoder, device)
+
+
+def encode_silence_chunks(seconds: float, audio_encoder: torch.nn.Module, device) -> List[torch.Tensor]:
+    return _encode_audio_samples([0.0] * int(seconds * 16000), audio_encoder, device)
+
+
+# === Streaming generation ===
+
+def _forward(model, tokens, input_pos, *, input_pos_maxp1, audio_feat):
+    return model(
+        tokens, None, 1, audio_feat, input_pos,
+        input_pos_maxp1=input_pos_maxp1,
+        audio_tokens_per_chunk=AUDIO_TOKENS_PER_CHUNK,
+    )
+
+
+def _init_input_pos_maxp1(model, prompt_size, device):
+    # input_pos_maxp1 introduces data-dependent shapes; skip if a Thunder module is involved.
+    if any(m.__class__.__name__ == "ThunderModule" for m in model.modules()):
+        return None
+    return torch.tensor(prompt_size, device=device)
+
+
+def _append_listening_block(token, input_pos, input_pos_maxp1, device):
+    """Append `[AUDIO_BEGIN, PAD*N, ASSISTANT]` to the running context."""
+    new_tokens = torch.LongTensor([AUDIO_BEGIN] + [PAD] * AUDIO_TOKENS_PER_CHUNK + [ASSISTANT]).to(device)
+    new_positions = input_pos[-1] + torch.arange(1, len(new_tokens) + 1, device=device)
+    token = torch.cat([token, new_tokens])
+    input_pos = torch.cat([input_pos, new_positions])
+    if input_pos_maxp1 is not None:
+        input_pos_maxp1.add_(len(new_tokens))
+    return token, input_pos, input_pos_maxp1
+
+
+def _advance_one(input_pos, input_pos_maxp1):
+    """Move input_pos forward by one (for the next single-token call)."""
+    new_pos = input_pos[-1].unsqueeze(0).add_(1)
+    if input_pos_maxp1 is not None:
+        input_pos_maxp1.add_(1)
+    return new_pos, input_pos_maxp1
+
+
+# === Pretty printing helpers ===
+
+class _Pretty:
+    """Tiny renderer for the streaming UI. Auto-disables color on non-TTY."""
+
+    # ANSI codes; emptied below when not a TTY.
+    DIM = "\033[2m"
+    BOLD = "\033[1m"
+    RESET = "\033[0m"
+    CYAN = "\033[36m"
+    GREEN = "\033[32m"
+    YELLOW = "\033[33m"
+    MAGENTA = "\033[35m"
+    GREY = "\033[90m"
+    CLEAR_LINE = "\033[2K\r"
+
+    BAR_WIDTH = 24
+
+    def __init__(self, stream=sys.stdout):
+        self.stream = stream
+        self.use_color = stream.isatty()
+        self.use_cr = stream.isatty()  # only do in-place updates on a real terminal
+        if not self.use_color:
+            for name in ("DIM", "BOLD", "RESET", "CYAN", "GREEN",
+                         "YELLOW", "MAGENTA", "GREY"):
+                setattr(self, name, "")
+            self.CLEAR_LINE = ""
+        # Track whether we're currently sitting on a transient status line
+        # (so the next permanent write knows to clear it first).
+        self._status_active = False
+
+    # --- low-level ---
+    def _write(self, s: str) -> None:
+        self.stream.write(s)
+        self.stream.flush()
+
+    def _clear_status(self) -> None:
+        if self._status_active and self.use_cr:
+            self._write(self.CLEAR_LINE)
+        self._status_active = False
+
+    # --- public ---
+    def header(self, round_idx: int, n_rounds: int, audio_path: str, n_chunks: int) -> None:
+        self._clear_status()
+        name = os.path.basename(audio_path)
+        bar = "━" * 60
+        self._write(
+            f"\n{self.CYAN}{bar}{self.RESET}\n"
+            f"{self.BOLD}{self.CYAN}▶ Round {round_idx}/{n_rounds}{self.RESET}  "
+            f"{self.DIM}{name}{self.RESET}  "
+            f"{self.GREY}[{n_chunks} chunks]{self.RESET}\n"
+            f"{self.CYAN}{bar}{self.RESET}\n"
+        )
+
+    def status(self, chunk_idx: int, n_chunks: int, silent_run: int, replied: int) -> None:
+        """Transient one-line progress; updates in place on a TTY."""
+        if not self.use_cr:
+            return  # don't spam a logfile with thousands of status lines
+        filled = int(self.BAR_WIDTH * chunk_idx / max(n_chunks, 1))
+        bar = (
+            f"{self.GREEN}{'█' * filled}{self.RESET}"
+            f"{self.GREY}{'·' * (self.BAR_WIDTH - filled)}{self.RESET}"
+        )
+        dots = f"{self.DIM}{'·' * min(silent_run, 40)}{self.RESET}" if silent_run else ""
+        replied_tag = (
+            f"  {self.GREEN}✓ {replied} reply{'ies' if replied != 1 else ''}{self.RESET}"
+            if replied else ""
+        )
+        line = (
+            f"{self.CLEAR_LINE}"
+            f"{self.DIM}listening{self.RESET} {bar} "
+            f"{self.BOLD}{chunk_idx:>3}/{n_chunks}{self.RESET}"
+            f"{replied_tag}  {dots}"
+        )
+        self._write(line)
+        self._status_active = True
+
+    def reply_begin(self) -> None:
+        """Promote to a permanent line: clear status, drop to a new line, print prefix."""
+        self._clear_status()
+        self._write(f"  {self.MAGENTA}▸{self.RESET} ")
+
+    def reply_token(self, text: str) -> None:
+        self._write(text)
+
+    def reply_emotion(self, token_id: int) -> None:
+        """Print the emotion kaomoji inline at the start of a reply."""
+        kao = EMOTION_KAOMOJI.get(token_id, f"[emo:{token_id}]")
+        self._write(f"{self.YELLOW}{kao}{self.RESET} ")
+
+    def reply_end(self) -> None:
+        self._write("\n")
+
+    def round_summary(self, replied: int, silent: int, total: int) -> None:
+        self._clear_status()
+        self._write(
+            f"  {self.GREY}└─ {replied} reply chunk(s), "
+            f"{silent} silent, {total} total{self.RESET}\n"
+        )
+
+    def finish(self) -> None:
+        self._clear_status()
+
+
+def streaming_generate(
+    model: GPT,
+    audio_encoder: torch.nn.Module,
+    tokenizer: Tokenizer,
+    prefix_ids: torch.Tensor,
+    *,
+    rounds: int = 10,
+    audio_paths: Optional[List[str]] = None,
+    max_returned_tokens: int = 4096,
+    temperature: float = 0,
+    top_k: Optional[int] = 1,
+    top_p: float = 1.0,
+):
+    """Stream audio→text. If `audio_paths` is given, run one round per path
+    non-interactively (offline); otherwise prompt stdin each round (online)."""
+    device = prefix_ids.device
+    ui = _Pretty()
+    ui._write(f"{ui.DIM}device: {device}{ui.RESET}\n")
+
+    token = prefix_ids
+    input_pos = torch.arange(0, prefix_ids.size(0), device=device, dtype=torch.int64)
+    input_pos_maxp1 = _init_input_pos_maxp1(model, prefix_ids.size(0), device)
+
+    turns: List[List[int]] = []  # one inner list per assistant turn (across all rounds)
+
+    if audio_paths is not None:
+        steps = []
+        for ap in audio_paths:
+            steps.append((ap, False))
+            steps.append((None, True))
+        n_real = len(audio_paths)
+    else:
+        steps = None
+        n_real = rounds
+
+    n_steps = len(steps) if steps is not None else rounds
+    real_idx = 0
+    acc_replied = 0
+    acc_silent = 0
+
+    for step_idx in range(n_steps):
+        if steps is not None:
+            audio_path, is_silence = steps[step_idx]
+        else:
+            is_silence = False
+            ui._clear_status()
+            audio_path = input(
+                f"{ui.BOLD}Round {real_idx + 1}/{n_real}{ui.RESET} — enter audio path: "
+            ).strip()
+
+        if is_silence:
+            audio_chunks = encode_silence_chunks(1.0, audio_encoder, device)
+        else:
+            audio_chunks = encode_audio_chunks(audio_path, audio_encoder, device)
+            real_idx += 1
+            acc_replied = 0
+            acc_silent = 0
+            ui.header(real_idx, n_real, audio_path, len(audio_chunks))
+
+        # Per-round counters for the progress line / summary.
+        replied_chunks = 0
+        silent_chunks = 0
+        silent_run = 0  # consecutive silent chunks since the last reply (for the dots)
+
+        # Buffer of decoded ids for the *current* assistant turn — emitted token-by-token.
+        # turn[0] is TEXT_BEGIN. turn[1] is *usually* an emotion tag, but if not we
+        # treat it as regular text.
+        current_turn: List[int] = []
+        text_started = False  # have we opened a reply line yet?
+
+        listening, audio_idx = True, -1
+        ui.status(0, len(audio_chunks), silent_run, replied_chunks)
+
+        for _ in range(max_returned_tokens - input_pos.numel()):
+            if listening:
+                audio_idx += 1
+                if audio_idx >= len(audio_chunks):
+                    break
+                token, input_pos, input_pos_maxp1 = _append_listening_block(
+                    token, input_pos, input_pos_maxp1, device
+                )
+                logits = _forward(
+                    model, token.view(1, -1), input_pos,
+                    input_pos_maxp1=input_pos_maxp1,
+                    audio_feat=audio_chunks[audio_idx].to(device),
+                )
+            else:
+                logits = _forward(
+                    model, token.view(1, -1), input_pos,
+                    input_pos_maxp1=input_pos_maxp1,
+                    audio_feat=None,
+                )
+
+            token = sample(logits, temperature=temperature, top_k=top_k, top_p=top_p).to(torch.int64)
+            int_token = token.item()
+            input_pos, input_pos_maxp1 = _advance_one(input_pos, input_pos_maxp1)
+
+            if listening:
+                if int_token == TEXT_BEGIN:
+                    listening = False
+                    current_turn = [int_token]
+                    text_started = False
+                    # Don't print yet — wait until we know whether the next token is an emotion.
+                elif int_token == KEEP_SILENCE:
+                    turns.append([int_token])
+                    silent_chunks += 1
+                    silent_run += 1
+                    ui.status(audio_idx + 1, len(audio_chunks), silent_run, replied_chunks)
+                else:
+                    raise ValueError(f"Unexpected token {int_token} while listening")
+            else:
+                current_turn.append(int_token)
+
+                if int_token == TEXT_END:
+                    if text_started:
+                        ui.reply_end()
+                    turns.append(current_turn)
+                    current_turn = []
+                    text_started = False
+                    replied_chunks += 1
+                    silent_run = 0
+                    listening = True
+                    ui.status(audio_idx + 1, len(audio_chunks), silent_run, replied_chunks)
+                else:
+                    # Layout: [TEXT_BEGIN, (optional EMOTION), t0, t1, ..., TEXT_END]
+                    # current_turn already has TEXT_BEGIN + this token in it.
+                    n = len(current_turn)
+                    if n == 2:
+                        # First token after TEXT_BEGIN — could be an emotion tag or
+                        # just regular text if the model skipped the emotion.
+                        ui.reply_begin()
+                        text_started = True
+                        if int_token in EMOTION_KAOMOJI:
+                            ui.reply_emotion(int_token)
+                        else:
+                            piece = tokenizer.decode(torch.tensor([int_token]))
+                            if piece:
+                                ui.reply_token(piece)
+                    elif n >= 3:
+                        # Stream this token's surface form immediately.
+                        piece = tokenizer.decode(torch.tensor([int_token]))
+                        if piece:
+                            ui.reply_token(piece)
+
+        if not listening and text_started:
+            ui.reply_end()
+            text_started = False
+
+        acc_replied += replied_chunks
+        acc_silent += silent_chunks
+        if steps is None or is_silence:
+            ui.round_summary(acc_replied, acc_silent, acc_replied + acc_silent)
+
+    ui.finish()
+    return turns

src/audiointeraction/model.py

@@ -0,0 +1,972 @@
+# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
+
+"""Full definition of a decoder-only transformer-based language model, all of it in this single file.
+
+Based on the nanoGPT implementation: https://github.com/karpathy/nanoGPT and
+https://github.com/EleutherAI/gpt-neox/tree/main/megatron/model.
+"""
+
+import math
+from typing import Any, Optional, Tuple, Union, List
+from functools import partial
+from transformers import AutoConfig, Qwen2_5OmniForConditionalGeneration
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing_extensions import Self
+import whisper
+from transformers import Qwen2AudioEncoder, Qwen2AudioConfig
+from src.audiointeraction.config import Config
+
+
+def qkv_reassemble(
+    param: torch.Tensor, config: Config
+) -> torch.Tensor:
+    """Reassemble from a normal to an interleaved placement in a QKV matrix.
+    [Q, K, V, Q, K, V, ...] --> [Q, Q, ..., K, K, ..., V, V, ...]
+    """
+    q_per_kv = config.n_head // config.n_query_groups
+    qs = []
+    ks = []
+    vs = []
+    for chunk in torch.chunk(param, config.n_query_groups):
+        split = torch.split(chunk, [config.head_size * q_per_kv, config.head_size, config.head_size])
+        qs.append(split[0])
+        ks.append(split[1])
+        vs.append(split[2])
+    q = torch.cat(qs)
+    k = torch.cat(ks)
+    v = torch.cat(vs)
+    return torch.cat((q, k, v))
+
+
+class GPT(nn.Module):
+    def __init__(self, config: Config) -> None:
+        super().__init__()
+        assert config.padded_vocab_size is not None
+        self.config = config
+    
+        self.lm_head = nn.Linear(
+            config.n_embd, config.padded_vocab_size, bias=config.lm_head_bias
+        )
+
+        self.transformer = nn.ModuleDict(
+            dict(
+                wte=nn.Embedding(config.padded_vocab_size, config.n_embd),
+                h=nn.ModuleList(
+                    Block(config, block_idx)
+                    for block_idx in range(config.n_layer)
+                ),
+                ln_f=config.norm_class(config.n_embd, eps=config.norm_eps),
+            )
+        )
+        self.mask_cache: Optional[torch.Tensor] = None
+        self.max_seq_length = self.config.block_size
+
+    @property
+    def max_seq_length(self) -> int:
+        return self._max_seq_length
+
+    @max_seq_length.setter
+    def max_seq_length(self, value: int) -> None:
+        """
+        When doing inference, the sequences used might be shorter than the model's context length.
+        This allows setting a smaller number to avoid allocating unused memory
+        """
+        if value > self.config.block_size:
+            raise ValueError(
+                f"Cannot attend to {value}, block size is only {self.config.block_size}."
+                " This is likely because the input text exceeds the supported context length of this model."
+            )
+        self._max_seq_length = value
+        if not hasattr(self, "cos"):
+            # first call
+            cos, sin = self.rope_cache()
+            self.register_buffer("cos", cos, persistent=False)
+            self.register_buffer("sin", sin, persistent=False)
+        # override
+        elif value != self.cos.size(0):
+            self.cos, self.sin = self.rope_cache(device=self.cos.device)
+        # the mask and kv cache size will get updated on `set_kv_cache`. we cannot update it here because we don't know
+        # if the kv cache is expected
+        if self.mask_cache is not None and self.mask_cache.shape[-1] < value:
+            print(f"Warning: KV cache has length {self.mask_cache.shape[-1]} < {value} = max_seq_length. Call 'set_kv_cache' before doing any forwards!")
+
+    def reset_parameters(self) -> None:
+        # Trigger resetting the rope-cache
+        self.cos, self.sin = self.rope_cache(device=self.cos.device)
+
+    def _init_weights(self, module: nn.Module) -> None:
+        """Meant to be used with `gpt.apply(gpt._init_weights)`."""
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+
+    def fill_in_audio_feature(self,
+                                input_embeddings: torch.Tensor,
+                                batch_size: int,
+                                audio_feats_list,
+                                audio_pos,
+                                tasks) -> torch.Tensor:
+        """Replace AUDIO_PAD positions in input_embeddings with precomputed audio features.
+
+        Two fill modes, dispatched per-sample by `tasks[batch_idx]`:
+
+          - "online":  audio is streamed in fixed 10-frame chunks. `audio_pos[i]`
+                       is a list of (start, end) tuples (each end-start == 10);
+                       slice the feature tensor by 10 per chunk.
+          - "offline": audio is one contiguous block. `audio_pos[i]` is a single
+                       (start, end) tuple covering `output_len` positions; place
+                       the entire feature tensor in one shot.
+        """
+        _, _, emb_dim = input_embeddings.shape
+
+        if not (batch_size == len(audio_feats_list) == len(audio_pos) == len(tasks)):
+            raise ValueError(
+                f"length mismatch: batch_size={batch_size}, "
+                f"feats={len(audio_feats_list)}, pos={len(audio_pos)}, tasks={len(tasks)}"
+            )
+
+        for batch_idx in range(batch_size):
+            audio_feats = audio_feats_list[batch_idx]
+            segments = audio_pos[batch_idx]
+            if segments is None or segments == -1:
+                continue
+
+            task = tasks[batch_idx]
+            if task == "offline":
+                # Single big block: place the whole feature tensor at the one segment.
+                start, end = segments[0]
+                if start >= self.max_seq_length:
+                    continue
+                end = min(end, self.max_seq_length)
+                input_embeddings[batch_idx, start:end, :] = audio_feats[: end - start]
+                continue
+
+            # Online: per-10-frame chunk placement.
+            for seg_idx, (start, end) in enumerate(segments):
+                if start > self.max_seq_length:
+                    continue
+                if end > self.max_seq_length:
+                    input_embeddings[batch_idx, start:self.max_seq_length, :] = torch.zeros(
+                        self.max_seq_length - start, emb_dim
+                    )
+                else:
+                    audio_feat = audio_feats[seg_idx * 10 : (seg_idx + 1) * 10]
+                    seg_len, feat_dim = audio_feat.shape
+                    expected_len = end - start
+                    if seg_len != expected_len or feat_dim != emb_dim:
+                        raise ValueError(
+                            f"Loaded feature shape {audio_feat.shape} does not match expected "
+                            f"({expected_len}, {emb_dim}) at batch {batch_idx}, segment {seg_idx}")
+                    # Overwrite the embedding segment
+                    input_embeddings[batch_idx, start:end, :] = audio_feat
+
+        return input_embeddings
+
+    def forward(
+        self,
+        idx: torch.Tensor,
+        tasks: Optional[List[str]],
+        batch_size: int,
+        audio_info: Optional[Union[dict, torch.Tensor]] = None,
+        input_pos: Optional[torch.Tensor] = None,
+        input_pos_maxp1: Optional[torch.Tensor] = None,
+        audio_tokens_per_chunk: int = 10,
+        lm_head_chunk_size: int = 0,
+    ) -> Union[torch.Tensor, List[torch.Tensor]]:
+        """
+        If `input_pos` is provided, the KV cache uses K and V vectors for
+        positions smaller than entries in `input_pos`. For efficiency, pass
+        `input_pos_maxp1` as `max(input_pos) + 1` if already available from
+        your forward algorithm. This slices the KV cache buffers and speeds
+        up multi-head attention.
+
+        Without `input_pos_maxp1`, the computation uses the full KV cache
+        (`max_seq_length`) with masking applied. Note that inferring
+        `input_pos_maxp1` from `input_pos` causes graph breaks and prevents
+        compilation.
+
+        Args:
+            idx: Token indices of input sequences, shape `(B, T)`, where `B`
+                is batch size.
+            input_pos: Optional. Positions of input tokens. The default is
+                `arange(T)`. Can have shape `(T,)` or `(B, T)` (batched index).
+            input_pos_maxp1: Optional. See above.
+            lm_head_chunk_size: Optional. If `lm_head_chunk_size > 0`, the final
+                `lm_head` computation is done in chunks of this size.
+
+        Returns:
+            Logit outputs, shape `(B, T, config.padded_vocab_size)`. If
+            `lm_head_chunk_size > 0`, this is a list of chunks of shape
+            `(B, lm_head_chunk_size, config.padded_vocab_size)`, the final
+            entry can be shorter.
+
+        """
+        T = idx.size(1)
+        if self.max_seq_length < T:
+            raise ValueError(f"Cannot forward sequence of length {T}, max seq length is only {self.max_seq_length}.")
+
+        if input_pos is not None:  # use the kv cache
+            if input_pos.dim() > 2:
+                # otherwise, things go wrong in `apply_rope`
+                raise ValueError(f"input_pos must have 1 or 2 dimensions, input_pos.shape = {input_pos.shape}")
+            if input_pos.shape[-1] != T:
+                raise ValueError(f"input_pos.shape[-1] = {input_pos.shape[-1]} != {T} = idx.shape[1], must be the same")
+            cos = batched_index_select(self.cos, 0, input_pos)
+            sin = batched_index_select(self.sin, 0, input_pos)
+            if input_pos.dim() == 1:
+                cos = cos.unsqueeze(0)
+                sin = sin.unsqueeze(0)
+            if self.mask_cache is None:
+                raise TypeError("You need to call `gpt.set_kv_cache()`")
+            mask = batched_index_select(self.mask_cache, 2, input_pos)
+            if mask.dim() > 4:
+                # the mask cache has a batch dim of 1 in addition to the one
+                # we get if input_pos has a batch dimension
+                mask = mask.view(*(mask.shape[0:1] + mask.shape[2:]))
+            if input_pos_maxp1 is not None:
+                # Shorten final dimension so it just covers all `input_pos` entries
+                if input_pos_maxp1 > self.max_seq_length:
+                    raise ValueError(f"Positions in 'input_pos' must be in [0,{self.max_seq_length})")
+                mask = mask[..., :input_pos_maxp1]
+        else:
+            # unsqueeze to have a batch dimension
+            cos = self.cos[:T].unsqueeze(0)
+            sin = self.sin[:T].unsqueeze(0)
+            # `cos`, `sin` have shape (1, T, config.rope_n_elem)
+            mask = None  # defaults to causal mask
+            input_pos_maxp1 = None
+
+        x = self.transformer.wte(idx)  # token embeddings of shape (B, T, n_embd)
+
+        # Audio feature injection — dispatch on input type. Encoder features are
+        # already n_embd-dim (projected by audio_tower.proj), so we place them
+        # directly into the input embeddings.
+        #   - dict   : training path, segment-based fill from precomputed features
+        #   - Tensor : inference path, streaming chunk replacement
+        #   - None   : no audio (e.g. text-only data or inter-token decoding step)
+        if isinstance(audio_info, dict):
+            # T_T (text-only) samples have audio_pos == None — nothing to fill.
+            if audio_info.get("audio_pos") is not None:
+                x = self.fill_in_audio_feature(
+                    x, batch_size, audio_info["feats_paths"], audio_info["audio_pos"], tasks,
+                )
+        elif torch.is_tensor(audio_info):
+            if T > audio_tokens_per_chunk:
+                if x.size(0) != 1:
+                    raise ValueError("inference mode, it is not supported for batch size > 1")
+                x[0, T - (audio_tokens_per_chunk + 1): T - 1, :] = audio_info
+
+        if self.config.scale_embeddings:
+            x = x * torch.tensor(self.config.n_embd ** 0.5, dtype=x.dtype)
+
+        for block in self.transformer.h:
+            x = block(x, cos, sin, mask, input_pos, input_pos_maxp1)
+
+        x = self.transformer.ln_f(x)
+        clamp_head = (
+            partial(do_softcapping, thresh=self.config.final_logit_softcapping)
+            if self.config.final_logit_softcapping is not None
+            else nn.Identity()
+        )
+        if lm_head_chunk_size > 0:
+            # chunk the lm head logits to reduce the peak memory used by autograd
+            return [
+                clamp_head(self.lm_head(x_i))
+                for x_i in x.split(lm_head_chunk_size, dim=1)
+            ]
+        else:
+            return clamp_head(self.lm_head(x))  # (B, T, padded_vocab_size)
+
+    def rope_cache(self, device: Optional[torch.device] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        if self.config.rope_adjustments is None:
+            extra_config = None
+
+        else:
+            adjusted_params_required = ["factor", "low_freq_factor", "high_freq_factor", "original_max_seq_len"]
+            params_present = [param in self.config.rope_adjustments for param in adjusted_params_required]
+            num_params_present = sum(params_present)
+
+            if num_params_present == 0:
+                extra_config = None  # uses standard RoPE
+            elif num_params_present == 4:
+                # These parameters should always be used together so that we don't interfere with standard rope
+                extra_config = {
+                    name: self.config.rope_adjustments[name]
+                    for name in adjusted_params_required
+                }
+            else:
+                # Some but not all parameters are specified; raise an error
+                missing_params = [
+                    param for param, present in zip(adjusted_params_required, params_present) if not present
+                ]
+                raise ValueError(
+                    f"The following adjusted RoPE parameters are missing in rope_adjustments: {', '.join(missing_params)}. "
+                    "All adjusted RoPE parameters must be specified together."
+                )
+
+        return build_rope_cache(
+            seq_len=self.max_seq_length,
+            n_elem=self.config.rope_n_elem,
+            device=device,
+            condense_ratio=self.config.rope_condense_ratio,
+            base=self.config.rope_base,
+            extra_config=extra_config,
+        )
+
+    def set_kv_cache(
+        self,
+        batch_size: int,
+        max_seq_length: Optional[int] = None,
+        rope_cache_length: Optional[int] = None,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> None:
+        if rope_cache_length is None:
+            rope_cache_length = self.cos.size(-1)
+
+        if max_seq_length is None:
+            max_seq_length = self.max_seq_length
+
+        # initialize the kv cache for all blocks
+        for block in self.transformer.h:
+            block.attn.kv_cache = block.attn.build_kv_cache(
+                batch_size,
+                max_seq_length,
+                rope_cache_length,
+                device,
+                dtype,
+            )
+
+        if self.mask_cache is None or self.mask_cache.size(3) != max_seq_length:
+            # passing `attn_mask` to SDPA disables the flash implementation. since we only need the mask
+            # for the kv-cache support (only during inference), we only create it in that situation
+            self.mask_cache = build_mask_cache(max_seq_length, device)
+
+    def clear_kv_cache(self) -> None:
+        self.mask_cache = None
+        for block in self.transformer.h:
+            block.attn.kv_cache = None
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        config: Config,
+        block_idx: int,
+    ) -> None:
+        super().__init__()
+        if not config.parallel_residual and config.shared_attention_norm:
+            raise NotImplementedError(
+                "No checkpoint amongst the ones we support uses this configuration"
+                " (non-parallel residual and shared attention norm)."
+            )
+
+        self.norm_1 = config.norm_class(config.n_embd, eps=config.norm_eps)
+        self.attn = CausalSelfAttention(config, block_idx)
+        self.post_attention_norm = (
+            config.norm_class(config.n_embd, eps=config.norm_eps) if config.post_attention_norm else nn.Identity()
+        )
+        self.norm_2 = None if config.shared_attention_norm else config.norm_class(config.n_embd, eps=config.norm_eps)
+        self.mlp = config.mlp_class(config)
+        self.post_mlp_norm = (
+            config.norm_class(config.n_embd, eps=config.norm_eps) if config.post_mlp_norm else nn.Identity()
+        )
+
+        self.config = config
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        input_pos: Optional[torch.Tensor] = None,
+        input_pos_maxp1: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Non-parallel residual       Parallel residual
+           ┌─ x                     ┌─ x ──────────────────┐             Note: if `shared_attention_norm` is True,
+           │  ↓                     │  ↓                   ↓                   the output from `norm_1` is reused
+           │  norm_1                │  norm_1  ───────►    norm_2
+           │  ↓                     │  ↓                   ↓
+           │  attn                  │  attn                MLP
+           │  ↓                     │  ↓                   ↓
+           |  post_attn_norm        |  post_attn_norm      post_mlp_norm
+           |  ↓                     |  ↓                   ↓
+        ┌─ └► +                     └► + ◄─────────────────┘
+        |     ↓
+        │     norm_2
+        │     ↓
+        │     MLP
+        │     ↓
+        |     post_mlp_norm
+        |     ↓
+        └───► +
+        """
+
+        x_normed = self.norm_1(x)
+        attention_output = self.attn(
+            x_normed, cos, sin, mask, input_pos, input_pos_maxp1
+        )
+        attention_output = self.post_attention_norm(attention_output)
+
+        if self.config.parallel_residual:
+            if not self.config.shared_attention_norm:
+                x_normed = self.norm_2(x)
+            x = attention_output + x
+        else:
+            x = attention_output + x
+            x_normed = self.norm_2(x)
+        return self.post_mlp_norm(self.mlp(x_normed)) + x
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config: Config, block_idx: int) -> None:
+        super().__init__()
+        # key, query and value projections for all heads, but in a batch
+        self.qkv = nn.Linear(
+            config.n_embd,
+            (config.n_head + 2 * config.n_query_groups) * config.head_size,  # support for grouped/multi queries
+            bias=config.bias or config.attn_bias,
+        )
+        # output projection
+        self.proj = nn.Linear(
+            config.head_size * config.n_head, config.n_embd, bias=config.bias
+        )
+        # disabled by default
+        self.kv_cache: Optional[KVCache] = None
+        self.apply_sliding_window_attention = (
+            config.sliding_window_size is not None and
+            block_idx % config.sliding_window_layer_stride == 0
+        )
+
+        if config.norm_qk:
+            self.norm_q = config.norm_class(config.head_size * config.n_head, eps=config.norm_eps)
+            self.norm_k = config.norm_class(config.head_size * config.n_query_groups, eps=config.norm_eps)
+        else:
+            self.norm_q = self.norm_k = None
+
+        self.config = config
+        self.block_idx = block_idx
+
+        # Attention capture flags (for analysis/visualization, disabled by default)
+        self.capture_attn: bool = False
+        self.captured_attn_weights: Optional[torch.Tensor] = None  # shape: (B, n_head, T_q, T_k)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        input_pos: Optional[torch.Tensor] = None,
+        input_pos_maxp1: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        # Notation:
+        # - B          | batch size
+        # - T          | time-step (sequence length)
+        # - C          | model's embeddings size (n_embd)
+        # - C*         | attentions's embeddings size
+        # - nh_(q,k,v) | number of heads for query, key and value
+        # - hs         | head size
+        head_size = self.config.head_size
+        n_head = self.config.n_head
+        n_query_groups = self.config.n_query_groups
+        rope_n_elem = self.config.rope_n_elem
+        B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)
+
+        # Perform a single multiplication operation using a combined QKV matrix to calculate `query`, `key`, and `value`
+        # instead of individually multiplying the input `x` with the respective weight matrices.
+        qkv = self.qkv(x)  # (B, T, 3xC*)
+
+        # Define query, key and value sizes.
+        # If grouped/multi query is enabled, these sizes are not equal (see the diagram in `lit_gpt/config.py::Config`).
+        query_size = n_head * head_size
+        key_size = value_size = n_query_groups * head_size
+        # Split qkv into query, key and value matrices.
+        q, k, v = qkv.split((query_size, key_size, value_size), dim=-1)  # 3x(B, T, C*)
+
+        if self.config.norm_qk:
+            q = self.norm_q(q)
+            k = self.norm_k(k)
+
+        # To place the num_heads (nh) dimension right after the batch (B) dimension, the first step is to decouple the
+        # embedding size (C) into num_heads (nh) and head_size (hs).
+        q = q.view(B, T, n_head, head_size)  # (B, T, nh_q, hs)
+        k = k.view(B, T, n_query_groups, head_size)  # (B, T, nh_k, hs)
+        v = v.view(B, T, n_query_groups, head_size)  # (B, T, nh_v, hs)
+
+        # The tensors `query`, `key`, and `value` are now accurately structured: within each batch element (B), there are
+        # multiple heads (nh), and within each head, there is a sequence of elements (T), each represented by a vector
+        # of size `hs`.
+        q = q.transpose(1, 2)  # (B, nh_q, T, hs)
+        k = k.transpose(1, 2)  # (B, nh_k, T, hs)
+        v = v.transpose(1, 2)  # (B, nh_v, T, hs)
+
+        # Unlike standard positional embeddings rotary embeddings must be applied at every layer.
+        q_roped = apply_rope(q[..., : rope_n_elem], cos, sin)
+        k_roped = apply_rope(k[..., : rope_n_elem], cos, sin)
+        q = torch.cat((q_roped, q[..., rope_n_elem :]), dim=-1)  # (B, nh_q, T, hs)
+        k = torch.cat((k_roped, k[..., rope_n_elem :]), dim=-1)  # (B, nh_k, T, hs)
+
+        # Apply kv-cache during inference.
+        if input_pos is not None:
+            if not isinstance(self.kv_cache, KVCache):
+                raise TypeError("You need to call `gpt.set_kv_cache()`")
+            k, v = self.kv_cache(input_pos, k, v)
+            if input_pos_maxp1 is not None:
+                # Subselect along sequence dimension
+                k = k[..., :input_pos_maxp1, :]
+                v = v[..., :input_pos_maxp1, :]
+            # k, v: (B, nh_k, input_pos_maxp1, hs)
+            # If input_pos_maxp1 is None -> max_seq_length
+
+        use_flash = (getattr(self.config, "use_flash_attention", True)
+                     and mask is None
+                     and n_query_groups == n_head
+                     )
+        if use_flash:
+            # FlashAttention: B H T D -> B T H D
+            q = q.transpose(1, 2).contiguous()  # (B, T, nh_q, hs)
+            k = k.transpose(1, 2).contiguous()  # (B, T, nh_k, hs)
+            v = v.transpose(1, 2).contiguous()  # (B, T, nh_v, hs)
+            from flash_attn.flash_attn_interface import flash_attn_func
+            y = flash_attn_func(q, k, v, dropout_p=0.0, causal=True)
+            y = y.transpose(1, 2) # back to B H T D
+
+        else:
+            # Grouped queries: balance the number of heads across all three matrices.
+            # NOTE: flash attention requires it in training mode.
+            # Multi-query: this step can be skipped since there is only 1 head, allowing us to use broadcasting.
+            if n_query_groups != n_head and (input_pos is None or n_query_groups != 1):
+                q_per_kv = n_head // n_query_groups
+                k = k.repeat_interleave(q_per_kv, dim=1)  # (B, nh_q, T, hs)
+                v = v.repeat_interleave(q_per_kv, dim=1)  # (B, nh_q, T, hs)
+
+            if self.apply_sliding_window_attention:
+                """
+                    Global Window              Sliding window             Sliding window
+                    attention mask      +            bias          =      attention mask
+                ┌────────────────────────┐  ┌───────────────────────┐  ┌─────────────────────────┐
+                │ True False False False │  │ True  True  True True │  │ True  False False False │
+                │ True True  False False │  │ True  True  True True │  │ True  True  False False │
+                │ True True  True  False │  │ False True  True True │  │ False True  True  False │
+                │ True True  True  True  │  │ False False True True │  │ False False True  True  │
+                └────────────────────────┘  └───────────────────────┘  └─────────────────────────┘
+                """
+                if mask is None:
+                    mask = torch.ones(T, T, dtype=q.dtype, device=q.device).triu(diagonal=1)
+                    mask.masked_fill_(mask.bool(), float("-inf"))
+                    mask = mask.view(1, 1, *mask.shape)
+                sliding_window_bias = torch.ones_like(mask).tril(diagonal=-self.config.sliding_window_size)
+                sliding_window_bias.masked_fill_(sliding_window_bias.bool(), float("-inf"))
+                mask += sliding_window_bias
+
+            # Efficient attention using Flash Attention CUDA kernels.
+            # NOTE: efficient implementation is disabled if `mask` is not None or softcapping is enabled.
+            # ↓ (B, nh, T, hs) @ (B, nh, T, hs).mT --> (B, nh, T, T) @ (B, nh, T, hs) --> (B, nh, T, hs)
+            y = self.scaled_dot_product_attention(q, k, v, mask)
+
+        # Re-assemble all head outputs side by side.
+        y = y.reshape(B, T, head_size * n_head)
+
+        # Output projection.
+        return self.proj(y)  # (B, T, C)
+
+    def scaled_dot_product_attention(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        scale = 1.0 / math.sqrt(self.config.attention_scores_scalar or self.config.head_size)
+
+        # with softcapping we cannot use SDPA
+        if self.config.attention_logit_softcapping is not None:
+            scores = q @ k.mT * scale
+            scores = do_softcapping(scores, self.config.attention_logit_softcapping)
+            if mask is None:
+                mask = torch.ones(q.size(2), q.size(2), dtype=q.dtype, device=q.device).triu(diagonal=1)
+                mask.masked_fill_(mask.bool(), torch.finfo(q.dtype).min)
+            scores = scores + mask
+            scores = F.softmax(scores, dim=-1, dtype=torch.float).to(dtype=q.dtype)
+            if self.capture_attn:
+                self.captured_attn_weights = scores.detach()
+            y = scores @ v
+        elif self.capture_attn:
+            # Manual attention computation to capture weights (bypasses fused SDPA kernel)
+            # q: (B, n_head, T_q, hs), k: (B, n_head, T_k, hs)
+            scores = torch.matmul(q.float(), k.float().transpose(-2, -1)) * scale
+            if mask is not None:
+                if mask.dtype == torch.bool:
+                    scores = scores.masked_fill(~mask, float('-inf'))
+                else:
+                    scores = scores + mask.float()
+            else:
+                # Apply causal mask manually when no mask is provided (training mode)
+                T_q, T_k = q.size(-2), k.size(-2)
+                causal = torch.ones(T_q, T_k, device=q.device, dtype=torch.bool).tril(diagonal=T_k - T_q)
+                scores = scores.masked_fill(~causal, float('-inf'))
+            attn_weights = F.softmax(scores, dim=-1)
+            self.captured_attn_weights = attn_weights.detach()
+            y = torch.matmul(attn_weights.to(dtype=v.dtype), v)
+            return y.transpose(1, 2)
+        else:
+            y = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=mask, dropout_p=0.0, scale=scale, is_causal=mask is None
+            )
+        return y.transpose(1, 2)
+
+    def build_kv_cache(
+        self,
+        batch_size: int,
+        max_seq_length: int,
+        rope_cache_length: Optional[int] = None,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> "KVCache":
+        v_shape = (batch_size, self.config.n_query_groups, max_seq_length, self.config.head_size)
+        if rope_cache_length is None:
+            if self.config.rotary_percentage != 1.0:
+                raise TypeError("Please pass the `rope_cache_length=gpt.cos.size(-1)` value")
+            k_shape = v_shape
+        else:
+            k_shape = (
+                batch_size,
+                self.config.n_query_groups,
+                max_seq_length,
+                rope_cache_length + self.config.head_size - self.config.rope_n_elem,
+            )
+        return KVCache(k_shape, v_shape, device=device, dtype=dtype)
+
+    def _load_from_state_dict(self, state_dict: dict, prefix: str, *args: Any, **kwargs: Any) -> None:
+        """For compatibility with legacy checkpoints."""
+
+        for attr in ("weight", "bias"):
+            legacy_key = f"{prefix}attn.{attr}"
+            current_key = f"{prefix}qkv.{attr}"
+            if legacy_key in state_dict:
+                state_dict[current_key] = qkv_reassemble(state_dict.pop(legacy_key), self.config)
+
+        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
+
+class GptNeoxMLP(nn.Module):
+    def __init__(self, config: Config) -> None:
+        super().__init__()
+        self.fc = nn.Linear(
+            config.n_embd, config.intermediate_size, bias=config.bias
+        )
+        self.proj = nn.Linear(
+            config.intermediate_size, config.n_embd, bias=config.bias
+        )
+        self.config = config
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc(x)
+        x = F.gelu(x, approximate=self.config.gelu_approximate)
+        return self.proj(x)
+
+
+class LLaMAMLP(nn.Module):
+    def __init__(self, config: Config) -> None:
+        super().__init__()
+        self.fc_1 = nn.Linear(
+            config.n_embd, config.intermediate_size, bias=config.bias
+        )
+        self.fc_2 = nn.Linear(
+            config.n_embd, config.intermediate_size, bias=config.bias
+        )
+        self.proj = nn.Linear(
+            config.intermediate_size, config.n_embd, bias=config.bias
+        )
+        self.config = config
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_fc_1 = self.fc_1(x)
+        x_fc_2 = self.fc_2(x)
+        x = F.silu(x_fc_1) * x_fc_2
+        return self.proj(x)
+
+
+class GemmaMLP(LLaMAMLP):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_fc_1 = self.fc_1(x)
+        x_fc_2 = self.fc_2(x)
+        x = F.gelu(x_fc_1, approximate=self.config.gelu_approximate) * x_fc_2
+        return self.proj(x)
+
+
+class LLaMAMoE(nn.Module):
+    def __init__(self, config: Config) -> None:
+        super().__init__()
+        self.gate = nn.Linear(config.n_embd, config.n_expert, bias=False)
+        self.experts = nn.ModuleList(LLaMAMLP(config) for _ in range(config.n_expert))
+        self.config = config
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Derived from: https://github.com/mistralai/mistral-src/blob/b46d6/moe_one_file_ref.py#L203-L219
+        See also figure 1 in https://arxiv.org/abs/2211.15841
+        """
+        B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)
+        x = x.view(-1, C)  # (B*T, C)
+        router = self.gate(x)  # (B*T, n_expert)
+        probs, indices = torch.topk(router, self.config.n_expert_per_token)  # (B*T, n_expert_per_token)
+        probs = probs.softmax(dim=1, dtype=torch.float).to(dtype=x.dtype)
+        masks = indices.unsqueeze(-1) == torch.arange(self.config.n_expert, device=x.device)
+        masks = masks.permute(2, 0, 1)  # (n_expert, B*T, n_expert_per_token)
+        y = torch.zeros_like(x)  # (B*T, C)
+        for mask, expert in zip(masks, self.experts):
+            token_idx, expert_idx = torch.where(mask)
+            y[token_idx] += probs[token_idx, expert_idx, None] * expert(x[token_idx])
+        return y.view(B, T, C)
+
+
+def build_rope_cache(
+    seq_len: int,
+    n_elem: int,
+    device: Optional[torch.device] = None,
+    base: int = 10000,
+    condense_ratio: int = 1,
+    extra_config: Optional[dict] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Enhanced Transformer with Rotary Position Embedding.
+
+    Args:
+        seq_len (int): Sequence length.
+        n_elem (int): Number of elements (head dimension).
+        device (torch.device, optional): Device for tensor allocations.
+        base (int, optional): Base for computing inverse frequencies.
+        condense_ratio (int, optional): Ratio to condense the position indices.
+        extra_config (dict, optional): Configuration parameters for frequency adjustments (used by Llama 3.1 and 3.2)
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Cosine and sine caches for RoPE.
+            Shapes are `(seq_len, n_elem)`.
+    """
+
+    # Compute the inverse frequencies theta
+    theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, device=device).float() / n_elem))
+
+    if extra_config is not None:
+        orig_context_len = extra_config["original_max_seq_len"]
+        factor = extra_config["factor"]
+        low_freq_factor = extra_config["low_freq_factor"]
+        high_freq_factor = extra_config["high_freq_factor"]
+
+        wavelen = 2 * torch.pi / theta
+        ratio = orig_context_len / wavelen
+        smooth_factor = (ratio - low_freq_factor) / (high_freq_factor - low_freq_factor)
+        smooth_factor = torch.clamp(smooth_factor, min=0.0, max=1.0)
+
+        # Compute adjusted_theta without masked indexing
+        adjusted_theta = (1 - smooth_factor) * (theta / factor) + smooth_factor * theta
+        theta = adjusted_theta
+
+    # Create position indices `[0, 1, ..., seq_len - 1]`
+    ### Zhifei fix bug 1:.
+    seq_idx = torch.arange(seq_len, device=device, dtype=torch.float16) / float(condense_ratio)
+    # seq_idx = torch.arange(seq_len, device=device) / condense_ratio
+
+    # Calculate the product of position index and $\theta_i$
+    idx_theta = torch.outer(seq_idx, theta).repeat(1, 2)
+    # If `n_elem` is odd, the final dimension of `idx_theta` has size
+    # `n_elem + 1`, so need to cut something off.
+    # Due to a current bug in Hugging Face, in the case `n_elem == 1`, we leave
+    # `idx_theta`, `cos`, `sin` as is. Things work out in `apply_rope` due to
+    # broadcasting. If we shorten `idx_theta`, unit tests comparing to
+    # Hugging Face fail.
+    # https://github.com/huggingface/transformers/issues/35233
+    if idx_theta.shape[-1] > n_elem > 1:
+        idx_theta = idx_theta[..., :n_elem]
+
+    return torch.cos(idx_theta), torch.sin(idx_theta)
+
+
+def batched_index_select(t, dim, idx):
+    """index_select for batched index and unbatched t"""
+    if idx.dim() == 1:
+        return torch.index_select(t, dim, idx)
+
+    *batch_shape, idx_size = idx.shape
+    res = torch.index_select(t, dim, idx.reshape(-1))  # flat index
+    # split out single batch idx
+    res = res.view(*t.shape[:dim], -1, idx_size, *t.shape[dim + 1 :])
+    if dim > 0:
+        # move batch dim to front, this is np.rollaxis(res, dim, 0) for tensors
+        dims = [dim] + list(range(res.dim()))
+        del dims[dim + 1]
+        res = res.permute(dims)
+    # unflatten batch dims
+    res = res.view(*batch_shape, *res.shape[1:])
+    return res
+
+
+def batched_index_copy_(t, dim, idx, val):
+    """Index copy for batched t, idx, val"""
+
+    if t.device.type == "mps":
+        # Normalize negative dimensions
+        if dim < 0:
+            dim = t.dim() + dim
+        if idx.dim() == 1:
+            idx_shape = [1] * val.dim()
+            idx_shape[dim] = -1
+            idx_expanded = idx.view(*idx_shape)
+            idx_expanded = idx_expanded.expand_as(val)
+            t.scatter_(dim, idx_expanded, val)
+            return t
+
+        elif idx.dim() == 2:
+            assert dim != 0, "Cannot index the batch dimension"
+            batch_size = idx.size(0)
+            idx_size = idx.size(1)
+            assert batch_size == t.size(0) == val.size(0)
+
+            idx_shape = [batch_size] + [1] * (val.dim() - 1)
+            idx_shape[dim] = idx_size
+            idx_expanded = idx.view(*idx_shape)
+            idx_expanded = idx_expanded.expand_as(val)
+
+            t.scatter_(dim, idx_expanded, val)
+            return t
+        else:
+            raise NotImplementedError(f"idx.dim() == {idx.dim()} not supported")
+
+    else:
+        if idx.dim() == 1:
+            return t.index_copy_(dim, idx, val)
+
+        assert idx.dim() == 2, f"multiple batch dims not yet {idx.shape=}"
+        assert dim != 0, f"cannot index batch dim {dim=}"
+        batch_size, idx_size = idx.shape
+        assert batch_size == t.size(0)
+        assert batch_size == val.size(0)
+
+        # if we can view the batch and indexed dimensions together, we could
+        # do index trickery. This is, sadly, not the case for kvcache so we
+        # fall back to for loop
+        for i in range(batch_size):
+            unbatched_dim = dim if dim < 0 else dim - 1
+            t[i].index_copy_(unbatched_dim, idx[i], val[i])
+        return t
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """
+    Applies RoPE transform to `x`. Note that `cos`, `sin` need to have a batch
+    dimension.
+
+    Args:
+        x: Input tensor, `(B, ..., T, head_size)`
+        cos: Cached cosines, `(B, T, head_size)` or `(1, T, head_size)`
+        sin: Cached sines, `(B, T, head_size)` or `(1, T, head_size)`
+
+    Returns:
+        Encoded tensor, `(B, ..., T, head_size)`
+    """
+    if cos.dim() != 3:
+        raise ValueError(f"cos must be three-dimensional, but shape is {cos.shape}")
+    if cos.shape != sin.shape:
+        raise ValueError(f"cos, sin must have same shape, but cos.shape={cos.shape}, sin.shape={sin.shape}")
+    head_size_half = x.size(-1) // 2
+    x1 = x[..., : head_size_half]  # (B, ..., T, head_size/2)
+    x2 = x[..., head_size_half :]  # (B, ..., T, head_size/2)
+    rotated = torch.cat((-x2, x1), dim=-1)  # (B, ..., T, head_size)
+    dims_diff = x.dim() - cos.dim()
+    if dims_diff > 0:
+        # Ensure that shapes of `x`, `cos`, `sin` align
+        new_shape = cos.shape[0:1] + (1,) * dims_diff + cos.shape[1:]
+        cos = cos.view(*new_shape)
+        sin = sin.view(*new_shape)
+
+    roped = (x * cos) + (rotated * sin)
+    return roped.to(dtype=x.dtype)
+
+
+def do_softcapping(x: torch.Tensor, thresh: float) -> torch.Tensor:
+    return torch.tanh(x / thresh) * thresh
+
+
+class KVCache(nn.Module):
+    """
+    Buffers `k`, `v` have shape
+    `(batch_size, n_query_groups, max_seq_length, head_size)`.
+    """
+    def __init__(
+        self,
+        k_shape: Tuple[int, int, int, int],
+        v_shape: Tuple[int, int, int, int],
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> None:
+        super().__init__()
+        self.register_buffer("k", torch.zeros(k_shape, device=device, dtype=dtype), persistent=False)
+        self.register_buffer("v", torch.zeros(v_shape, device=device, dtype=dtype), persistent=False)
+
+    def forward(self, input_pos: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Writes new values `k` and `v` into the cache at the positions specified
+        by `input_pos` along the sequence dimension (`max_seq_length`). The batch
+        size of `k` and `v` (`bs`) must be smaller or equal to `KVCache` batch
+        size. Returns the full buffers, adjusted to the batch size `bs`.
+
+        Args:
+            input_pos: Position index, `(bs, T)` or `(T,)`
+            k: New values, `(bs, n_query_groups, T, head_size)`
+            v: New values, `(bs, n_query_groups, T, head_size)`
+
+        Returns:
+            k_full, v_full, `(bs, n_query_groups, max_seq_length, head_size)`
+
+        """
+        # move the buffer to the activation dtype for when AMP is used
+        self.k = self.k.to(k.dtype)
+        self.v = self.v.to(v.dtype)
+        # update the cache
+        bs = k.size(0)
+        k = batched_index_copy_(self.k[:bs, ...], -2, input_pos, k)
+        v = batched_index_copy_(self.v[:bs, ...], -2, input_pos, v)
+        return k, v
+
+    def reset_parameters(self) -> None:
+        torch.nn.init.zeros_(self.k)
+        torch.nn.init.zeros_(self.v)
+
+
+def build_mask_cache(max_seq_length: int, device: Optional[torch.device] = None) -> torch.Tensor:
+    ones = torch.ones((max_seq_length, max_seq_length), device=device, dtype=torch.bool)
+    return torch.tril(ones).unsqueeze(0).unsqueeze(0)
+
+
+class RMSNorm(torch.nn.Module):
+    """Root Mean Square Layer Normalization.
+
+    Derived from https://github.com/bzhangGo/rmsnorm/blob/master/rmsnorm_torch.py. BSD 3-Clause License:
+    https://github.com/bzhangGo/rmsnorm/blob/master/LICENSE.
+    """
+
+    def __init__(self, size: int, dim: int = -1, eps: float = 1e-6, add_unit_offset: bool = False) -> None:
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(size))
+        self.eps = eps
+        self.dim = dim
+        self.add_unit_offset = add_unit_offset
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        x = x.float()
+        # NOTE: the original RMSNorm paper implementation is not equivalent
+        norm_x = torch.mean(x * x, dim=self.dim, keepdim=True)
+        x_normed = x * torch.rsqrt(norm_x + self.eps)
+        weight = (1 + self.weight) if self.add_unit_offset else self.weight
+        return (x_normed * weight.float()).to(dtype=dtype)
+
+    def reset_parameters(self) -> None:
+        torch.nn.init.ones_(self.weight)

src/audiointeraction/tokenizer.py

@@ -0,0 +1,179 @@
+# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
+
+import json
+from pathlib import Path
+from typing import Optional, Union, Iterable, Iterator
+
+from src.audiointeraction.utils import fix_and_load_json
+import torch
+
+
+class Tokenizer:
+    def __init__(self, checkpoint_dir: Union[Path, str]) -> None:
+        checkpoint_dir = Path(checkpoint_dir)
+        if not checkpoint_dir.exists():
+            raise NotADirectoryError(f"The checkpoint directory does not exist: {str(checkpoint_dir)}")
+
+        self.model_name = checkpoint_dir.stem
+        self.use_bos = self.check_if_bos_token_used(checkpoint_dir)
+        self.bos_id = None
+        self.eos_id = None
+
+        # some checkpoints have both files, `.json` takes precedence
+        if (vocabulary_path := checkpoint_dir / "tokenizer.json").is_file():
+            from tokenizers import Tokenizer as HFTokenizer
+
+            self.processor = HFTokenizer.from_file(str(vocabulary_path))
+            self.backend = "huggingface"
+
+            if (special_tokens_path := checkpoint_dir / "tokenizer_config.json").is_file():
+                with open(special_tokens_path, encoding="utf-8") as fp:
+                    config = json.load(fp)
+                bos_token = config.get("bos_token")
+                eos_token = config.get("eos_token")
+                if bos_token is not None and isinstance(bos_token, dict):
+                    bos_token = bos_token.get("content")
+                if eos_token is not None and isinstance(eos_token, dict):
+                    eos_token = eos_token.get("content")
+                self.bos_id = self.token_to_id(bos_token) if bos_token is not None else None
+                self.eos_id = self.token_to_id(eos_token) if eos_token is not None else None
+            if (special_tokens_path := checkpoint_dir / "generation_config.json").is_file():
+                try:
+                    with open(special_tokens_path, encoding="utf-8") as fp:
+                        config = json.load(fp)
+                except json.JSONDecodeError:  # Some files like the Llama 3.2 one have bugs
+                    with open(special_tokens_path, encoding="utf-8") as fp:
+                        json_string = fp.read()
+                        config = fix_and_load_json(json_string)
+                if self.bos_id is None:
+                    self.bos_id = config.get("bos_token_id")
+                if self.eos_id is None:
+                    self.eos_id = config.get("eos_token_id")
+
+        elif (vocabulary_path := checkpoint_dir / "tokenizer.model").is_file():
+            from sentencepiece import SentencePieceProcessor
+
+            self.processor = SentencePieceProcessor(model_file=str(vocabulary_path))
+            self.backend = "sentencepiece"
+            self.bos_id = self.processor.bos_id()
+            self.eos_id = self.processor.eos_id()
+        else:
+            raise NotImplementedError
+
+        # NOTE: A temporary fix until it's resolved on Tokenizers side.
+        # LlaMA tokenizer strips leading spaces if to decode a single token at a time.
+        # https://github.com/huggingface/transformers/issues/31643
+        self.apply_decoding_fix = None
+        if (config_path := checkpoint_dir / "tokenizer_config.json").is_file():
+            with open(config_path, encoding="utf-8") as fp:
+                self.apply_decoding_fix = "LlamaTokenizer" in json.load(fp)["tokenizer_class"]
+
+    @property
+    def vocab_size(self) -> int:
+        if self.backend == "huggingface":
+            return self.processor.get_vocab_size(with_added_tokens=False)
+        if self.backend == "sentencepiece":
+            return self.processor.vocab_size()
+        raise RuntimeError
+
+    def token_to_id(self, token: str) -> int:
+        if self.backend == "huggingface":
+            id_ = self.processor.token_to_id(token)
+        elif self.backend == "sentencepiece":
+            id_ = self.processor.piece_to_id(token)
+        else:
+            raise RuntimeError
+        if id_ is None:
+            raise ValueError(f"token {token!r} not found in the collection.")
+        return id_
+
+    def check_if_bos_token_used(self, checkpoint_dir: Path) -> bool:
+        if not (tokenizer_config_path := checkpoint_dir / "tokenizer_config.json").is_file():
+            return False
+        with open(tokenizer_config_path, encoding="utf-8") as fp:
+            config = json.load(fp)
+        # for LlaMA-3 tokenizer there is no `add_bos_token` at all and `tokenizer_class` is only
+        # `PreTrainedTokenizerFast`
+        if checkpoint_dir.stem.startswith(("Meta-Llama-3", "Llama-3")):
+            return True
+        if checkpoint_dir.stem.startswith("SmolLM2") and checkpoint_dir.name.endswith("Instruct"):
+            return True
+        if "add_bos_token" in config:
+            return config["add_bos_token"]
+        # if `add_bos_token` isn't in the config file, but LLaMA tokenizer is used - return True.
+        # ex: https://huggingface.co/stabilityai/StableBeluga2/blob/main/tokenizer_config.json#L2
+        return config.get("tokenizer_class") == "LlamaTokenizer"
+
+    def encode(
+        self,
+        string: str,
+        device: Optional[torch.device] = None,
+        bos: Optional[bool] = None,
+        eos: bool = False,
+        max_length: int = -1,
+    ) -> torch.Tensor:
+        if self.backend == "huggingface":
+            tokens = self.processor.encode(string).ids
+        elif self.backend == "sentencepiece":
+            tokens = self.processor.encode(string)
+        else:
+            raise RuntimeError(f"`{self.backend}` is not supported.")
+        if tokens is None:
+            raise ValueError("`self.processor` returned tokens of None value.")
+
+        if bos or (bos is None and self.use_bos):
+            if self.bos_id is None:
+                raise NotImplementedError("This tokenizer does not have a defined bos token.")
+            if not tokens or tokens[0] != self.bos_id:
+                tokens = [self.bos_id] + tokens
+        # if the processor misbehaves and adds `bos` token no matter what
+        elif tokens and tokens[0] == self.bos_id:
+            tokens = tokens[1:]
+
+        if eos and (not tokens or tokens[-1] != self.eos_id):
+            tokens = tokens + [self.eos_id]
+        # if the processor misbehaves and adds `eos` token no matter what
+        elif tokens and tokens[-1] == self.eos_id:
+            tokens = tokens[:-1]
+
+        if max_length > 0:
+            tokens = tokens[:max_length]
+        return torch.tensor(tokens, dtype=torch.int, device=device)
+
+    def decode(self, tensor: torch.Tensor) -> str:
+        tokens = [tensor.item()] if tensor.ndim == 0 else tensor.tolist()
+        if len(tokens) == 1 and self.apply_decoding_fix:
+            dummy_token_id = 33  # \x1e
+            dummy_token = self.processor.decode([dummy_token_id])
+            if dummy_token != "\x1e":
+                dummy_token_id = 165 # \x1e is different in salamandra tokenizers
+                dummy_token = self.processor.decode([dummy_token_id])
+            return self.processor.decode([dummy_token_id] + tokens)[len(dummy_token) :]
+        return self.processor.decode(tokens)
+
+    def decode_stream(self, token_stream: Iterable[torch.Tensor], device: Optional[torch.device] = None) -> Iterator[str]:
+        if self.backend == "huggingface":
+            try:
+                for token in token_stream:
+                    yield self.decode(token)
+            except KeyboardInterrupt:
+                return
+        elif self.backend == "sentencepiece":
+            # TODO: Is there a way to not have to do this?
+            # This may actually affect our tokens per second.
+
+            # sentencepiece does not support decoding token-by-token because it adds spaces based on the surrounding tokens
+            # meaning that we need to decode everything each time
+            so_far = torch.tensor([], dtype=torch.long, device=device)
+            decoded_so_far = ""
+            try:
+                for token in token_stream:
+                    so_far = so_far.to(device=token.device)
+                    so_far = torch.cat((so_far, token.view(-1)))
+                    decoded_new = self.decode(so_far)
+                    yield decoded_new[len(decoded_so_far) :]
+                    decoded_so_far = decoded_new
+            except KeyboardInterrupt:
+                return
+        else:
+            raise NotImplementedError(self.backend)

src/audiointeraction/utils.py

@@ -0,0 +1,290 @@
+"""Training / inference helpers."""
+
+import inspect
+import json
+import os
+import re
+import shutil
+import subprocess
+from pathlib import Path
+from typing import Any, Iterable, List, Literal, Optional, Union
+
+import lightning as L
+import torch
+import torch.nn as nn
+from lightning.fabric.loggers import CSVLogger, TensorBoardLogger
+from lightning.fabric.strategies import FSDPStrategy
+from lightning.fabric.utilities.load import _lazy_load as lazy_load
+from lightning.pytorch.cli import instantiate_class
+from lightning.pytorch.loggers import WandbLogger
+from typing_extensions import Self
+
+
+# === Output dir / checkpoint resume ===
+
+def init_out_dir(out_dir: Path) -> Path:
+    if not isinstance(out_dir, Path):
+        out_dir = Path(out_dir)
+    if not out_dir.is_absolute() and "LIGHTNING_ARTIFACTS_DIR" in os.environ:
+        return Path(os.getenv("LIGHTNING_ARTIFACTS_DIR")) / out_dir
+    return out_dir
+
+
+def find_resume_path(resume: Union[bool, Literal["auto"], Path], out_dir: Path) -> Optional[Path]:
+    if not resume or isinstance(resume, Path):
+        return resume or None
+    # `resume` is True or "auto"; find the newest lit_model.pth under out_dir.
+    candidates = list(out_dir.rglob("lit_model.pth"))
+    if not candidates:
+        if resume == "auto":
+            return None
+        raise FileNotFoundError(
+            f"resume=True but no `lit_model.pth` under {out_dir}; set resume='auto' to skip or False to start fresh."
+        )
+    return max(candidates, key=lambda p: p.stat().st_mtime)
+
+
+# === Misc math / introspection ===
+
+def find_multiple(n: int, k: int) -> int:
+    if n % k == 0:
+        return n
+    return n + k - (n % k)
+
+
+def num_parameters(module: nn.Module, requires_grad: Optional[bool] = None) -> int:
+    total = 0
+    for p in module.parameters():
+        if requires_grad is None or p.requires_grad == requires_grad:
+            total += p.numel()
+    return total
+
+
+# === Cross-entropy ===
+
+def chunked_cross_entropy(
+    logits: Union[torch.Tensor, List[torch.Tensor]],
+    targets: torch.Tensor,
+    chunk_size: int = 128,
+    ignore_index: int = -100,
+) -> torch.Tensor:
+    """Memory-friendly cross-entropy: optionally chunk along the seq axis to bound the activation peak."""
+    if isinstance(logits, list):
+        if chunk_size == 0:
+            logits = torch.cat(logits, dim=1).reshape(-1, logits[0].size(-1))
+            return torch.nn.functional.cross_entropy(logits, targets.reshape(-1), ignore_index=ignore_index)
+        logit_chunks = [c.reshape(-1, c.size(-1)) for c in logits]
+        target_chunks = [c.reshape(-1) for c in targets.split(logits[0].size(1), dim=1)]
+        loss_chunks = [
+            torch.nn.functional.cross_entropy(lc, tc, ignore_index=ignore_index, reduction="none")
+            for lc, tc in zip(logit_chunks, target_chunks)
+        ]
+        non_masked = (targets != ignore_index).sum()
+        return torch.cat(loss_chunks).sum() / non_masked.maximum(torch.ones_like(non_masked))
+
+    logits = logits.reshape(-1, logits.size(-1))
+    targets = targets.reshape(-1)
+    if chunk_size == 0:
+        return torch.nn.functional.cross_entropy(logits, targets, ignore_index=ignore_index)
+    logit_chunks = logits.split(chunk_size)
+    target_chunks = targets.split(chunk_size)
+    loss_chunks = [
+        torch.nn.functional.cross_entropy(lc, tc, ignore_index=ignore_index, reduction="none")
+        for lc, tc in zip(logit_chunks, target_chunks)
+    ]
+    non_masked = (targets != ignore_index).sum()
+    return torch.cat(loss_chunks).sum() / non_masked.maximum(torch.ones_like(non_masked))
+
+
+# === Precision / checkpoint IO ===
+
+def get_default_supported_precision(training: bool) -> str:
+    """Return bf16 if the GPU supports it, otherwise 16-bit. Picks `-mixed` for training, `-true` for inference."""
+    if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
+        return "bf16-mixed" if training else "bf16-true"
+    return "16-mixed" if training else "16-true"
+
+
+def load_checkpoint(fabric: L.Fabric, model: nn.Module, checkpoint_path: Path, strict: bool = True) -> None:
+    if isinstance(fabric.strategy, FSDPStrategy):
+        fabric.load_raw(checkpoint_path, model, strict=strict)
+    else:
+        state_dict = lazy_load(checkpoint_path)
+        state_dict = state_dict.get("model", state_dict)
+        model.load_state_dict(state_dict, strict=strict)
+
+
+# === Iteration ===
+
+class CycleIterator:
+    """Indefinitely cycles through an iterable, exposing the current epoch count."""
+
+    def __init__(self, iterable: Iterable) -> None:
+        self.iterable = iterable
+        self.epoch = 0
+        self._iterator = None
+
+    def __next__(self) -> Any:
+        if self._iterator is None:
+            self._iterator = iter(self.iterable)
+        try:
+            return next(self._iterator)
+        except StopIteration:
+            self._iterator = iter(self.iterable)
+            self.epoch += 1
+            return next(self._iterator)
+
+    def __iter__(self) -> Self:
+        return self
+
+
+# === Setup helpers (out dir, logger, devices, optimizer) ===
+
+def copy_config_files(source_dir: Path, out_dir: Path) -> None:
+    """Copy model_config.yaml and tokenizer files (when present) into the output directory."""
+    for name in ("config.json", "generation_config.json", "model_config.yaml",
+                 "tokenizer.json", "tokenizer.model", "tokenizer_config.json"):
+        src = source_dir / name
+        if src.exists():
+            shutil.copy(src, out_dir)
+
+
+def parse_devices(devices: Union[str, int]) -> int:
+    if devices in (-1, "auto"):
+        return torch.cuda.device_count() or 1
+    if isinstance(devices, int) and devices > 0:
+        return devices
+    raise ValueError(f"Devices must be 'auto' or a positive integer, got: {devices!r}")
+
+
+def choose_logger(
+    logger_name: Literal["csv", "tensorboard", "wandb"],
+    out_dir: Path,
+    name: str,
+    log_interval: int = 1,
+    resume: Optional[bool] = None,
+    **kwargs: Any,
+):
+    if logger_name == "csv":
+        return CSVLogger(root_dir=(out_dir / "logs"), name="csv", flush_logs_every_n_steps=log_interval, **kwargs)
+    if logger_name == "tensorboard":
+        return TensorBoardLogger(root_dir=(out_dir / "logs"), name="tensorboard", **kwargs)
+    if logger_name == "wandb":
+        return WandbLogger(project=name, resume=resume, **kwargs)
+    raise ValueError(f"`--logger_name={logger_name}` is not a valid option. Choose from 'csv', 'tensorboard', 'wandb'.")
+
+
+def instantiate_torch_optimizer(optimizer, model_parameters, **kwargs):
+    """Resolve `optimizer` (str name like 'AdamW' or a dict with class_path/init_args) to an actual optimizer."""
+    if isinstance(optimizer, str):
+        class_module, class_name = optimizer.rsplit(".", 1) if "." in optimizer else ("torch.optim", optimizer)
+        module = __import__(class_module, fromlist=[class_name])
+        optimizer_cls = getattr(module, class_name)
+        valid = set(inspect.signature(optimizer_cls).parameters)
+        kwargs = {k: v for k, v in kwargs.items() if k in valid}
+        return optimizer_cls(model_parameters, **kwargs)
+    if isinstance(optimizer, dict):
+        optimizer = dict(optimizer)
+        class_module, class_name = optimizer["class_path"].rsplit(".", 1)
+        module = __import__(class_module, fromlist=[class_name])
+        valid = set(inspect.signature(getattr(module, class_name)).parameters)
+        optimizer["init_args"].update({k: v for k, v in kwargs.items() if k in valid})
+        return instantiate_class(model_parameters, optimizer)
+    raise ValueError(f'Unrecognized "optimizer" value: {optimizer}')
+
+
+# === GPU topology (nvlink / xgmi) ===
+
+def check_nvlink_connectivity(fabric=None):
+    """Best-effort check that the multi-GPU interconnect is fast. Prints a warning if not."""
+    custom_print = fabric.print if fabric is not None else print
+    if os.getenv("RANK", "0") != "0":
+        return
+    try:
+        if not torch.cuda.is_available():
+            custom_print("No GPUs available")
+            return
+        gpu_name = torch.cuda.get_device_properties(0).name.lower()
+        if "nvidia" in gpu_name:
+            _check_nvidia_connectivity(custom_print)
+        elif "advanced micro devices" in gpu_name or "amd" in gpu_name:
+            _check_amd_connectivity(custom_print)
+        else:
+            custom_print(f"Unrecognized GPU vendor: {gpu_name}")
+    except Exception as e:
+        custom_print(f"An error occurred while checking GPU connectivity: {e}")
+
+
+def _check_nvidia_connectivity(custom_print):
+    result = subprocess.run(["nvidia-smi", "topo", "-m"], stdout=subprocess.PIPE, text=True)
+    if result.returncode != 0:
+        custom_print("Failed to run nvidia-smi"); return
+    lines = result.stdout.strip().split("\n")
+    start = next((i for i, line in enumerate(lines) if "GPU0" in line), None)
+    if start is None:
+        custom_print("Failed to parse nvidia-smi output"); return
+    headers = lines[start].split()
+    gpu_count = len([h for h in headers if re.match(r"^GPU\d+$", h)])
+    all_nvlink = all(
+        all("NV" in conn for conn in row.split()[1:1 + gpu_count] if conn != "X")
+        for row in lines[start + 1: start + 1 + gpu_count]
+    )
+    if all_nvlink:
+        custom_print("All GPUs are fully connected via NVLink.")
+    else:
+        custom_print("Warning: not all GPUs are fully connected via NVLink — multi-GPU throughput may be suboptimal.")
+
+
+def _check_amd_connectivity(custom_print):
+    result = subprocess.run(["rocm-smi", "--showtopotype"], stdout=subprocess.PIPE, text=True)
+    if result.returncode != 0:
+        custom_print("Failed to run rocm-smi"); return
+    lines = result.stdout.strip().split("\n")
+    gpu_idx = next((i for i, line in enumerate(lines) if re.match(r"^\s*GPU0", line)), None)
+    if gpu_idx is None or gpu_idx == 0:
+        custom_print("Failed to parse rocm-smi output (no GPU headers found)"); return
+    headers = lines[gpu_idx - 1].strip().split()
+    gpu_count = len([h for h in headers if re.match(r"^GPU\d+$", h)])
+    gpu_lines = [l.strip() for l in lines[gpu_idx: gpu_idx + gpu_count] if re.match(r"^\s*GPU\d+", l)]
+    if len(gpu_lines) != gpu_count:
+        custom_print("Mismatch in GPU count when parsing rocm-smi output"); return
+    all_xgmi = all(
+        all(conn in ("XGMI", "0") for conn in row.split()[1:1 + gpu_count])
+        for row in gpu_lines
+    )
+    if all_xgmi:
+        custom_print("All GPUs are fully connected via XGMI.")
+    else:
+        custom_print("Warning: not all GPUs are fully connected via XGMI — multi-GPU throughput may be suboptimal.")
+
+
+# === Misc ===
+
+def fix_and_load_json(s: str):
+    """Tolerant JSON parser: strips trailing commas and inserts missing ones before retrying."""
+    s = re.sub(r',(\s*[}\]])', r'\1', s)
+    s = re.sub(r'(?<=[}\]0-9truefalsenull"])\s*(\n\s*)"', r',\1"', s)
+    try:
+        return json.loads(s)
+    except json.JSONDecodeError as e:
+        raise ValueError(f"Failed to parse JSON after fixing: {e}")
+
+
+def create_finetuning_performance_report(training_time: float, device_type: str) -> str:
+    lines = [
+        "",
+        "| ------------------------------------------------------",
+        f"| Training Time : {training_time:.2f} s",
+    ]
+    if device_type == "cuda":
+        memory_used = torch.cuda.max_memory_allocated() / 1e9
+        lines.append(f"| Peak Memory   : {memory_used:.2f} GB")
+    lines.append("| ------------------------------------------------------")
+    return "\n".join(lines)
+
+
+
+
+
+
+

utils.py

@@ -0,0 +1,99 @@
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+from transformers import AutoConfig
+from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder
+
+from src.audiointeraction.generate.base import AUDIO_TOKENS_PER_CHUNK  # noqa: F401  (re-export for callers)
+from src.audiointeraction.model import GPT, Config
+from src.audiointeraction.utils import load_checkpoint
+
+
+def set_seed(seed: int = 1337) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+import json
+from pathlib import Path
+
+from safetensors.torch import load_file
+
+
+def load_model(fabric, model_config_dir, checkpoint_dir):
+    """Load a GPT from a local sharded safetensors directory.
+
+    `checkpoint_dir` must contain:
+        model.safetensors.index.json
+        model-00001-of-0000N.safetensors
+        ...
+    """
+    config = Config.from_file(Path(model_config_dir) / "model_config.yaml")
+    with fabric.init_module(empty_init=(fabric.world_size > 1)):
+        model = GPT(config)
+    model = fabric.setup(model)
+
+    checkpoint_dir = Path(checkpoint_dir)
+    index_path = checkpoint_dir / "model.safetensors.index.json"
+    if not index_path.is_file():
+        raise FileNotFoundError(
+            f"No model.safetensors.index.json under {checkpoint_dir}. "
+            f"Expected a sharded safetensors directory."
+        )
+
+    with open(index_path) as f:
+        index = json.load(f)
+    shard_files = sorted(set(index["weight_map"].values()))
+
+    state_dict = {}
+    for shard in shard_files:
+        state_dict.update(load_file(str(checkpoint_dir / shard), device="cpu"))
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=True)
+    if missing or unexpected:
+        print(f"[load_model] missing={missing[:3]}… unexpected={unexpected[:3]}…")
+    return model
+
+def load_audio_encoder(qwen_omni_ckpt, audio_tower_ckpt, device):
+    print(qwen_omni_ckpt)
+    cfg = AutoConfig.from_pretrained(qwen_omni_ckpt)
+    # Omni 的 config 是嵌套的：thinker_config.audio_config 才是 audio encoder 的配置
+    audio_cfg = cfg.thinker_config.audio_config
+    encoder = Qwen2_5OmniAudioEncoder._from_config(audio_cfg)
+    state_dict = torch.load(audio_tower_ckpt, map_location=device)
+    encoder.load_state_dict(state_dict)
+    encoder.to(device).requires_grad_(False).eval()
+    return encoder
+
+
+def resolve_checkpoint_paths(checkpoint_dir: str):
+    """Map a single checkpoint root → (model_config_dir, trained_checkpoint,
+    qwen_omni_ckpt, audio_tower_ckpt). The release layout is:
+
+        <checkpoint_dir>/
+            model_config.yaml + tokenizer.json + ...   ← model_config_dir = root
+            audiointeraction_LM.pt
+            audiointeraction_ChunkwisedEncoder.pth
+            qwen_2_5_omni_config/
+    """
+
+    ckpt = Path(checkpoint_dir)
+    return (
+        str(ckpt),
+        str(ckpt),
+        str(ckpt / "qwen25OmniConfig"),
+        str(ckpt / "audiointeraction_ChunkwisedEncoder.pth"),
+    )
+
+def get_best_device():
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+