MOSS-Transcribe-Diarize-0.9B / infer_moss_axengine.py
Nnow2024
Upload folder using huggingface_hub
a2f0cc7 verified 2 months ago
33.5 kB
#!/usr/bin/env python3
"""Run MOSS-Transcribe-Diarize with local AXEngine AXModels.
The audio front end is executed by whisper_encoder_vq_adaptor.axmodel.  The
Qwen3 LLM is deliberately driven through utils.infer_func.InferManager, using
the same prefill/decode call pattern as infer_axmodel.py:
    imer = InferManager(QWEN3_AX_CONFIG, llm_dir, max_seq_len=max_seq_len)
    token_ids = imer.prefill(tokenizer, token_ids, prompt_embeddings, slice_len=slice_len)
    imer.decode(tokenizer, token_ids, embed_matrix, slice_len=slice_len, ...)
All deployment files are resolved relative to this directory by default, so
the directory can be copied to an AX650 board as a standalone repository.
"""

from __future__ import annotations

import argparse
import atexit
import json
import re
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Mapping, Sequence

import numpy as np
from ml_dtypes import bfloat16
from tokenizers import Tokenizer

from utils.infer_func import InferManager


SCRIPT_DIR = Path(__file__).resolve().parent
DEFAULT_MODEL_DIR = SCRIPT_DIR
DEFAULT_AUDIO = SCRIPT_DIR / "wav" / "R8001_M8004_MS801.wav"
DEFAULT_AUDIO_MODEL = DEFAULT_MODEL_DIR / "whisper_encoder_vq_adaptor.axmodel"
DEFAULT_LLM_DIR = Path("./")
DEFAULT_TOKENIZER = DEFAULT_MODEL_DIR / "tokenizer.json"
DEFAULT_PREPROCESSOR_CONFIG = DEFAULT_MODEL_DIR / "preprocessor_config.json"
DEFAULT_PROMPT = (
    "请将音频转写为文本，每一段需以起始时间戳和说话人编号（[S01]、[S02]、[S03]…）开头，"
    "正文为对应的语音内容，并在段末标注结束时间戳，以清晰标明该段语音范围。"
)
DEFAULT_SYSTEM_PROMPT = "You are a helpful assistant."

AUDIO_START_TOKEN = "<|audio_start|>"
AUDIO_PAD_TOKEN = "<|audio_pad|>"
AUDIO_END_TOKEN = "<|audio_end|>"
EOS_TOKEN = "<|im_end|>"

# Values from preprocessor_config.json and the exported WhisperVQAdaptor's
# fixed 4x temporal merge.
SAMPLE_RATE = 16000
WHISPER_ENCODER_STRIDE = 2
AUDIO_MERGE_SIZE = 4
DEFAULT_AUDIO_TOKENS_PER_SECOND = 12.5
DEFAULT_TIME_MARKER_EVERY_SECONDS = 5

# ---------------------------------------------------------------------------
# Globally configurable LLM runtime settings
# ---------------------------------------------------------------------------
#
# Change these three values together for a differently compiled LLM.  In
# particular, slice_len must match the pN part of every LLM layer AXModel
# filename (for example, qwen3_p128_l0_together.axmodel).
#
# prefill_max_len limits only the user prompt (including injected audio
# tokens); max_seq_len is passed to InferManager as its KV-cache/decode limit.
slice_len = 256
prefill_max_len = 8448
max_seq_len = 10240


@dataclass(frozen=True)
class Qwen3AxConfig:
    """The Qwen3 fields required by utils.infer_func.InferManager."""

    hidden_size: int = 1024
    num_attention_heads: int = 16
    num_key_value_heads: int = 8
    head_dim: int = 128
    num_hidden_layers: int = 28


QWEN3_AX_CONFIG = Qwen3AxConfig()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "audio",
        nargs="?",
        type=Path,
        default=DEFAULT_AUDIO,
        help="Input audio readable by libsndfile (default: bundled WAV).",
    )
    parser.add_argument(
        "--model-dir",
        type=Path,
        default=DEFAULT_MODEL_DIR,
        help="Standalone MOSS-Transcribe-Diarize-0.9B directory (default: script directory).",
    )
    parser.add_argument(
        "--audio-model",
        type=Path,
        default=None,
        help="Whisper+VQ AXModel (default: MODEL_DIR/whisper_encoder_vq_adaptor.axmodel).",
    )
    parser.add_argument(
        "--llm-dir",
        type=Path,
        default=None,
        help="Split Qwen3 AXModel directory (default: MODEL_DIR/MOSS_qwen3_0.6b-ax650).",
    )
    parser.add_argument(
        "--tokenizer",
        type=Path,
        default=None,
        help="Tokenizer JSON containing the required audio special tokens.",
    )
    parser.add_argument(
        "--preprocessor-config",
        type=Path,
        default=None,
        help="Whisper feature-extractor configuration JSON.",
    )
    parser.add_argument("--prompt", default=DEFAULT_PROMPT, help="Transcription instruction.")
    parser.add_argument(
        "--system-prompt",
        default=DEFAULT_SYSTEM_PROMPT,
        help="System message in the Qwen chat prompt.",
    )
    parser.add_argument(
        "--time-marker-every-seconds",
        type=int,
        default=DEFAULT_TIME_MARKER_EVERY_SECONDS,
        help="Insert numeric time anchors inside the audio span at this interval (default: 5).",
    )
    parser.add_argument(
        "--disable-time-markers",
        action="store_true",
        help="Do not insert numeric time anchors into the audio span.",
    )
    parser.add_argument(
        "--stream",
        action="store_true",
        help="Stream decoded text from InferManager while generating.",
    )
    parser.add_argument(
        "--print-io",
        action="store_true",
        help="Print audio AXModel tensor metadata and the configured LLM settings.",
    )
    parser.add_argument(
        "--save-response",
        type=Path,
        default=None,
        help="Optional JSON file for transcript, token IDs, and timing.",
    )
    return parser.parse_args()


def validate_runtime_settings() -> None:
    """Validate the three global settings before allocating an InferManager."""

    for name, value in (
        ("slice_len", slice_len),
        ("prefill_max_len", prefill_max_len),
        ("max_seq_len", max_seq_len),
    ):
        if not isinstance(value, int) or value <= 0:
            raise ValueError(f"Global {name} must be a positive integer, got {value!r}")
    if prefill_max_len >= max_seq_len:
        raise ValueError(
            "Global prefill_max_len must be smaller than max_seq_len so the deployment has decode capacity: "
            f"{prefill_max_len} >= {max_seq_len}"
        )


def validate_llm_artifacts(llm_dir: Path) -> Path:
    """Validate the files InferManager will discover without loading AXEngine.
    InferManager discovers layer files by prefix.  Rejecting mixed pN sizes
    here prevents a difficult-to-diagnose shape-group failure later and makes
    the global slice_len setting explicit.
    """

    if not llm_dir.is_dir():
        raise FileNotFoundError(f"LLM AXModel directory not found: {llm_dir}")

    pattern = re.compile(r"^(?P<prefix>.+)_p(?P<prefill>\d+)_l(?P<index>\d+)_together\.axmodel$")
    by_prefix: dict[str, list[tuple[int, int]]] = {}
    for path in llm_dir.glob("*_p*_l*_together.axmodel"):
        match = pattern.match(path.name)
        if match:
            by_prefix.setdefault(match.group("prefix"), []).append(
                (int(match.group("prefill")), int(match.group("index")))
            )
    if not by_prefix:
        raise FileNotFoundError(f"No split layer AXModels found in {llm_dir}")

    # Match InferManager's choice: the prefix with the most layer files.
    prefix = max(by_prefix, key=lambda item: len(by_prefix[item]))
    layer_specs = by_prefix[prefix]
    compiled_slice_lengths = sorted({prefill for prefill, _ in layer_specs})
    if compiled_slice_lengths != [slice_len]:
        raise RuntimeError(
            f"Configured slice_len={slice_len} does not match the AXModels in {llm_dir}. "
            f"Prefix {prefix!r} contains p{compiled_slice_lengths}; set the global slice_len to the "
            "compiled pN value, or provide AXModels compiled for the configured value."
        )

    indices = sorted(index for _, index in layer_specs)
    expected_indices = list(range(QWEN3_AX_CONFIG.num_hidden_layers))
    if indices != expected_indices:
        raise RuntimeError(
            f"Expected Qwen3 layer indexes {expected_indices}, but found {indices} for prefix {prefix!r}"
        )

    post_path = llm_dir / f"{prefix}_post.axmodel"
    if not post_path.is_file():
        raise FileNotFoundError(f"Post AXModel not found: {post_path}")
    embedding_path = llm_dir / "model.embed_tokens.weight.bfloat16.bin"
    if not embedding_path.is_file():
        raise FileNotFoundError(f"BF16 embedding matrix not found: {embedding_path}")
    return embedding_path


def validate_prompt_length(prompt_len: int) -> None:
    """Ensure the prompt fits both configured prefill and KV-cache lengths."""

    if prompt_len > prefill_max_len:
        raise RuntimeError(
            f"Prompt length {prompt_len} exceeds prefill_max_len={prefill_max_len}. "
            "Use a shorter audio/prompt or increase the global prefill_max_len and supply compatible AXModels."
        )
    if prompt_len >= max_seq_len:
        raise RuntimeError(
            f"Prompt length {prompt_len} must be smaller than max_seq_len={max_seq_len}. "
            "Use a shorter audio/prompt or increase the global max_seq_len with compatible AXModels."
        )


def dtype_from_axengine(dtype: Any) -> np.dtype:
    """Map AXEngine's NumPy-like dtype metadata to a NumPy dtype."""

    name = str(dtype).lower()
    if "bfloat16" in name or "bf16" in name:
        return np.dtype(bfloat16)
    if "float32" in name or "fp32" in name:
        return np.dtype(np.float32)
    if "float16" in name or "fp16" in name:
        return np.dtype(np.float16)
    if "uint32" in name or "u32" in name:
        return np.dtype(np.uint32)
    if "int32" in name or "i32" in name:
        return np.dtype(np.int32)
    if "uint16" in name or "u16" in name:
        return np.dtype(np.uint16)
    if "int16" in name or "i16" in name:
        return np.dtype(np.int16)
    if "uint8" in name or "u8" in name:
        return np.dtype(np.uint8)
    if "int8" in name or "i8" in name:
        return np.dtype(np.int8)
    raise ValueError(f"Unsupported AXEngine dtype: {dtype!r}")


def ensure_finite(name: str, value: np.ndarray) -> None:
    """Fail early on malformed accelerator output rather than decoding garbage."""

    if not np.isfinite(np.asarray(value, dtype=np.float32)).all():
        raise RuntimeError(f"{name} contains NaN or Inf")


def release_ax_inference_session(session: Any) -> None:
    """Unload one AXEngine session, including axengine 0.1.x bindings."""

    inner = getattr(session, "_sess", None)
    unload = getattr(inner, "_unload", None)
    if not callable(unload):
        return
    try:
        unload()
    except Exception as exc:  # pragma: no cover - depends on AX runtime teardown
        print(f"[WARN] AXEngine session cleanup failed: {exc}")
    finally:
        # AXEngine 0.1.x can call _unload again from __del__.
        try:
            inner._unload = lambda: None
        except Exception:
            pass


def close_infer_manager(imer: InferManager) -> None:
    """Release sessions that are owned by the unmodified InferManager."""

    for session in getattr(imer, "decoder_sessions", ()):
        release_ax_inference_session(session)
    post_session = getattr(imer, "post_process_session", None)
    if post_session is not None:
        release_ax_inference_session(post_session)


def describe_nodes(nodes: Sequence[Any]) -> str:
    return ", ".join(f"{node.name}:{tuple(node.shape)}/{node.dtype}" for node in nodes)


def load_audio(path: Path, target_rate: int = SAMPLE_RATE) -> np.ndarray:
    """Read a mono waveform and resample it to the Whisper sample rate."""

    import soundfile as sf
    from scipy.signal import resample_poly

    waveform, sample_rate = sf.read(str(path), dtype="float32", always_2d=False)
    waveform = np.asarray(waveform, dtype=np.float32)
    if waveform.ndim == 2:
        waveform = waveform.mean(axis=1, dtype=np.float32)
    if waveform.ndim != 1 or waveform.size == 0:
        raise ValueError(f"Expected non-empty mono/stereo audio, got shape {waveform.shape}")
    if int(sample_rate) != int(target_rate):
        gcd = int(np.gcd(int(sample_rate), int(target_rate)))
        waveform = resample_poly(waveform, target_rate // gcd, int(sample_rate) // gcd).astype(np.float32)
    return np.ascontiguousarray(waveform, dtype=np.float32)


def load_feature_extractor(preprocessor_config: Path | None):
    """Create the matching Whisper log-Mel frontend without loading weights."""

    from transformers import WhisperFeatureExtractor

    if preprocessor_config is not None and preprocessor_config.is_file():
        return WhisperFeatureExtractor.from_pretrained(str(preprocessor_config.parent), local_files_only=True)

    # Portable fallback for deployments containing only the compiled models.
    return WhisperFeatureExtractor(
        feature_size=80,
        sampling_rate=SAMPLE_RATE,
        hop_length=160,
        chunk_length=30,
        n_fft=400,
        padding_value=0.0,
        dither=0.0,
        return_attention_mask=False,
    )


def compute_audio_token_length(num_samples: int, feature_extractor: Any) -> int:
    stride = int(feature_extractor.hop_length) * WHISPER_ENCODER_STRIDE * AUDIO_MERGE_SIZE
    return (int(num_samples) - 1) // stride + 1


def make_audio_features(waveform: np.ndarray, feature_extractor: Any) -> tuple[np.ndarray, list[int]]:
    """Split audio into padded 30-second chunks and compute valid token sizes."""

    n_samples = int(feature_extractor.n_samples)
    chunks: list[np.ndarray] = []
    token_lengths: list[int] = []
    for start in range(0, int(waveform.size), n_samples):
        chunk = waveform[start : start + n_samples]
        token_lengths.append(compute_audio_token_length(int(chunk.size), feature_extractor))
        if chunk.size < n_samples:
            chunk = np.pad(chunk, (0, n_samples - chunk.size))
        chunks.append(np.asarray(chunk, dtype=np.float32))
    if not chunks:
        raise ValueError("Audio must contain at least one sample")

    batch = feature_extractor(
        chunks,
        sampling_rate=int(feature_extractor.sampling_rate),
        padding="max_length",
        return_tensors="np",
    )
    input_features = np.ascontiguousarray(batch["input_features"], dtype=np.float32)
    expected_shape = (len(chunks), int(feature_extractor.feature_size), int(feature_extractor.nb_max_frames))
    if tuple(input_features.shape) != expected_shape:
        raise RuntimeError(f"Unexpected Whisper feature shape {input_features.shape}; expected {expected_shape}")
    return input_features, token_lengths


def build_chat_prompt(system_prompt: str, prompt: str) -> str:
    """Reproduce the repository's one-audio Qwen chat template."""

    return (
        f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
        "<|im_start|>user\n"
        f"{AUDIO_START_TOKEN}{AUDIO_PAD_TOKEN}{AUDIO_END_TOKEN}\n"
        f"{prompt}<|im_end|>\n"
        "<|im_start|>assistant\n"
    )


def get_digit_token_ids(tokenizer: Tokenizer) -> dict[str, int]:
    ids: dict[str, int] = {}
    for digit in "0123456789":
        encoded = tokenizer.encode(digit, add_special_tokens=False).ids
        if len(encoded) != 1:
            raise ValueError(f"Digit {digit!r} must be exactly one tokenizer token, got {encoded}")
        ids[digit] = int(encoded[0])
    return ids


def make_audio_span_ids(
    audio_token_id: int,
    audio_seq_len: int,
    digit_token_ids: Mapping[str, int],
    *,
    audio_tokens_per_second: float = DEFAULT_AUDIO_TOKENS_PER_SECOND,
    time_marker_every_seconds: int = DEFAULT_TIME_MARKER_EVERY_SECONDS,
    enable_time_markers: bool = True,
) -> list[int]:
    """Match MossTranscribeDiarizeProcessor._audio_span_ids."""

    audio_seq_len = int(audio_seq_len)
    if not enable_time_markers or audio_seq_len <= 0 or time_marker_every_seconds <= 0:
        return [int(audio_token_id)] * max(audio_seq_len, 0)

    tokens_per_marker = int(audio_tokens_per_second * time_marker_every_seconds)
    if tokens_per_marker <= 0:
        return [int(audio_token_id)] * audio_seq_len

    duration = audio_seq_len / float(audio_tokens_per_second)
    result: list[int] = []
    consumed = 0
    for second in range(time_marker_every_seconds, int(duration) + 1, time_marker_every_seconds):
        position = (second // time_marker_every_seconds) * tokens_per_marker
        segment_len = position - consumed
        if segment_len > 0:
            result.extend([int(audio_token_id)] * segment_len)
            consumed += segment_len
        result.extend(digit_token_ids[digit] for digit in str(second))
    remainder = audio_seq_len - consumed
    if remainder > 0:
        result.extend([int(audio_token_id)] * remainder)
    return result


def build_prompt_token_ids(
    tokenizer: Tokenizer,
    system_prompt: str,
    prompt: str,
    audio_token_length: int,
    *,
    time_marker_every_seconds: int,
    enable_time_markers: bool,
) -> tuple[list[int], int, int]:
    """Tokenise the chat prompt and expand its single audio placeholder."""

    audio_token_id = tokenizer.token_to_id(AUDIO_PAD_TOKEN)
    if audio_token_id is None:
        raise ValueError(f"Tokenizer is missing required {AUDIO_PAD_TOKEN!r}")
    token_ids = tokenizer.encode(build_chat_prompt(system_prompt, prompt), add_special_tokens=False).ids
    positions = [index for index, token_id in enumerate(token_ids) if token_id == audio_token_id]
    if len(positions) != 1:
        raise RuntimeError(f"Expected exactly one audio placeholder, found {len(positions)}")

    audio_span = make_audio_span_ids(
        int(audio_token_id),
        audio_token_length,
        get_digit_token_ids(tokenizer),
        time_marker_every_seconds=time_marker_every_seconds,
        enable_time_markers=enable_time_markers,
    )
    position = positions[0]
    expanded = token_ids[:position] + audio_span + token_ids[position + 1 :]
    audio_positions = sum(token_id == audio_token_id for token_id in expanded)
    if audio_positions != int(audio_token_length):
        raise RuntimeError(
            f"Audio placeholder count {audio_positions} does not match valid audio embedding count {audio_token_length}"
        )
    return [int(token_id) for token_id in expanded], int(audio_token_id), len(audio_span)


class InferTokenizerAdapter:
    """Adapt tokenizers.Tokenizer to the tokenizer interface InferManager uses."""

    def __init__(self, tokenizer: Tokenizer, eos_token_id: int) -> None:
        self._tokenizer = tokenizer
        self.eos_token_id = int(eos_token_id)

    def decode(self, token_ids: int | np.integer | Sequence[int], skip_special_tokens: bool = True) -> str:
        if isinstance(token_ids, (int, np.integer)):
            token_ids = [int(token_ids)]
        return self._tokenizer.decode(
            [int(token_id) for token_id in token_ids],
            skip_special_tokens=skip_special_tokens,
        )


class MossAxEngineRunner:
    """Own the audio AXModel, tokenizer, and tied embedding matrix."""

    def __init__(
        self,
        audio_model_path: Path,
        tokenizer_path: Path,
        embedding_path: Path,
        *,
        print_io: bool = False,
    ) -> None:
        try:
            from axengine import InferenceSession
        except ImportError as exc:  # pragma: no cover - requires target AX runtime
            raise RuntimeError(
                "Cannot import axengine. Install the board-matched AXEngine Python package and ensure "
                "the AX runtime shared libraries are available."
            ) from exc

        if not audio_model_path.is_file():
            raise FileNotFoundError(audio_model_path)
        if not tokenizer_path.is_file():
            raise FileNotFoundError(f"Tokenizer JSON not found: {tokenizer_path}")

        self.hidden_size = QWEN3_AX_CONFIG.hidden_size
        self.hidden_dtype = np.dtype(bfloat16)
        self.tokenizer = Tokenizer.from_file(str(tokenizer_path))

        raw_embeddings = np.memmap(embedding_path, mode="r", dtype=np.uint16)
        if raw_embeddings.size % self.hidden_size:
            raise RuntimeError(
                f"Embedding file has {raw_embeddings.size} BF16 values, not divisible by hidden size {self.hidden_size}"
            )
        self._embedding_storage = raw_embeddings
        self.embed_matrix = raw_embeddings.view(bfloat16).reshape(-1, self.hidden_size)
        self.vocab_size = int(self.embed_matrix.shape[0])

        eos_token_id = self.tokenizer.token_to_id(EOS_TOKEN)
        if eos_token_id is None:
            raise ValueError(f"Tokenizer is missing end-of-sequence token {EOS_TOKEN!r}")
        self.eos_token_id = int(eos_token_id)

        self.audio_session = InferenceSession(str(audio_model_path))
        self._closed = False
        atexit.register(self.close)

        self.audio_input_nodes = self.audio_session.get_inputs(shape_group=0)
        self.audio_output_nodes = self.audio_session.get_outputs(shape_group=0)
        if len(self.audio_input_nodes) != 1 or len(self.audio_output_nodes) != 1:
            raise RuntimeError(
                "Audio AXModel must have exactly one input and one output; "
                f"got input=[{describe_nodes(self.audio_input_nodes)}], "
                f"output=[{describe_nodes(self.audio_output_nodes)}]"
            )
        self.audio_input_name = self.audio_input_nodes[0].name
        self.audio_input_shape = tuple(self.audio_input_nodes[0].shape)
        self.audio_input_dtype = dtype_from_axengine(self.audio_input_nodes[0].dtype)
        self.audio_output_shape = tuple(self.audio_output_nodes[0].shape)

        if print_io:
            self.print_io()

    def print_io(self) -> None:
        """Print accelerator-visible audio I/O and LLM global settings."""

        print(f"[Audio IO] input={describe_nodes(self.audio_input_nodes)}")
        print(f"[Audio IO] output={describe_nodes(self.audio_output_nodes)}")
        print(
            f"[LLM] InferManager config: layers={QWEN3_AX_CONFIG.num_hidden_layers} "
            f"hidden={QWEN3_AX_CONFIG.hidden_size} slice_len={slice_len} "
            f"prefill_max_len={prefill_max_len} max_seq_len={max_seq_len} vocab={self.vocab_size}"
        )

    def close(self) -> None:
        if self._closed:
            return
        release_ax_inference_session(self.audio_session)
        self._closed = True

    def decode_text(self, token_ids: Sequence[int], *, skip_special_tokens: bool = True) -> str:
        return self.tokenizer.decode([int(token_id) for token_id in token_ids], skip_special_tokens=skip_special_tokens)

    def encode_audio_features(self, input_features: np.ndarray) -> np.ndarray:
        """Execute the static batch=1 audio AXModel once per 30-second chunk."""

        if input_features.ndim != 3:
            raise ValueError(f"Expected [chunks, 80, 3000] features, got {input_features.shape}")
        expected_shape = (1, *input_features.shape[1:])
        if tuple(self.audio_input_shape) != expected_shape:
            raise RuntimeError(
                f"Audio AXModel expects {self.audio_input_shape}, but feature extractor produced {expected_shape}"
            )

        outputs: list[np.ndarray] = []
        for chunk_index, chunk in enumerate(input_features):
            feed = np.ascontiguousarray(chunk[None], dtype=self.audio_input_dtype)
            output = self.audio_session.run(None, {self.audio_input_name: feed})[0]
            if tuple(output.shape) != self.audio_output_shape:
                raise RuntimeError(
                    f"Audio AXModel chunk {chunk_index} returned {output.shape}, expected {self.audio_output_shape}"
                )
            ensure_finite(f"audio chunk {chunk_index}", output)
            outputs.append(np.asarray(output))
        return np.concatenate(outputs, axis=0)

    def build_prompt_embeddings(
        self,
        token_ids: Sequence[int],
        audio_token_id: int,
        chunk_embeddings: np.ndarray,
        chunk_token_lengths: Sequence[int],
    ) -> np.ndarray:
        """Replace audio-pad token embeddings with trimmed audio AXModel output."""

        ids = np.asarray(token_ids, dtype=np.int64)
        if ids.size == 0:
            raise ValueError("Prompt must contain at least one token")
        if int(ids.min()) < 0 or int(ids.max()) >= self.vocab_size:
            raise ValueError(f"Prompt includes a token outside vocabulary [0, {self.vocab_size})")
        if chunk_embeddings.ndim != 3 or chunk_embeddings.shape[0] != len(chunk_token_lengths):
            raise ValueError(
                f"Chunk embeddings {chunk_embeddings.shape} do not match {len(chunk_token_lengths)} chunk lengths"
            )
        if chunk_embeddings.shape[-1] != self.hidden_size:
            raise ValueError(
                f"Audio embedding width {chunk_embeddings.shape[-1]} differs from LLM hidden size {self.hidden_size}"
            )

        audio_pieces: list[np.ndarray] = []
        for chunk_index, token_length in enumerate(chunk_token_lengths):
            token_length = int(token_length)
            if token_length <= 0 or token_length > chunk_embeddings.shape[1]:
                raise ValueError(
                    f"Invalid audio token length for chunk {chunk_index}: {token_length}; "
                    f"AXModel has {chunk_embeddings.shape[1]} positions"
                )
            audio_pieces.append(chunk_embeddings[chunk_index, :token_length])
        audio_embeddings = np.concatenate(audio_pieces, axis=0)
        audio_positions = np.flatnonzero(ids == int(audio_token_id))
        if audio_positions.size != audio_embeddings.shape[0]:
            raise RuntimeError(
                f"Audio placeholder count {audio_positions.size} does not match "
                f"audio embedding count {audio_embeddings.shape[0]}"
            )

        embeddings = np.asarray(self.embed_matrix[ids], dtype=self.hidden_dtype)
        embeddings[audio_positions, :] = audio_embeddings.astype(self.hidden_dtype, copy=False)
        return embeddings


# The model emits "[0.31][S01]试错的过程很简单[2.15]" per segment, all runs
# together on one line.  Bracket contents never nest, so every field is matched
# as bracket-free text.
TRANSCRIPT_SEGMENT_PATTERN = re.compile(
    r"\[(?P<start>[^\[\]]*)\]\s*"
    r"\[(?P<speaker>[^\[\]]*)\]\s*"
    r"(?P<text>[^\[\]]*?)\s*"
    r"\[(?P<end>[^\[\]]*)\]"
)


def parse_transcript_segments(transcript: str) -> list[dict[str, str]]:
    """Split the raw transcript into start/end/speaker/text segments."""

    return [
        {
            "start": match.group("start").strip(),
            "end": match.group("end").strip(),
            "speaker": match.group("speaker").strip(),
            "text": match.group("text").strip(),
        }
        for match in TRANSCRIPT_SEGMENT_PATTERN.finditer(transcript)
    ]


def format_transcript(segments: Sequence[Mapping[str, str]]) -> str:
    """Render segments as "[start][end][speaker]text", one per line."""

    return "\n".join(
        f"[{segment['start']}][{segment['end']}][{segment['speaker']}]{segment['text']}"
        for segment in segments
    )


def main() -> None:
    args = parse_args()
    if args.time_marker_every_seconds < 0:
        raise ValueError("--time-marker-every-seconds must be non-negative")
    validate_runtime_settings()

    model_dir = args.model_dir.resolve()
    audio_model = (args.audio_model or model_dir / DEFAULT_AUDIO_MODEL.name).resolve()
    llm_dir = (args.llm_dir or model_dir / DEFAULT_LLM_DIR.name).resolve()
    tokenizer_path = (args.tokenizer or model_dir / DEFAULT_TOKENIZER.name).resolve()
    preprocessor_config = (args.preprocessor_config or model_dir / DEFAULT_PREPROCESSOR_CONFIG.name).resolve()
    audio_path = args.audio.resolve()
    if not audio_path.is_file():
        raise FileNotFoundError(audio_path)
    embedding_path = validate_llm_artifacts(llm_dir)

    timings: dict[str, float] = {}
    print(
        f"[LLM] slice_len={slice_len} prefill_max_len={prefill_max_len} max_seq_len={max_seq_len}",
        flush=True,
    )
    print(f"[Audio] loading {audio_path}", flush=True)
    started = time.perf_counter()
    waveform = load_audio(audio_path)
    feature_extractor = load_feature_extractor(preprocessor_config)
    input_features, token_lengths = make_audio_features(waveform, feature_extractor)
    timings["audio_preprocess_s"] = time.perf_counter() - started
    duration_s = waveform.size / float(feature_extractor.sampling_rate)
    print(
        f"[Audio] duration={duration_s:.3f}s chunks={len(token_lengths)} "
        f"valid_audio_tokens={sum(token_lengths)} features={tuple(input_features.shape)} "
        f"preprocess={timings['audio_preprocess_s']:.3f}s",
        flush=True,
    )

    runner: MossAxEngineRunner | None = None
    imer: InferManager | None = None
    prompt_len = 0
    generated_ids: list[int] = []
    try:
        print(f"[AXEngine] loading audio model: {audio_model}", flush=True)
        runner = MossAxEngineRunner(
            audio_model,
            tokenizer_path,
            embedding_path,
            print_io=args.print_io,
        )

        started = time.perf_counter()
        chunk_embeddings = runner.encode_audio_features(input_features)
        timings["audio_axmodel_s"] = time.perf_counter() - started
        print(
            f"[Audio AXModel] embeddings={tuple(chunk_embeddings.shape)} "
            f"cost={timings['audio_axmodel_s']:.3f}s",
            flush=True,
        )

        started = time.perf_counter()
        prompt_token_ids, audio_token_id, expanded_audio_span_len = build_prompt_token_ids(
            runner.tokenizer,
            args.system_prompt,
            args.prompt,
            sum(token_lengths),
            time_marker_every_seconds=args.time_marker_every_seconds,
            enable_time_markers=not args.disable_time_markers,
        )
        prompt_embeddings = runner.build_prompt_embeddings(
            prompt_token_ids,
            audio_token_id,
            chunk_embeddings,
            token_lengths,
        )
        prompt_len = len(prompt_token_ids)
        validate_prompt_length(prompt_len)
        timings["prompt_build_s"] = time.perf_counter() - started
        print(
            f"[Prompt] tokens={prompt_len} audio_span={expanded_audio_span_len} "
            f"audio_embeddings={sum(token_lengths)} build={timings['prompt_build_s']:.3f}s",
            flush=True,
        )

        print(f"[AXEngine] loading LLM models through InferManager: {llm_dir}", flush=True)
        started = time.perf_counter()
        imer = InferManager(QWEN3_AX_CONFIG, str(llm_dir), max_seq_len=max_seq_len)
        timings["llm_load_s"] = time.perf_counter() - started

        tokenizer = InferTokenizerAdapter(runner.tokenizer, runner.eos_token_id)
        print("[Generate] decoding until EOS or InferManager max_seq_len.", flush=True)
        started = time.perf_counter()
        token_ids = imer.prefill(
            tokenizer,
            prompt_token_ids,
            prompt_embeddings,
            slice_len=slice_len,
        )
        timings["llm_prefill_s"] = time.perf_counter() - started

        started = time.perf_counter()
        imer.decode(
            tokenizer,
            token_ids,
            runner.embed_matrix,
            slice_len=slice_len,
            eos_token_id=[runner.eos_token_id],
            stream=args.stream,
        )
        timings["llm_generate_s"] = time.perf_counter() - started
        timings["llm_total_s"] = timings["llm_prefill_s"] + timings["llm_generate_s"]
        generated_ids = [int(token_id) for token_id in token_ids[prompt_len:]]
    finally:
        if imer is not None:
            close_infer_manager(imer)
        if runner is not None:
            runner.close()

    assert runner is not None
    raw_text = runner.decode_text(generated_ids, skip_special_tokens=False)
    transcript = runner.decode_text(generated_ids, skip_special_tokens=True)
    segments = parse_transcript_segments(transcript)
    # Fall back to the model's own layout when nothing parses as a segment.
    formatted_transcript = format_transcript(segments) if segments else transcript.strip()
    print(f"\n[Generate] generated={len(generated_ids)} cost={timings['llm_generate_s']:.3f}s", flush=True)
    print(f"\n[Transcript] segments={len(segments)}")
    print(formatted_transcript)

    if args.save_response is not None:
        output = args.save_response.resolve()
        output.parent.mkdir(parents=True, exist_ok=True)
        output.write_text(
            json.dumps(
                {
                    "audio": str(audio_path),
                    "prompt": args.prompt,
                    "system_prompt": args.system_prompt,
                    "generated_token_ids": generated_ids,
                    "raw_text": raw_text,
                    "text": transcript,
                    "formatted_text": formatted_transcript,
                    "segments": segments,
                    "audio_chunk_token_lengths": token_lengths,
                    "runtime": {
                        "slice_len": slice_len,
                        "prefill_max_len": prefill_max_len,
                        "max_seq_len": max_seq_len,
                        "prompt_len": prompt_len,
                    },
                    "timings": timings,
                },
                ensure_ascii=False,
                indent=2,
            )
            + "\n",
            encoding="utf-8",
        )
        print(f"[Save] {output}")


if __name__ == "__main__":
    main()