diff --git a/families/parakeet_tdt/README.md b/families/parakeet_tdt/README.md new file mode 100644 index 0000000000..8b84c9e959 --- /dev/null +++ b/families/parakeet_tdt/README.md @@ -0,0 +1,40 @@ +# Parakeet TDT 0.6B v3 + +This family targets the native Hugging Face Transformers checkpoint +`nvidia/parakeet-tdt-0.6b-v3` at revision +`541d1f99c6b0c3cd0b11a95167540bb8edefd82b` (publisher license: CC BY 4.0). +It exposes the public `SpeechTranscription` Task for one audio request at a time, +with FP16 or FP32 TensorRT engines and a fixed 30-second input window. Runtime +input accepts interleaved PCM, downmixes channels, and resamples to 16 kHz. +Longer audio, language controls, streaming, quantization and multi-GPU execution +are rejected or not exposed. NeMo archives are not a build input for this port. + +The migration has CPU contract coverage; GPU conversion and transcript parity +remain unverified. Do not interpret the presence of the manifest as qualification. + +## Validation + +Install this family's `requirements.txt` in its isolated build/reference +environment. The pinned Transformers version provides `AutoModelForTDT`, which +is not available in the repository's base Transformers 5.2.0 environment. + +Build with `TRTMC_BUILD_TESTS=ON` to produce `test_parakeet_tdt_sdk_cpp`. +After caching the pinned Hugging Face snapshot, select the E2E explicitly: + +```sh +TRTMC_RUNTIME_ROOT=/path/to/runtime \ +TRTMC_NATIVE_BUILD_DIR=/path/to/native/build \ +PYTHONPATH=core/builder:. python -m pytest families/parakeet_tdt/tests/test_e2e.py \ + --e2e-testcase parakeet-tdt-0.6b-v3 -v +``` + +`--e2e-model parakeet_tdt` selects all eight audio cases. An explicit selection +fails when dependencies, artifacts or CUDA are unavailable; an ordinary CPU +test run reports these E2Es as skipped. The test invokes the public C++ SDK +consumer, records native/reference evidence, and requires exact transcript +agreement after the original case/whitespace normalization. It does not relax +the old PR's exact-transcript gate to a shape-only or approximate comparison. + +`TRTMC_PARAKEET_TDT_MODEL_DIR` may point to a prepared copy of that same pinned +snapshot. Record its provenance when using the override. The tests do not +download model weights automatically. Audio fixture generation is family-local. diff --git a/families/parakeet_tdt/__init__.py b/families/parakeet_tdt/__init__.py new file mode 100644 index 0000000000..1f89c21a5a --- /dev/null +++ b/families/parakeet_tdt/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Family-owned Parakeet TDT implementation.""" diff --git a/families/parakeet_tdt/checkpoint.py b/families/parakeet_tdt/checkpoint.py new file mode 100644 index 0000000000..5262f39766 --- /dev/null +++ b/families/parakeet_tdt/checkpoint.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Checkpoint readers and tensor normalization for Parakeet TDT.""" + +from __future__ import annotations + +import numpy as np + + +class WeightDict(dict): + """Normalized family-owned build tensors.""" + + +def _array(value) -> np.ndarray: + if hasattr(value, "detach"): + value = value.detach().cpu().numpy() + return np.ascontiguousarray(np.asarray(value, dtype=np.float32)) + + +def _transpose(value) -> np.ndarray: + array = _array(value) + if array.ndim != 2: + raise ValueError(f"expected rank-2 weight, got {array.shape}") + return np.ascontiguousarray(array.T) + + +def _transpose_2d(value, name: str, precision: str = "fp32") -> np.ndarray: + """Family graph-builder adapter with the repository's mapper signature.""" + del name, precision + return _transpose(value) diff --git a/families/parakeet_tdt/config.py b/families/parakeet_tdt/config.py new file mode 100644 index 0000000000..1a89e8429b --- /dev/null +++ b/families/parakeet_tdt/config.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed configuration for the pinned Parakeet TDT checkpoint contract.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + + +@dataclass(frozen=True) +class ParakeetTDTConfig: + model_type: str + architectures: tuple[str, ...] + vocab_size: int + blank_id: int + pad_id: int + durations: tuple[int, ...] + max_symbols_per_step: int + encoder_hidden_size: int + encoder_layers: int + encoder_heads: int + encoder_ffn_size: int + encoder_conv_kernel_size: int + encoder_max_positions: int + num_mel_bins: int + subsampling_channels: int + subsampling_factor: int + decoder_hidden_size: int + decoder_layers: int + joint_activation: str + raw: dict + + @classmethod + def from_json(cls, text: str) -> "ParakeetTDTConfig": + raw = json.loads(text) + enc = raw.get("encoder_config") or {} + return cls( + model_type=str(raw.get("model_type", "")), + architectures=tuple(str(item) for item in raw.get("architectures", ())), + vocab_size=int(raw.get("vocab_size", 0)), + blank_id=int(raw.get("blank_token_id", -1)), + pad_id=int(raw.get("pad_token_id", -1)), + durations=tuple(int(item) for item in raw.get("durations", ())), + max_symbols_per_step=int(raw.get("max_symbols_per_step", 0)), + encoder_hidden_size=int(enc.get("hidden_size", 0)), + encoder_layers=int(enc.get("num_hidden_layers", 0)), + encoder_heads=int(enc.get("num_attention_heads", 0)), + encoder_ffn_size=int(enc.get("intermediate_size", 0)), + encoder_conv_kernel_size=int(enc.get("conv_kernel_size", 0)), + encoder_max_positions=int(enc.get("max_position_embeddings", 0)), + num_mel_bins=int(enc.get("num_mel_bins", 0)), + subsampling_channels=int(enc.get("subsampling_conv_channels", 0)), + subsampling_factor=int(enc.get("subsampling_factor", 0)), + decoder_hidden_size=int(raw.get("decoder_hidden_size", 0)), + decoder_layers=int(raw.get("num_decoder_layers", 0)), + joint_activation=str(raw.get("hidden_act", "")), + raw=raw, + ) + + @classmethod + def from_dir(cls, model_dir: str | Path) -> "ParakeetTDTConfig": + config_path = Path(model_dir) / "config.json" + if not config_path.is_file(): + raise FileNotFoundError( + f"Parakeet TDT model directory is missing config.json: {config_path}" + ) + return cls.from_json(config_path.read_text(encoding="utf-8")) + + def validate_supported_checkpoint(self) -> None: + expected = { + "model_type": (self.model_type, "parakeet_tdt"), + "architecture": (self.architectures, ("ParakeetForTDT",)), + "vocab_size": (self.vocab_size, 8193), + "blank_token_id": (self.blank_id, 8192), + "durations": (self.durations, (0, 1, 2, 3, 4)), + "encoder_hidden_size": (self.encoder_hidden_size, 1024), + "encoder_layers": (self.encoder_layers, 24), + "encoder_heads": (self.encoder_heads, 8), + "encoder_ffn_size": (self.encoder_ffn_size, 4096), + "encoder_conv_kernel_size": (self.encoder_conv_kernel_size, 9), + "num_mel_bins": (self.num_mel_bins, 128), + "subsampling_factor": (self.subsampling_factor, 8), + "decoder_hidden_size": (self.decoder_hidden_size, 640), + "decoder_layers": (self.decoder_layers, 2), + "joint_activation": (self.joint_activation, "relu"), + "max_symbols_per_step": (self.max_symbols_per_step, 10), + } + mismatches = [ + f"{name}={actual!r}, expected {wanted!r}" + for name, (actual, wanted) in expected.items() + if actual != wanted + ] + if mismatches: + raise ValueError("Unsupported Parakeet TDT checkpoint contract: " + "; ".join(mismatches)) diff --git a/families/parakeet_tdt/engines.py b/families/parakeet_tdt/engines.py new file mode 100644 index 0000000000..0f58c2134a --- /dev/null +++ b/families/parakeet_tdt/engines.py @@ -0,0 +1,615 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Family-local TensorRT graphs ported from PR #1060; native HF input only.""" + +from __future__ import annotations +import sys +from pathlib import Path +import numpy as np +import tensorrt as trt +from . import graph_ops +from .checkpoint import WeightDict, _transpose_2d +from .model_config import ModelConfig +from .fastconformer import _build_encoder, _compute_enc_seq_len, _compute_causal_enc_seq_len, _relative_pe, _to_np + + +def _load_hf_as_nemo(model_dir: str): + """Expose the official HF safetensors under the NeMo-style encoder keys. + + The graph builders consume one normalized, family-private layout. Keeping + this conversion explicit avoids choosing a topology from a checkpoint name. + """ + from safetensors.numpy import load_file + + root = Path(model_dir) + state = dict(load_file(root / "model.safetensors")) + if any(key.startswith("prompt_kernel.") for key in state): + raise ValueError("Parakeet TDT native checkpoint does not support prompt_kernel") + from .config import ParakeetTDTConfig + + cfg = ParakeetTDTConfig.from_dir(root) + cfg.validate_supported_checkpoint() + normalized = dict(state) + aliases = { + "encoder.pre_encode.conv.0.weight": "encoder.subsampling.layers.0.weight", + "encoder.pre_encode.conv.0.bias": "encoder.subsampling.layers.0.bias", + "encoder.pre_encode.conv.2.weight": "encoder.subsampling.layers.2.weight", + "encoder.pre_encode.conv.2.bias": "encoder.subsampling.layers.2.bias", + "encoder.pre_encode.conv.3.weight": "encoder.subsampling.layers.3.weight", + "encoder.pre_encode.conv.3.bias": "encoder.subsampling.layers.3.bias", + "encoder.pre_encode.conv.5.weight": "encoder.subsampling.layers.5.weight", + "encoder.pre_encode.conv.5.bias": "encoder.subsampling.layers.5.bias", + "encoder.pre_encode.conv.6.weight": "encoder.subsampling.layers.6.weight", + "encoder.pre_encode.conv.6.bias": "encoder.subsampling.layers.6.bias", + "encoder.pre_encode.out.weight": "encoder.subsampling.linear.weight", + "encoder.pre_encode.out.bias": "encoder.subsampling.linear.bias", + "decoder.prediction.embed.weight": "decoder.embedding.weight", + "joint.enc.weight": "encoder_projector.weight", + "joint.enc.bias": "encoder_projector.bias", + "joint.pred.weight": "decoder.decoder_projector.weight", + "joint.pred.bias": "decoder.decoder_projector.bias", + "joint.joint_net.2.weight": "joint.head.weight", + "joint.joint_net.2.bias": "joint.head.bias", + } + for layer in range(cfg.decoder_layers): + for field in ("weight_ih", "weight_hh", "bias_ih", "bias_hh"): + aliases[f"decoder.prediction.dec_rnn.{field}_l{layer}"] = ( + f"decoder.lstm.{field}_l{layer}") + for layer in range(cfg.encoder_layers): + nk = f"encoder.layers.{layer}" + for nemo, hf in (("linear_q", "q_proj"), ("linear_k", "k_proj"), + ("linear_v", "v_proj"), ("linear_out", "o_proj"), + ("linear_pos", "relative_k_proj")): + aliases[f"{nk}.self_attn.{nemo}.weight"] = f"{nk}.self_attn.{hf}.weight" + aliases[f"{nk}.self_attn.pos_bias_u"] = f"{nk}.self_attn.bias_u" + aliases[f"{nk}.self_attn.pos_bias_v"] = f"{nk}.self_attn.bias_v" + aliases[f"{nk}.conv.batch_norm.weight"] = f"{nk}.conv.norm.weight" + aliases[f"{nk}.conv.batch_norm.bias"] = f"{nk}.conv.norm.bias" + aliases[f"{nk}.conv.batch_norm.running_mean"] = f"{nk}.conv.norm.running_mean" + aliases[f"{nk}.conv.batch_norm.running_var"] = f"{nk}.conv.norm.running_var" + missing = [(dst, src) for dst, src in aliases.items() if src not in state] + if missing: + detail = ", ".join(f"{dst} <- {src}" for dst, src in missing[:8]) + raise KeyError(f"HF Parakeet checkpoint is missing normalized tensors: {detail}") + normalized.update({dst: state[src] for dst, src in aliases.items()}) + nemo_cfg = { + "encoder": {"d_model": cfg.encoder_hidden_size, + "n_layers": cfg.encoder_layers, + "n_heads": cfg.encoder_heads, + "ff_expansion_factor": cfg.encoder_ffn_size // cfg.encoder_hidden_size, + "conv_kernel_size": cfg.encoder_conv_kernel_size, + "subsampling_conv_channels": cfg.subsampling_channels, + "att_context_size": [[-1, -1]], + "conv_norm_type": "batch_norm", "conv_context_size": "symmetric"}, + "preprocessor": {"features": cfg.num_mel_bins}, + "decoder": {"blank_idx": cfg.blank_id, + "prednet": {"pred_hidden": cfg.decoder_hidden_size, + "pred_rnn_layers": cfg.decoder_layers, + "rnn_hidden_size": cfg.decoder_hidden_size}}, + "joint": {"jointnet": {"joint_hidden": cfg.decoder_hidden_size, + "activation": cfg.joint_activation}}, + "decoding": {"max_symbols_per_step": cfg.max_symbols_per_step}, + "tdt_durations": list(cfg.durations), + } + return normalized, nemo_cfg + + +def _cfg_int(*values, default: int) -> int: + for value in values: + if value is not None: + try: + return int(value) + except (TypeError, ValueError): + pass + return default + + +def _cfg_dict(root: dict, *path: str) -> dict: + cur = root + for item in path: + if not isinstance(cur, dict): + return {} + cur = cur.get(item, {}) + return cur if isinstance(cur, dict) else {} + + +def _find_tensor(sd: dict, candidates: list[str], label: str): + for key in candidates: + if key in sd: + return sd[key] + suffix_matches = [] + for key in sd: + for suffix in candidates: + if key.endswith(suffix): + suffix_matches.append(key) + break + if len(suffix_matches) == 1: + return sd[suffix_matches[0]] + if len(suffix_matches) > 1: + raise KeyError(f"Ambiguous tensor for {label}: {suffix_matches}") + raise KeyError(f"Missing tensor for {label}; tried {candidates}") + + +def _find_joint_linear(sd: dict, prefix: str, label: str): + candidates = [ + f"{prefix}.joint_net.1.weight", + f"{prefix}.joint_net.2.weight", + f"{prefix}.joint_net.0.weight", + ] + for key in candidates: + if key in sd: + bias_key = key[:-6] + "bias" + if bias_key not in sd: + raise KeyError(f"Missing tensor for {label} bias: {bias_key}") + return sd[key], sd[bias_key] + + suffixes = [".joint_net.1.weight", ".joint_net.2.weight", ".joint_net.0.weight"] + matches = [key for key in sd if key.startswith(prefix) and any(key.endswith(s) for s in suffixes)] + if len(matches) == 1: + key = matches[0] + bias_key = key[:-6] + "bias" + if bias_key not in sd: + raise KeyError(f"Missing tensor for {label} bias: {bias_key}") + return sd[key], sd[bias_key] + raise KeyError(f"Missing tensor for {label}; tried {candidates}") + + +def _precision_dtypes(precision: str) -> tuple[type[np.generic], object]: + if precision == "fp16": + return np.float16, trt.float16 + if precision == "fp32": + return np.float32, trt.float32 + raise ValueError( + f"Unsupported Parakeet TDT precision {precision!r}; " + "expected fp32 or fp16") + + +def _add_lstm_cell( + network, x, h_prev, c_prev, weights, pfx: str, hidden: int, + dtype=np.float32): + w_ih = graph_ops.add_constant( + network, (hidden, 4 * hidden), weights[f"{pfx}.w_ih_t"], dtype=dtype) + w_hh = graph_ops.add_constant( + network, (hidden, 4 * hidden), weights[f"{pfx}.w_hh_t"], dtype=dtype) + bias = graph_ops.add_constant( + network, (1, 4 * hidden), weights[f"{pfx}.bias"], dtype=dtype) + + xw = network.add_matrix_multiply(x, trt.MatrixOperation.NONE, w_ih, trt.MatrixOperation.NONE) + hw = network.add_matrix_multiply(h_prev, trt.MatrixOperation.NONE, w_hh, trt.MatrixOperation.NONE) + gates = network.add_elementwise(xw.get_output(0), hw.get_output(0), trt.ElementWiseOperation.SUM) + gates = network.add_elementwise(gates.get_output(0), bias, trt.ElementWiseOperation.SUM) + + gate_i = network.add_slice(gates.get_output(0), start=(0, 0), shape=(1, hidden), stride=(1, 1)) + gate_f = network.add_slice(gates.get_output(0), start=(0, hidden), shape=(1, hidden), stride=(1, 1)) + gate_g = network.add_slice( + gates.get_output(0), start=(0, 2 * hidden), shape=(1, hidden), stride=(1, 1)) + gate_o = network.add_slice( + gates.get_output(0), start=(0, 3 * hidden), shape=(1, hidden), stride=(1, 1)) + + i_t = network.add_activation(gate_i.get_output(0), trt.ActivationType.SIGMOID).get_output(0) + f_t = network.add_activation(gate_f.get_output(0), trt.ActivationType.SIGMOID).get_output(0) + g_t = network.add_activation(gate_g.get_output(0), trt.ActivationType.TANH).get_output(0) + o_t = network.add_activation(gate_o.get_output(0), trt.ActivationType.SIGMOID).get_output(0) + + forget = network.add_elementwise(f_t, c_prev, trt.ElementWiseOperation.PROD).get_output(0) + update = network.add_elementwise(i_t, g_t, trt.ElementWiseOperation.PROD).get_output(0) + c_new = network.add_elementwise(forget, update, trt.ElementWiseOperation.SUM).get_output(0) + tanh_c = network.add_activation(c_new, trt.ActivationType.TANH).get_output(0) + h_new = network.add_elementwise(o_t, tanh_c, trt.ElementWiseOperation.PROD).get_output(0) + return h_new, c_new + + +def _build_predictor( + weights: WeightDict, *, precision: str = "fp32", + verbose: bool = False) -> bytes: + pred_hidden = int(weights["_pred_hidden"]) + pred_layers = int(weights["_pred_layers"]) + vocab_total = int(weights["_vocab_total"]) + work_np_dtype, work_trt_dtype = _precision_dtypes(precision) + + logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + config = builder.create_builder_config() + config.builder_optimization_level = 1 + config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 256 << 20) + + token_id = network.add_input("token_id", trt.int32, (1,)) + embedding = graph_ops.add_constant( + network, (vocab_total, pred_hidden), weights["pred_embedding"], + dtype=work_np_dtype) + hidden = network.add_gather(embedding, token_id, 0).get_output(0) + + next_h = [] + next_c = [] + for layer in range(pred_layers): + h_in = network.add_input(f"state_h_{layer}", trt.float32, (1, pred_hidden)) + c_in = network.add_input(f"state_c_{layer}", trt.float32, (1, pred_hidden)) + if work_trt_dtype != trt.float32: + h_in = network.add_cast(h_in, work_trt_dtype).get_output(0) + c_in = network.add_cast(c_in, work_trt_dtype).get_output(0) + hidden, c_new = _add_lstm_cell(network, hidden, h_in, c_in, weights, f"pred.{layer}", + pred_hidden, dtype=work_np_dtype) + next_h.append(hidden) + next_c.append(c_new) + + # Keep pred_output distinct from the last recurrent state even in FP32; + # TensorRT output names belong to tensors, so aliasing would rename the + # same tensor from pred_output to next_h_{last_layer}. + pred_output = network.add_identity(hidden).get_output(0) + if pred_output.dtype != trt.float32: + pred_output = network.add_cast(pred_output, trt.float32).get_output(0) + pred_output.name = "pred_output" + network.mark_output(pred_output) + for layer in range(pred_layers): + h_output = next_h[layer] + c_output = next_c[layer] + if h_output.dtype != trt.float32: + h_output = network.add_cast(h_output, trt.float32).get_output(0) + c_output = network.add_cast(c_output, trt.float32).get_output(0) + h_output.name = f"next_h_{layer}" + c_output.name = f"next_c_{layer}" + network.mark_output(h_output) + network.mark_output(c_output) + + if verbose: + print(f"[trtmc build] Building TDT predictor ({pred_layers}L, h={pred_hidden})", + file=sys.stderr) + plan = builder.build_serialized_network(network, config) + if plan is None: + raise RuntimeError("TDT predictor build failed") + return bytes(plan) + + +def _build_joint( + weights: WeightDict, *, precision: str = "fp32", + verbose: bool = False) -> bytes: + enc_hidden = int(weights["_hidden"]) + pred_hidden = int(weights["_pred_hidden"]) + joint_hidden = int(weights["_joint_hidden"]) + vocab_total = int(weights["_vocab_total"]) + joint_output_size = int(weights["_joint_output_size"]) + duration_count = len(weights["_duration_values"]) + activation = str(weights["_joint_activation"]).lower() + work_np_dtype, work_trt_dtype = _precision_dtypes(precision) + + logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + config = builder.create_builder_config() + config.builder_optimization_level = 1 + config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 256 << 20) + + enc = network.add_input("encoder_frame", trt.float32, (1, enc_hidden)) + pred = network.add_input("pred_output", trt.float32, (1, pred_hidden)) + if work_trt_dtype != trt.float32: + enc = network.add_cast(enc, work_trt_dtype).get_output(0) + pred = network.add_cast(pred, work_trt_dtype).get_output(0) + enc_proj = graph_ops.add_bias_sum( + network, + graph_ops.add_matmul_rhs_constant(network, enc, enc_hidden, joint_hidden, + weights["joint_enc_w"], + dtype=work_np_dtype), + joint_hidden, + weights["joint_enc_b"], + dtype=work_np_dtype, + ) + pred_proj = graph_ops.add_bias_sum( + network, + graph_ops.add_matmul_rhs_constant(network, pred, pred_hidden, joint_hidden, + weights["joint_pred_w"], + dtype=work_np_dtype), + joint_hidden, + weights["joint_pred_b"], + dtype=work_np_dtype, + ) + joint = network.add_elementwise(enc_proj, pred_proj, trt.ElementWiseOperation.SUM).get_output(0) + if activation == "relu": + joint = network.add_activation(joint, trt.ActivationType.RELU).get_output(0) + elif activation == "tanh": + joint = network.add_activation(joint, trt.ActivationType.TANH).get_output(0) + elif activation == "sigmoid": + joint = network.add_activation(joint, trt.ActivationType.SIGMOID).get_output(0) + else: + raise ValueError(f"Unsupported TDT joint activation: {activation}") + + logits = graph_ops.add_bias_sum( + network, + graph_ops.add_matmul_rhs_constant(network, joint, joint_hidden, joint_output_size, + weights["joint_out_w"], + dtype=work_np_dtype), + joint_output_size, + weights["joint_out_b"], + dtype=work_np_dtype, + ) + if logits.dtype != trt.float32: + logits = network.add_cast(logits, trt.float32).get_output(0) + token_slice = network.add_slice(logits, (0, 0), (1, vocab_total), (1, 1)).get_output(0) + duration_slice = network.add_slice( + logits, (0, vocab_total), (1, duration_count), (1, 1)).get_output(0) + token_slice.name = "token_logits" + duration_slice.name = "duration_logits" + network.mark_output(token_slice) + network.mark_output(duration_slice) + + if verbose: + print(f"[trtmc build] Building TDT joint (enc={enc_hidden}, pred={pred_hidden}, " + f"joint={joint_hidden}, vocab={vocab_total})", file=sys.stderr) + plan = builder.build_serialized_network(network, config) + if plan is None: + raise RuntimeError("TDT joint build failed") + return bytes(plan) + + +def _build_mel_filterbank(weights: WeightDict, *, verbose: bool = False) -> bytes | None: + num_mel_bins = int(weights.get("_mel_bins", 128)) + n_fft = 512 + sampling_rate = 16000 + n_freq_bins = 1 + n_fft // 2 + import librosa + filters = librosa.filters.mel( + sr=sampling_rate, n_fft=n_fft, n_mels=num_mel_bins, + fmin=0.0, fmax=sampling_rate / 2.0, norm="slaney", + ).T + filters_flat = np.ascontiguousarray(filters, dtype=np.float32) + header = np.array([n_freq_bins, num_mel_bins], dtype=np.int32) + if verbose: + print(f"[trtmc build] TDT mel filterbank: {n_freq_bins}x{num_mel_bins}", + file=sys.stderr) + return header.tobytes() + filters_flat.tobytes() + + +def _load_weights(model_dir: str, config: ModelConfig, *, precision: str = "fp32") -> tuple[WeightDict, dict]: + del precision + w = WeightDict() + sd, ncfg = _load_hf_as_nemo(model_dir) + + ec = ncfg.get("encoder", {}) + defaults = ncfg.get("model_defaults", {}) + dec_cfg = ncfg.get("decoder", {}) + prednet = dec_cfg.get("prednet", _cfg_dict(dec_cfg, "config_dict", "prednet")) + joint_cfg = ncfg.get("joint", {}) + jointnet = joint_cfg.get("jointnet", _cfg_dict(joint_cfg, "config_dict", "jointnet")) + + hidden = _cfg_int(ec.get("d_model"), defaults.get("enc_hidden"), default=1024) + mel_bins = _cfg_int(ncfg.get("preprocessor", {}).get("features"), ec.get("feat_in"), default=128) + kern = _cfg_int(ec.get("conv_kernel_size"), default=9) + conv_norm_type = str(ec.get("conv_norm_type", "batch_norm")).lower() + conv_context_size = str(ec.get("conv_context_size", "symmetric")).lower() + causal_downsampling = bool(ec.get("causal_downsampling", False)) + enc_heads = _cfg_int(ec.get("n_heads"), default=8) + enc_ffn = _cfg_int(ec.get("ff_expansion_factor"), default=4) * hidden + sub_ch = _cfg_int(ec.get("subsampling_conv_channels"), default=256) + head_dim = hidden // enc_heads + + enc_layers = max(int(k.split(".")[2]) for k in sd if k.startswith("encoder.layers.")) + 1 + mel_length = _cfg_int(config.raw.get("mel_length"), ncfg.get("trtmc_mel_length"), + default=3000) + enc_seq = (_compute_causal_enc_seq_len(mel_length) + if causal_downsampling else _compute_enc_seq_len(mel_length)) + att_contexts = ec.get("att_context_size") or [[-1, -1]] + att_context = att_contexts[0] if isinstance(att_contexts, list) and att_contexts else [70, 13] + att_left = _cfg_int(att_context[0] if len(att_context) > 0 else None, default=70) + att_right = _cfg_int(att_context[1] if len(att_context) > 1 else None, default=13) + + # Streaming knobs are checkpoint-defined. The NeMo config exposes the full + # list of supported att_context_size pairs; drive the per-right-context + # engine set from that list. + streaming_right_contexts = [] + streaming_cache_left = att_left + w["_streaming_right_contexts"] = streaming_right_contexts + w["_streaming_cache_left"] = streaming_cache_left + + pred_hidden = _cfg_int(prednet.get("pred_hidden"), defaults.get("pred_hidden"), default=640) + pred_layers = _cfg_int(prednet.get("pred_rnn_layers"), default=1) + rnn_hidden = _cfg_int(prednet.get("rnn_hidden_size"), default=pred_hidden) + if rnn_hidden != pred_hidden: + raise ValueError( + "Parakeet TDT TDT currently supports predictor LSTMs without " + f"projection (pred_hidden={pred_hidden}, rnn_hidden_size={rnn_hidden})." + ) + joint_hidden = _cfg_int(jointnet.get("joint_hidden"), defaults.get("joint_hidden"), + default=pred_hidden) + joint_activation = str(jointnet.get("activation", "relu")).lower() + + # --- Encoder: same FastConformer tensor layout as the native Parakeet TDT encoder path. --- + w["_enc_layers"] = enc_layers + w["_enc_heads"] = enc_heads + w["_enc_ffn"] = enc_ffn + w["_hidden"] = hidden + w["_mel_bins"] = mel_bins + w["_kern"] = kern + w["_mel_length"] = mel_length + w["_enc_seq"] = enc_seq + w["_sub_ch"] = sub_ch + w["_head_dim"] = head_dim + w["_conv_norm_type"] = conv_norm_type + w["_conv_context_size"] = conv_context_size + w["_causal_downsampling"] = causal_downsampling + w["_encoder_attention_mask_2d"] = True + + w["enc_sub_conv0_w"] = _to_np(sd["encoder.pre_encode.conv.0.weight"]) + w["enc_sub_conv0_b"] = _to_np(sd["encoder.pre_encode.conv.0.bias"]) + for s, (di, pi) in enumerate([(2, 3), (5, 6)]): + w[f"enc_sub_dw{s}_w"] = _to_np(sd[f"encoder.pre_encode.conv.{di}.weight"]) + w[f"enc_sub_dw{s}_b"] = _to_np(sd[f"encoder.pre_encode.conv.{di}.bias"]) + w[f"enc_sub_pw{s}_w"] = _to_np(sd[f"encoder.pre_encode.conv.{pi}.weight"]) + w[f"enc_sub_pw{s}_b"] = _to_np(sd[f"encoder.pre_encode.conv.{pi}.bias"]) + w["enc_sub_out_w"] = _transpose_2d(_to_np(sd["encoder.pre_encode.out.weight"]), "sub") + w["enc_sub_out_b"] = _to_np(sd["encoder.pre_encode.out.bias"]) + + for i in range(enc_layers): + nk = f"encoder.layers.{i}" + pk = f"el.{i}" + for p, n in [("w_q", "linear_q"), ("w_k", "linear_k"), ("w_v", "linear_v"), + ("w_o", "linear_out")]: + w[f"{pk}.{p}"] = _transpose_2d(_to_np(sd[f"{nk}.self_attn.{n}.weight"]), p) + bk = f"{nk}.self_attn.{n}.bias" + w[f"{pk}.b_{p[-1]}"] = _to_np(sd[bk]) if bk in sd else np.zeros(hidden, dtype=np.float32) + w[f"{pk}.pos_bias_u"] = _to_np(sd[f"{nk}.self_attn.pos_bias_u"]) + w[f"{pk}.pos_bias_v"] = _to_np(sd[f"{nk}.self_attn.pos_bias_v"]) + w[f"{pk}.w_pos"] = _transpose_2d(_to_np(sd[f"{nk}.self_attn.linear_pos.weight"]), "pos") + w[f"{pk}.norm_sa"] = _to_np(sd[f"{nk}.norm_self_att.weight"]) + w[f"{pk}.norm_sa_b"] = _to_np(sd[f"{nk}.norm_self_att.bias"]) + for fn, fk in [("ff1", "feed_forward1"), ("ff2", "feed_forward2")]: + w[f"{pk}.{fn}.w1"] = _transpose_2d(_to_np(sd[f"{nk}.{fk}.linear1.weight"]), f"{fn}1") + b1 = f"{nk}.{fk}.linear1.bias" + w[f"{pk}.{fn}.b1"] = _to_np(sd[b1]) if b1 in sd else np.zeros(enc_ffn, dtype=np.float32) + w[f"{pk}.{fn}.w2"] = _transpose_2d(_to_np(sd[f"{nk}.{fk}.linear2.weight"]), f"{fn}2") + b2 = f"{nk}.{fk}.linear2.bias" + w[f"{pk}.{fn}.b2"] = _to_np(sd[b2]) if b2 in sd else np.zeros(hidden, dtype=np.float32) + nm = "norm_feed_forward1" if fn == "ff1" else "norm_feed_forward2" + w[f"{pk}.{fn}.norm"] = _to_np(sd[f"{nk}.{nm}.weight"]) + w[f"{pk}.{fn}.norm_b"] = _to_np(sd[f"{nk}.{nm}.bias"]) + w[f"{pk}.cpw1_w"] = _to_np(sd[f"{nk}.conv.pointwise_conv1.weight"]) + cpw1_b = f"{nk}.conv.pointwise_conv1.bias" + w[f"{pk}.cpw1_b"] = _to_np(sd[cpw1_b]) if cpw1_b in sd else np.zeros(2 * hidden, dtype=np.float32) + w[f"{pk}.cdw_w"] = _to_np(sd[f"{nk}.conv.depthwise_conv.weight"]) + cdw_b = f"{nk}.conv.depthwise_conv.bias" + w[f"{pk}.cdw_b"] = _to_np(sd[cdw_b]) if cdw_b in sd else np.zeros(hidden, dtype=np.float32) + w[f"{pk}.bn_w"] = _to_np(sd[f"{nk}.conv.batch_norm.weight"]) + w[f"{pk}.bn_b"] = _to_np(sd[f"{nk}.conv.batch_norm.bias"]) + w[f"{pk}.bn_m"] = _to_np(sd[f"{nk}.conv.batch_norm.running_mean"]) if f"{nk}.conv.batch_norm.running_mean" in sd else np.zeros(hidden, dtype=np.float32) + w[f"{pk}.bn_v"] = _to_np(sd[f"{nk}.conv.batch_norm.running_var"]) if f"{nk}.conv.batch_norm.running_var" in sd else np.ones(hidden, dtype=np.float32) + w[f"{pk}.cpw2_w"] = _to_np(sd[f"{nk}.conv.pointwise_conv2.weight"]) + cpw2_b = f"{nk}.conv.pointwise_conv2.bias" + w[f"{pk}.cpw2_b"] = _to_np(sd[cpw2_b]) if cpw2_b in sd else np.zeros(hidden, dtype=np.float32) + w[f"{pk}.norm_conv"] = _to_np(sd[f"{nk}.norm_conv.weight"]) + w[f"{pk}.norm_conv_b"] = _to_np(sd[f"{nk}.norm_conv.bias"]) + w[f"{pk}.norm_out"] = _to_np(sd[f"{nk}.norm_out.weight"]) + w[f"{pk}.norm_out_b"] = _to_np(sd[f"{nk}.norm_out.bias"]) + + rpe = _relative_pe(enc_seq, hidden) + for i in range(enc_layers): + proj = rpe @ w[f"el.{i}.w_pos"] + w[f"el.{i}.rpe_proj"] = proj.reshape(2 * enc_seq - 1, enc_heads, head_dim) + + # --- TDT predictor. --- + embed = _to_np(_find_tensor(sd, ["decoder.prediction.embed.weight"], "predictor embedding")) + vocab_total = int(embed.shape[0]) + blank_id = _cfg_int(dec_cfg.get("blank_idx"), default=vocab_total - 1) + if blank_id >= vocab_total: + raise ValueError( + f"TDT blank_id={blank_id} is outside predictor embedding rows={vocab_total}; " + "blank_as_pad=False checkpoints are not supported yet." + ) + vocab = blank_id + w["pred_embedding"] = embed + w["_pred_hidden"] = pred_hidden + w["_pred_layers"] = pred_layers + w["_vocab"] = vocab + w["_vocab_total"] = vocab_total + w["_blank_id"] = blank_id + + for i in range(pred_layers): + pfx = f"pred.{i}" + base_candidates = [ + f"decoder.prediction.dec_rnn.weight_{{kind}}_l{i}", + f"decoder.prediction.dec_rnn.lstm.weight_{{kind}}_l{i}", + f"decoder.prediction.dec_rnn.rnn.weight_{{kind}}_l{i}", + f"decoder.prediction.dec_rnn._rnn.weight_{{kind}}_l{i}", + ] + b_candidates = [ + f"decoder.prediction.dec_rnn.bias_{{kind}}_l{i}", + f"decoder.prediction.dec_rnn.lstm.bias_{{kind}}_l{i}", + f"decoder.prediction.dec_rnn.rnn.bias_{{kind}}_l{i}", + f"decoder.prediction.dec_rnn._rnn.bias_{{kind}}_l{i}", + ] + w_ih = _to_np(_find_tensor(sd, [c.format(kind="ih") for c in base_candidates], + f"predictor layer {i} weight_ih")) + w_hh = _to_np(_find_tensor(sd, [c.format(kind="hh") for c in base_candidates], + f"predictor layer {i} weight_hh")) + b_ih = _to_np(_find_tensor(sd, [c.format(kind="ih") for c in b_candidates], + f"predictor layer {i} bias_ih")) + b_hh = _to_np(_find_tensor(sd, [c.format(kind="hh") for c in b_candidates], + f"predictor layer {i} bias_hh")) + if w_ih.shape != (4 * pred_hidden, pred_hidden) or w_hh.shape != (4 * pred_hidden, pred_hidden): + raise ValueError( + f"Unsupported predictor LSTM layer {i} shapes: " + f"w_ih={w_ih.shape}, w_hh={w_hh.shape}, expected " + f"{(4 * pred_hidden, pred_hidden)}." + ) + w[f"{pfx}.w_ih_t"] = np.ascontiguousarray(w_ih.T.astype(np.float32)) + w[f"{pfx}.w_hh_t"] = np.ascontiguousarray(w_hh.T.astype(np.float32)) + w[f"{pfx}.bias"] = (b_ih + b_hh).astype(np.float32).reshape(1, -1) + + # --- TDT joint. --- + joint_prefix = "joint" + w["joint_enc_w"] = _transpose_2d( + _to_np(_find_tensor(sd, [f"{joint_prefix}.enc.weight"], "joint encoder projection")), + "joint_enc") + w["joint_enc_b"] = _to_np(_find_tensor(sd, [f"{joint_prefix}.enc.bias"], "joint encoder bias")) + w["joint_pred_w"] = _transpose_2d( + _to_np(_find_tensor(sd, [f"{joint_prefix}.pred.weight"], "joint predictor projection")), + "joint_pred") + w["joint_pred_b"] = _to_np(_find_tensor(sd, [f"{joint_prefix}.pred.bias"], "joint predictor bias")) + out_w, out_b = _find_joint_linear(sd, joint_prefix, "joint output") + w["joint_out_w"] = _transpose_2d(_to_np(out_w), "joint_out") + w["joint_out_b"] = _to_np(out_b) + duration_values = tuple(int(x) for x in ncfg.get("tdt_durations", [0, 1, 2, 3, 4])) + joint_output_size = int(w["joint_out_b"].shape[0]) + if joint_output_size != vocab_total + len(duration_values): + raise ValueError( + f"TDT joint output has {joint_output_size} rows; expected " + f"{vocab_total} token rows plus {len(duration_values)} duration rows") + w["_joint_output_size"] = joint_output_size + w["_duration_values"] = duration_values + w["_joint_hidden"] = joint_hidden + w["_joint_activation"] = joint_activation + + config.hidden_size = pred_hidden + config.vocab_size = vocab_total + config.num_hidden_layers = pred_layers + config.num_attention_heads = 1 + config.num_key_value_heads = 1 + + frontend = { + "num_mel_bins": mel_bins, + "max_source_positions": enc_seq, + "encoder_layers": enc_layers, + "mel_length": mel_length, + "subsampling_factor": 8, + "sample_rate": 16000, + } + runtime = { + "hidden_size": pred_hidden, + "num_hidden_layers": pred_layers, + "num_attention_heads": 1, + "num_key_value_heads": 1, + "vocab_size": vocab_total, + "tdt_encoder_hidden_size": hidden, + "tdt_pred_hidden_size": pred_hidden, + "tdt_pred_num_layers": pred_layers, + "tdt_encoder_layers": enc_layers, + "tdt_joint_hidden_size": joint_hidden, + "tdt_vocab_size": vocab, + "tdt_blank_id": blank_id, + "tdt_duration_values": list(duration_values), + "tdt_max_symbols_per_step": _cfg_int(ncfg.get("decoding", {}).get("max_symbols_per_step"), + default=10), + "tdt_causal_downsampling": causal_downsampling, + "tdt_att_context_left": att_left, + "tdt_att_context_right": att_right, + } + return w, {**frontend, **runtime} + + +def compile_engines(model_dir: Path, *, precision: str, verbose: bool): + config = ModelConfig.from_dir(model_dir) + weights, runtime = _load_weights(str(model_dir), config, precision=precision) + mel = _build_mel_filterbank(weights, verbose=verbose) + if mel is None: + raise RuntimeError("Parakeet TDT requires a mel filterbank") + plans = { + "encoder.plan": _build_encoder(config, weights, precision=precision, verbose=verbose), + "predictor.plan": _build_predictor(weights, precision=precision, verbose=verbose), + "joint.plan": _build_joint(weights, precision=precision, verbose=verbose), + "mel_filterbank": mel, + } + runtime.update({ + "tensor_parallel_size": 1, "mel_n_fft": 512, "mel_win_length": 400, + "mel_hop_length": 160, "mel_chunk_length": 30, "mel_sampling_rate": 16000, + "mel_preemph": 0.97, "mel_normalize": "per_feature", + }) + return plans, runtime diff --git a/families/parakeet_tdt/fastconformer.py b/families/parakeet_tdt/fastconformer.py new file mode 100644 index 0000000000..5c990fe1ac --- /dev/null +++ b/families/parakeet_tdt/fastconformer.py @@ -0,0 +1,543 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations +import math +import sys +import numpy as np +import tensorrt as trt +from . import graph_ops + + +def _to_np(t) -> np.ndarray: + if hasattr(t, "numpy"): + return t.detach().cpu().numpy().astype(np.float32) + return np.asarray(t, dtype=np.float32) + + +def _relative_pe(seq_len: int, d_model: int, max_len: int = 5000) -> np.ndarray: + """Compute relative PE for Transformer-XL attention. + + Matches NeMo RelPositionalEncoding: builds table with positive and + negative position encodings, where negative uses sin(-k*d) = -sin(k*d). + """ + if seq_len <= 0: + raise ValueError(f"seq_len must be positive, got {seq_len}") + if seq_len > max_len: + raise ValueError(f"seq_len={seq_len} exceeds relative-position max_len={max_len}") + pos = np.arange(0, max_len, dtype=np.float32)[:, np.newaxis] + div = np.exp(np.arange(0, d_model, 2, dtype=np.float32) * -(math.log(10000.0) / d_model)) + + pe_pos = np.zeros((max_len, d_model), dtype=np.float32) + pe_pos[:, 0::2] = np.sin(pos * div) + pe_pos[:, 1::2] = np.cos(pos * div) + pe_pos = pe_pos[::-1].copy() + + pe_neg = np.zeros((max_len, d_model), dtype=np.float32) + pe_neg[:, 0::2] = np.sin(-pos * div) + pe_neg[:, 1::2] = np.cos(-pos * div) + pe_neg = pe_neg[1:] + + pe_full = np.concatenate([pe_pos, pe_neg], axis=0) + start = max_len - seq_len + end = max_len + seq_len - 1 + return pe_full[start:end] + + +def _compute_enc_seq_len(mel_length: int) -> int: + """Encoder time output after 3 CausalConv2D stride-2 stages. + + Time dim uses symmetric padding (left=1, right=1): (t+2-3)//2+1 = (t-1)//2+1 + """ + t = mel_length + for _ in range(3): + t = (t + 2 - 3) // 2 + 1 + return t + + +def _compute_causal_enc_seq_len(mel_length: int) -> int: + t = mel_length + for _ in range(3): + t = t // 2 + 1 + return t + + +def _build_subsampling( + network, mel_input, weights, sub_ch, hidden, num_mel_bins, mel_length, dtype=np.float32 +): + causal_downsampling = bool(weights.get("_causal_downsampling", False)) + + def add_subsample_conv(inp, weight, bias, out_channels, *, groups=1): + if causal_downsampling: + pad = network.add_padding_nd(inp, pre_padding=(2, 2), post_padding=(1, 1)) + inp = pad.get_output(0) + padding = (0, 0) + else: + padding = (1, 1) + return graph_ops.add_conv2d( + network, + inp, + weight=weight, + bias=bias, + out_channels=out_channels, + kernel_size=(3, 3), + stride=(2, 2), + padding=padding, + groups=groups, + dtype=dtype, + ) + + # NeMo ConformerEncoder passes audio as [B, T, F] to pre_encode. + # MaskedConvSequential.forward unsqueezes to [B, 1, T, F]. + # So Conv2d input is [1, 1, mel_length, mel_bins] (time=H, features=W). + # Our mel_input is [mel_bins, mel_length] = [F, T]. Transpose to [T, F]. + tr_mel = network.add_shuffle(mel_input) + tr_mel.first_transpose = trt.Permutation([1, 0]) # [F,T] -> [T,F] + ri = network.add_shuffle(tr_mel.get_output(0)) + ri.reshape_dims = (1, 1, mel_length, num_mel_bins) # [1, 1, T, F] + x = ri.get_output(0) + # Standard Conv2d with symmetric padding + ReLU (NOT SiLU, NOT causal) + x = add_subsample_conv(x, weights["enc_sub_conv0_w"], weights["enc_sub_conv0_b"], sub_ch) + x = graph_ops.add_activation(network, x, "relu", dtype=dtype) + for s in range(2): + x = add_subsample_conv( + x, weights[f"enc_sub_dw{s}_w"], weights[f"enc_sub_dw{s}_b"], sub_ch, groups=sub_ch + ) + x = graph_ops.add_conv2d( + network, + x, + weight=weights[f"enc_sub_pw{s}_w"], + bias=weights[f"enc_sub_pw{s}_b"], + out_channels=sub_ch, + kernel_size=(1, 1), + dtype=dtype, + ) + x = graph_ops.add_activation(network, x, "relu", dtype=dtype) + # After convs: [1, C, T_out, F_out] where T=time, F=features + time_out = int(weights.get("_enc_seq", _compute_enc_seq_len(mel_length))) + sub_out_in = int(weights["enc_sub_out_w"].shape[0]) + feat_out = sub_out_in // sub_ch + # NeMo: x.transpose(1,2).reshape(B,T,-1) on [B,C,T,F] + # = permute(0,2,1,3) -> [B,T,C,F], reshape -> [T, C*F] + tr = network.add_shuffle(x) + tr.first_transpose = trt.Permutation([0, 2, 1, 3]) # [B,C,T,F] -> [B,T,C,F] + tr.reshape_dims = (time_out, sub_ch * feat_out) # [T, C*F] + out = graph_ops.add_matmul_rhs_constant( + network, tr.get_output(0), sub_ch * feat_out, hidden, weights["enc_sub_out_w"], dtype=dtype + ) + return graph_ops.add_bias_sum(network, out, hidden, weights["enc_sub_out_b"], dtype=dtype) + + +def _rel_shift(network, x, H, S, dtype=np.float32): + zeros = graph_ops.add_constant( + network, (H, S, 1), np.zeros((H, S, 1), dtype=dtype), dtype=dtype + ) + padded = network.add_concatenation([zeros, x]) + padded.axis = 2 + rs1 = network.add_shuffle(padded.get_output(0)) + rs1.reshape_dims = (H, 2 * S, S) + sl1 = network.add_slice( + rs1.get_output(0), start=(0, 1, 0), shape=(H, 2 * S - 1, S), stride=(1, 1, 1) + ) + rs2 = network.add_shuffle(sl1.get_output(0)) + rs2.reshape_dims = (H, S, 2 * S - 1) + sl2 = network.add_slice(rs2.get_output(0), start=(0, 0, 0), shape=(H, S, S), stride=(1, 1, 1)) + return sl2.get_output(0) + + +def _add_rel_pos_attention( + network, hs, weights, pfx, hidden, H, D, S, rel_pe_proj, eps, enc_mask=None, dtype=np.float32 +): + normed = graph_ops.add_layer_norm( + network, + hs, + hidden, + weights[f"{pfx}.norm_sa"], + weights[f"{pfx}.norm_sa_b"], + eps, + dtype=dtype, + ) + q = graph_ops.add_bias_sum( + network, + graph_ops.add_matmul_rhs_constant( + network, normed, hidden, hidden, weights[f"{pfx}.w_q"], dtype=dtype + ), + hidden, + weights[f"{pfx}.b_q"], + dtype=dtype, + ) + k = graph_ops.add_bias_sum( + network, + graph_ops.add_matmul_rhs_constant( + network, normed, hidden, hidden, weights[f"{pfx}.w_k"], dtype=dtype + ), + hidden, + weights[f"{pfx}.b_k"], + dtype=dtype, + ) + v = graph_ops.add_bias_sum( + network, + graph_ops.add_matmul_rhs_constant( + network, normed, hidden, hidden, weights[f"{pfx}.w_v"], dtype=dtype + ), + hidden, + weights[f"{pfx}.b_v"], + dtype=dtype, + ) + qr = network.add_shuffle(q) + qr.reshape_dims = (S, H, D) + kr = network.add_shuffle(k) + kr.reshape_dims = (S, H, D) + vr = network.add_shuffle(v) + vr.reshape_dims = (S, H, D) + bu = graph_ops.add_constant(network, (1, H, D), weights[f"{pfx}.pos_bias_u"], dtype=dtype) + bv = graph_ops.add_constant(network, (1, H, D), weights[f"{pfx}.pos_bias_v"], dtype=dtype) + qu = network.add_elementwise(qr.get_output(0), bu, trt.ElementWiseOperation.SUM).get_output(0) + qv = network.add_elementwise(qr.get_output(0), bv, trt.ElementWiseOperation.SUM).get_output(0) + qu_t = network.add_shuffle(qu) + qu_t.first_transpose = trt.Permutation([1, 0, 2]) + qv_t = network.add_shuffle(qv) + qv_t.first_transpose = trt.Permutation([1, 0, 2]) + k_t = network.add_shuffle(kr.get_output(0)) + k_t.first_transpose = trt.Permutation([1, 0, 2]) + v_t = network.add_shuffle(vr.get_output(0)) + v_t.first_transpose = trt.Permutation([1, 0, 2]) + cs = network.add_matrix_multiply( + qu_t.get_output(0), + trt.MatrixOperation.NONE, + k_t.get_output(0), + trt.MatrixOperation.TRANSPOSE, + ).get_output(0) + rp_t = network.add_shuffle(rel_pe_proj) + rp_t.first_transpose = trt.Permutation([1, 0, 2]) + ps_raw = network.add_matrix_multiply( + qv_t.get_output(0), + trt.MatrixOperation.NONE, + rp_t.get_output(0), + trt.MatrixOperation.TRANSPOSE, + ).get_output(0) + ps = _rel_shift(network, ps_raw, H, S, dtype=dtype) + total = network.add_elementwise(cs, ps, trt.ElementWiseOperation.SUM).get_output(0) + sc = graph_ops.add_constant( + network, (1, 1, 1), np.array([1.0 / math.sqrt(D)], dtype=dtype), dtype=dtype + ) + scaled = network.add_elementwise(total, sc, trt.ElementWiseOperation.PROD).get_output(0) + # Apply encoder sequence mask: [1, 1, S] added to scores [H, S, S] + if enc_mask is not None: + scaled = network.add_elementwise(scaled, enc_mask, trt.ElementWiseOperation.SUM).get_output( + 0 + ) + # Conformer relative-position attention uses a rel-shifted Q*R term in + # the logits, which native IAttention cannot represent as a plain mask. + sm = network.add_softmax(scaled) + sm.axes = 1 << 2 + ao = network.add_matrix_multiply( + sm.get_output(0), trt.MatrixOperation.NONE, v_t.get_output(0), trt.MatrixOperation.NONE + ).get_output(0) + at = network.add_shuffle(ao) + at.first_transpose = trt.Permutation([1, 0, 2]) + af = network.add_shuffle(at.get_output(0)) + af.reshape_dims = (S, hidden) + return graph_ops.add_bias_sum( + network, + graph_ops.add_matmul_rhs_constant( + network, af.get_output(0), hidden, hidden, weights[f"{pfx}.w_o"], dtype=dtype + ), + hidden, + weights[f"{pfx}.b_o"], + dtype=dtype, + ) + + +def _add_causal_depthwise_conv1d(network, x, weights, pfx, hidden, kern, dtype=np.float32): + pad = kern - 1 + if pad > 0: + zeros = graph_ops.add_constant( + network, (1, hidden, pad), np.zeros((1, hidden, pad), dtype=dtype), dtype=dtype + ) + cat = network.add_concatenation([zeros, x]) + cat.axis = 2 + x = cat.get_output(0) + return graph_ops.add_conv1d( + network, + x, + weight=weights[f"{pfx}.cdw_w"], + bias=weights[f"{pfx}.cdw_b"], + out_channels=hidden, + kernel_size=kern, + groups=hidden, + dtype=dtype, + ) + + +def _add_conv_norm(network, x, weights, pfx, hidden, S, eps, conv_norm_type, dtype=np.float32): + if conv_norm_type == "layer_norm": + r1 = network.add_shuffle(x) + r1.reshape_dims = (hidden, S) + r2 = network.add_shuffle(r1.get_output(0)) + r2.first_transpose = trt.Permutation([1, 0]) + normed = graph_ops.add_layer_norm( + network, + r2.get_output(0), + hidden, + weights[f"{pfx}.bn_w"], + weights[f"{pfx}.bn_b"], + eps, + dtype=dtype, + ) + r3 = network.add_shuffle(normed) + r3.first_transpose = trt.Permutation([1, 0]) + r4 = network.add_shuffle(r3.get_output(0)) + r4.reshape_dims = (1, hidden, S) + return r4.get_output(0) + + bn = network.add_shuffle(x) + bn.reshape_dims = (1, hidden, 1, S) + x = graph_ops.add_batch_norm_2d( + network, + bn.get_output(0), + hidden, + gamma=weights[f"{pfx}.bn_w"], + beta=weights[f"{pfx}.bn_b"], + running_mean=weights[f"{pfx}.bn_m"], + running_var=weights[f"{pfx}.bn_v"], + dtype=dtype, + ) + bo = network.add_shuffle(x) + bo.reshape_dims = (1, hidden, S) + return bo.get_output(0) + + +def _add_conv_module( + network, + hs, + weights, + pfx, + hidden, + kern, + S, + eps, + conv_norm_type="batch_norm", + conv_context_size="symmetric", + dtype=np.float32, +): + normed = graph_ops.add_layer_norm( + network, + hs, + hidden, + weights[f"{pfx}.norm_conv"], + weights[f"{pfx}.norm_conv_b"], + eps, + dtype=dtype, + ) + r1 = network.add_shuffle(normed) + r1.first_transpose = trt.Permutation([1, 0]) + r2 = network.add_shuffle(r1.get_output(0)) + r2.reshape_dims = (1, hidden, S) + x = graph_ops.add_conv1d( + network, + r2.get_output(0), + weight=weights[f"{pfx}.cpw1_w"], + bias=weights[f"{pfx}.cpw1_b"], + out_channels=2 * hidden, + kernel_size=1, + dtype=dtype, + ) + xa = network.add_slice(x, start=(0, 0, 0), shape=(1, hidden, S), stride=(1, 1, 1)).get_output(0) + xb = network.add_slice( + x, start=(0, hidden, 0), shape=(1, hidden, S), stride=(1, 1, 1) + ).get_output(0) + gate = network.add_activation(xb, trt.ActivationType.SIGMOID).get_output(0) + x = network.add_elementwise(xa, gate, trt.ElementWiseOperation.PROD).get_output(0) + if conv_context_size == "causal": + x = _add_causal_depthwise_conv1d(network, x, weights, pfx, hidden, kern, dtype=dtype) + else: + x = graph_ops.add_conv1d( + network, + x, + weight=weights[f"{pfx}.cdw_w"], + bias=weights[f"{pfx}.cdw_b"], + out_channels=hidden, + kernel_size=kern, + padding=kern // 2, + groups=hidden, + dtype=dtype, + ) + x = _add_conv_norm(network, x, weights, pfx, hidden, S, eps, conv_norm_type, dtype=dtype) + x = graph_ops.add_activation(network, x, "silu", dtype=dtype) + x = graph_ops.add_conv1d( + network, + x, + weight=weights[f"{pfx}.cpw2_w"], + bias=weights[f"{pfx}.cpw2_b"], + out_channels=hidden, + kernel_size=1, + dtype=dtype, + ) + r3 = network.add_shuffle(x) + r3.reshape_dims = (hidden, S) + r4 = network.add_shuffle(r3.get_output(0)) + r4.first_transpose = trt.Permutation([1, 0]) + return r4.get_output(0) + + +def _add_half_ffn(network, hs, weights, pfx, hidden, ffn, eps, dtype=np.float32): + normed = graph_ops.add_layer_norm( + network, hs, hidden, weights[f"{pfx}.norm"], weights[f"{pfx}.norm_b"], eps, dtype=dtype + ) + fc1 = graph_ops.add_bias_sum( + network, + graph_ops.add_matmul_rhs_constant( + network, normed, hidden, ffn, weights[f"{pfx}.w1"], dtype=dtype + ), + ffn, + weights[f"{pfx}.b1"], + dtype=dtype, + ) + act = graph_ops.add_activation(network, fc1, "silu", dtype=dtype) + fc2 = graph_ops.add_bias_sum( + network, + graph_ops.add_matmul_rhs_constant( + network, act, ffn, hidden, weights[f"{pfx}.w2"], dtype=dtype + ), + hidden, + weights[f"{pfx}.b2"], + dtype=dtype, + ) + half = graph_ops.add_constant(network, (1, 1), np.array([0.5], dtype=dtype), dtype=dtype) + return network.add_elementwise(fc2, half, trt.ElementWiseOperation.PROD).get_output(0) + + +def _add_conformer_block( + network, + hs, + weights, + pfx, + hidden, + H, + D, + ffn, + kern, + S, + rpe, + eps, + enc_mask=None, + conv_norm_type="batch_norm", + conv_context_size="symmetric", + dtype=np.float32, +): + ffn1 = _add_half_ffn(network, hs, weights, f"{pfx}.ff1", hidden, ffn, eps, dtype=dtype) + hs = network.add_elementwise(hs, ffn1, trt.ElementWiseOperation.SUM).get_output(0) + attn = _add_rel_pos_attention( + network, hs, weights, pfx, hidden, H, D, S, rpe, eps, enc_mask, dtype=dtype + ) + hs = network.add_elementwise(hs, attn, trt.ElementWiseOperation.SUM).get_output(0) + conv = _add_conv_module( + network, + hs, + weights, + pfx, + hidden, + kern, + S, + eps, + conv_norm_type=conv_norm_type, + conv_context_size=conv_context_size, + dtype=dtype, + ) + hs = network.add_elementwise(hs, conv, trt.ElementWiseOperation.SUM).get_output(0) + ffn2 = _add_half_ffn(network, hs, weights, f"{pfx}.ff2", hidden, ffn, eps, dtype=dtype) + hs = network.add_elementwise(hs, ffn2, trt.ElementWiseOperation.SUM).get_output(0) + return graph_ops.add_layer_norm( + network, + hs, + hidden, + weights[f"{pfx}.norm_out"], + weights[f"{pfx}.norm_out_b"], + eps, + dtype=dtype, + ) + + +def _build_encoder(config, weights, *, precision="fp32", verbose=False): + el = weights["_enc_layers"] + eh = weights["_enc_heads"] + h = weights["_hidden"] + hd = weights["_head_dim"] + ef = weights["_enc_ffn"] + k = weights["_kern"] + mb = weights["_mel_bins"] + ml = weights["_mel_length"] + es = weights["_enc_seq"] + sc = weights["_sub_ch"] + conv_norm_type = str(weights.get("_conv_norm_type", "batch_norm")).lower() + conv_context_size = str(weights.get("_conv_context_size", "symmetric")).lower() + if precision == "fp16": + work_np_dtype, work_trt_dtype = np.float16, trt.float16 + elif precision == "fp32": + work_np_dtype, work_trt_dtype = np.float32, trt.float32 + else: + raise ValueError( + f"Unsupported streaming encoder precision {precision!r}; expected fp32 or fp16" + ) + + log = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) + b = trt.Builder(log) + net = b.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + tc = b.create_builder_config() + tc.builder_optimization_level = 1 + tc.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 2 << 30) + + eps = graph_ops.add_constant( + net, (1, 1), np.array([1e-5], dtype=work_np_dtype), dtype=work_np_dtype + ) + mel = net.add_input("mel_features", trt.float32, (mb, ml)) + # Encoder attention mask: [1, 1, enc_seq] -- 0.0 valid, -10000.0 padded. + # Applied additively to self-attention scores before softmax. + mask_shape = ( + (1, es, es) if bool(weights.get("_encoder_attention_mask_2d", False)) else (1, 1, es) + ) + enc_mask = net.add_input("encoder_mask", trt.float32, mask_shape) + if work_trt_dtype != trt.float32: + mel = net.add_cast(mel, work_trt_dtype).get_output(0) + enc_mask = net.add_cast(enc_mask, work_trt_dtype).get_output(0) + + hs = _build_subsampling(net, mel, weights, sc, h, mb, ml, dtype=work_np_dtype) + for li in range(el): + pfx = f"el.{li}" + rpe = graph_ops.add_constant( + net, (2 * es - 1, eh, hd), weights[f"{pfx}.rpe_proj"], dtype=work_np_dtype + ) + hs = _add_conformer_block( + net, + hs, + weights, + pfx, + h, + eh, + hd, + ef, + k, + es, + rpe, + eps, + enc_mask, + conv_norm_type=conv_norm_type, + conv_context_size=conv_context_size, + dtype=work_np_dtype, + ) + + output = hs + if output.dtype != trt.float32: + output = net.add_cast(output, trt.float32).get_output(0) + output.name = "encoder_output" + net.mark_output(output) + if verbose: + print( + f"[trtmc build] Building Parakeet TDT encoder ({el}L, h={h}, heads={eh}, seq={es})", + file=sys.stderr, + ) + plan = b.build_serialized_network(net, tc) + if plan is None: + raise RuntimeError("Parakeet TDT encoder build failed") + return bytes(plan) diff --git a/families/parakeet_tdt/graph_ops.py b/families/parakeet_tdt/graph_ops.py new file mode 100644 index 0000000000..2c7a4701a3 --- /dev/null +++ b/families/parakeet_tdt/graph_ops.py @@ -0,0 +1,423 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Family-owned TensorRT graph operations for Python engine builds. + +Tensor names and shapes must stay compatible with the C++ bundle runtime. +""" + +from __future__ import annotations + +import numpy as np +import tensorrt as trt + + +def _cast_back_to_trt_dtype( + network: trt.INetworkDefinition, + tensor: trt.ITensor, + target_dtype: trt.DataType, +) -> trt.ITensor: + """Cast a tensor back to the original TRT runtime dtype after FP32 compute.""" + if tensor.dtype == target_dtype: + return tensor + return network.add_cast(tensor, target_dtype).get_output(0) + + +def add_constant( + network: trt.INetworkDefinition, + shape: tuple[int, ...], + values: np.ndarray, + dtype: np.dtype = np.float32, +) -> trt.ITensor: + """Add a constant tensor in the given *dtype* (default float32).""" + weights = trt.Weights(np.ascontiguousarray(values, dtype=dtype)) + layer = network.add_constant(shape, weights) + return layer.get_output(0) + + +def add_matmul_rhs_constant( + network: trt.INetworkDefinition, + lhs: trt.ITensor, + lhs_width: int, + rhs_width: int, + rhs_weights: np.ndarray, + dtype: np.dtype = np.float32, +) -> trt.ITensor: + """Matrix multiply: lhs @ rhs_constant. rhs is [lhs_width, rhs_width].""" + rank = len(tuple(lhs.shape)) + rhs_shape = (lhs_width, rhs_width) if rank <= 2 else (1,) * (rank - 2) + (lhs_width, rhs_width) + rhs = add_constant( + network, + rhs_shape, + np.asarray(rhs_weights).reshape(rhs_shape), + dtype=dtype, + ) + rhs = _cast_back_to_trt_dtype(network, rhs, lhs.dtype) + mm = network.add_matrix_multiply( + lhs, + trt.MatrixOperation.NONE, + rhs, + trt.MatrixOperation.NONE, + ) + return _cast_back_to_trt_dtype(network, mm.get_output(0), lhs.dtype) + + +def add_bias_sum( + network: trt.INetworkDefinition, + inp: trt.ITensor, + width: int, + bias: np.ndarray, + dtype: np.dtype = np.float32, +) -> trt.ITensor: + """Element-wise add a bias broadcast over all non-feature axes.""" + rank = len(tuple(inp.shape)) + bias_shape = (width,) if rank <= 1 else (1,) * (rank - 1) + (width,) + bias_t = add_constant(network, bias_shape, np.asarray(bias).reshape(bias_shape), dtype=dtype) + bias_t = _cast_back_to_trt_dtype(network, bias_t, inp.dtype) + s = network.add_elementwise(inp, bias_t, trt.ElementWiseOperation.SUM) + return _cast_back_to_trt_dtype(network, s.get_output(0), inp.dtype) + + +def add_layer_norm( + network: trt.INetworkDefinition, + inp: trt.ITensor, + hidden_size: int, + gamma: np.ndarray, + beta: np.ndarray, + eps_tensor: trt.ITensor, + dtype: np.dtype = np.float32, +) -> trt.ITensor: + """LayerNorm: gamma * ((x - mean) / sqrt(var + eps)) + beta. + + FP32 precision boundary: when dtype != float32, casts to FP32 before + norm computation for numerical stability, then casts back. + """ + need_cast = dtype != np.float32 + output_dtype = inp.dtype + if need_cast: + inp = network.add_cast(inp, trt.float32).get_output(0) + eps_tensor = network.add_cast(eps_tensor, trt.float32).get_output(0) + # mean = reduce_mean(x) + mean = network.add_reduce(inp, trt.ReduceOperation.AVG, 1 << 1, keep_dims=True) + # x - mean + centered = network.add_elementwise(inp, mean.get_output(0), trt.ElementWiseOperation.SUB) + # variance = mean((x - mean)^2) + sq = network.add_elementwise( + centered.get_output(0), centered.get_output(0), trt.ElementWiseOperation.PROD + ) + var = network.add_reduce(sq.get_output(0), trt.ReduceOperation.AVG, 1 << 1, keep_dims=True) + # sqrt(var + eps) + denom_in = network.add_elementwise(var.get_output(0), eps_tensor, trt.ElementWiseOperation.SUM) + sqrt_l = network.add_unary(denom_in.get_output(0), trt.UnaryOperation.SQRT) + recip = network.add_unary(sqrt_l.get_output(0), trt.UnaryOperation.RECIP) + # normalized = (x - mean) / sqrt(var + eps) + normalized = network.add_elementwise( + centered.get_output(0), recip.get_output(0), trt.ElementWiseOperation.PROD + ) + # gamma * normalized + beta + gamma_t = add_constant(network, (1, hidden_size), gamma, dtype=np.float32) + scaled = network.add_elementwise( + normalized.get_output(0), gamma_t, trt.ElementWiseOperation.PROD + ) + beta_t = add_constant(network, (1, hidden_size), beta, dtype=np.float32) + result = network.add_elementwise(scaled.get_output(0), beta_t, trt.ElementWiseOperation.SUM) + result = result.get_output(0) + if need_cast: + result = _cast_back_to_trt_dtype(network, result, output_dtype) + return result + + +def add_gelu_new( + network: trt.INetworkDefinition, + inp: trt.ITensor, + dtype: np.dtype = np.float32, +) -> trt.ITensor: + """GELU (tanh approximation): 0.5*x*(1+tanh(sqrt(2/pi)*(x+0.044715*x^3))). + + Constants are cast to ``inp.dtype`` so the elementwise ops are valid in + a STRONGLY_TYPED network when ``inp`` is bf16 (storage np_dtype is + fp16, runtime trt_dtype is bfloat16) or any other non-matching combo. + """ + target_dtype = inp.dtype + const_shape = (1,) * max(1, len(tuple(inp.shape))) + + def _const(name, value): + c = add_constant(network, const_shape, np.array([value], dtype=np.float32), dtype=dtype) + return _cast_back_to_trt_dtype(network, c, target_dtype) + + # x^3 + x_sq = network.add_elementwise(inp, inp, trt.ElementWiseOperation.PROD) + x_cu = network.add_elementwise(x_sq.get_output(0), inp, trt.ElementWiseOperation.PROD) + # 0.044715 * x^3 + coeff = _const("coeff", 0.044715) + scaled_cube = network.add_elementwise(x_cu.get_output(0), coeff, trt.ElementWiseOperation.PROD) + # x + 0.044715 * x^3 + inner_sum = network.add_elementwise( + inp, scaled_cube.get_output(0), trt.ElementWiseOperation.SUM + ) + # sqrt(2/pi) * (x + 0.044715 * x^3) + sqrt_2_over_pi = _const("sqrt_2_over_pi", np.sqrt(2.0 / np.pi)) + tanh_arg = network.add_elementwise( + sqrt_2_over_pi, inner_sum.get_output(0), trt.ElementWiseOperation.PROD + ) + # tanh(...) + tanh_l = network.add_activation(tanh_arg.get_output(0), trt.ActivationType.TANH) + # 1 + tanh(...) + one = _const("one", 1.0) + one_plus_tanh = network.add_elementwise(one, tanh_l.get_output(0), trt.ElementWiseOperation.SUM) + # 0.5 * x + half = _const("half", 0.5) + half_x = network.add_elementwise(half, inp, trt.ElementWiseOperation.PROD) + # 0.5 * x * (1 + tanh(...)) + result = network.add_elementwise( + half_x.get_output(0), one_plus_tanh.get_output(0), trt.ElementWiseOperation.PROD + ) + return result.get_output(0) + + +def add_activation( + network: trt.INetworkDefinition, + inp: trt.ITensor, + activation_type: str, + dtype: np.dtype = np.float32, +) -> trt.ITensor: + """Dispatch activation by name: 'silu', 'gelu_new', 'gelu', 'relu', 'relu2'/'squared_relu'.""" + if activation_type in ("gelu_new", "gelu"): + return add_gelu_new(network, inp, dtype=dtype) + elif activation_type == "relu": + act = network.add_activation(inp, trt.ActivationType.RELU) + return act.get_output(0) + elif activation_type in ("relu2", "squared_relu"): + relu = network.add_activation(inp, trt.ActivationType.RELU) + sq = network.add_elementwise( + relu.get_output(0), relu.get_output(0), trt.ElementWiseOperation.PROD + ) + return sq.get_output(0) + elif activation_type == "silu": + sigmoid = network.add_activation(inp, trt.ActivationType.SIGMOID) + swish = network.add_elementwise(inp, sigmoid.get_output(0), trt.ElementWiseOperation.PROD) + return swish.get_output(0) + else: + raise ValueError(f"Unsupported activation: {activation_type}") + + +# --------------------------------------------------------------------------- +# Conv / Norm / Resize ops for segmentation and audio models +# --------------------------------------------------------------------------- + + +def add_conv2d( + network: trt.INetworkDefinition, + inp: trt.ITensor, + weight: np.ndarray, + bias: np.ndarray | None, + out_channels: int, + kernel_size: tuple[int, int], + stride: tuple[int, int] = (1, 1), + padding: tuple[int, int] = (0, 0), + groups: int = 1, + dtype: np.dtype = np.float32, +) -> trt.ITensor: + """2D convolution wrapper. + + Input: [N, C_in, H, W] + Weight: [C_out, C_in/groups, kH, kW] + Output: [N, C_out, H', W'] + """ + conv_w = trt.Weights(np.ascontiguousarray(weight, dtype=dtype)) + conv_b = trt.Weights() + if bias is not None: + conv_b = trt.Weights(np.ascontiguousarray(bias, dtype=dtype)) + + conv = network.add_convolution_nd( + inp, + num_output_maps=out_channels, + kernel_shape=kernel_size, + kernel=conv_w, + bias=conv_b, + ) + conv.stride_nd = stride + conv.padding_nd = padding + conv.num_groups = groups + return conv.get_output(0) + + +def add_batch_norm_2d( + network: trt.INetworkDefinition, + inp: trt.ITensor, + num_channels: int, + gamma: np.ndarray, + beta: np.ndarray, + running_mean: np.ndarray, + running_var: np.ndarray, + eps: float = 1e-5, + dtype: np.dtype = np.float32, +) -> trt.ITensor: + """Fused BatchNorm2d: gamma * (x - mean) / sqrt(var + eps) + beta. + + Input: [N, C, H, W] + Output: same shape + + FP32 precision boundary: when dtype != float32, casts to FP32 before + norm computation for numerical stability, then casts back. + """ + need_cast = dtype != np.float32 + output_dtype = inp.dtype + if need_cast: + inp = network.add_cast(inp, trt.float32).get_output(0) + # Fuse into scale + shift + scale = gamma / np.sqrt(running_var + eps) + shift = beta - running_mean * scale + + scale_t = add_constant( + network, + (1, num_channels, 1, 1), + scale.reshape(1, -1, 1, 1).astype(np.float32), + dtype=np.float32, + ) + shift_t = add_constant( + network, + (1, num_channels, 1, 1), + shift.reshape(1, -1, 1, 1).astype(np.float32), + dtype=np.float32, + ) + + scaled = network.add_elementwise(inp, scale_t, trt.ElementWiseOperation.PROD) + result = network.add_elementwise( + scaled.get_output(0), shift_t, trt.ElementWiseOperation.SUM + ).get_output(0) + if need_cast: + result = _cast_back_to_trt_dtype(network, result, output_dtype) + return result + + +def add_conv1d( + network: trt.INetworkDefinition, + inp: trt.ITensor, + weight: np.ndarray, + bias: np.ndarray | None, + out_channels: int, + kernel_size: int, + stride: int = 1, + padding: int = 0, + groups: int = 1, + dtype: np.dtype = np.float32, +) -> trt.ITensor: + """1D convolution via 2D convolution with height=1. + + Input: [N, C_in, L] + Weight: [C_out, C_in/groups, K] + Output: [N, C_out, L'] + """ + # Reshape to [N, C_in, 1, L] + n, c_in, length = inp.shape + reshape_in = network.add_shuffle(inp) + reshape_in.reshape_dims = (n, c_in, 1, length) + + # Weight: [C_out, C_in/groups, K] -> [C_out, C_in/groups, 1, K] + w_4d = weight.reshape(out_channels, -1, 1, kernel_size) + result = add_conv2d( + network, + reshape_in.get_output(0), + w_4d, + bias, + out_channels, + kernel_size=(1, kernel_size), + stride=(1, stride), + padding=(0, padding), + groups=groups, + dtype=dtype, + ) + + # Reshape back to [N, C_out, L'] + out_length = result.shape[3] + reshape_out = network.add_shuffle(result) + reshape_out.reshape_dims = (n, out_channels, out_length) + return reshape_out.get_output(0) + + +# Alias: add_gelu_tanh is the same as add_gelu_new (tanh approximation) +add_gelu_tanh = add_gelu_new + + +def add_attention_core( + network: trt.INetworkDefinition, + q_4d: trt.ITensor, + k_4d: trt.ITensor, + v_4d: trt.ITensor, + causal: bool = False, + mask: trt.ITensor | None = None, + scale: float | None = None, + fp32_accumulation: bool = False, +) -> trt.ITensor: + """Scaled dot-product attention via TRT native IAttention layer. + + Replaces the manual Q@K^T -> scale -> softmax -> @V chain. TRT 10 fuses + this into a single kernel when a compatible implementation is available; + decomposable=True ensures a correct fallback to primitives otherwise. + + NOTE: TRT IAttention computes raw BMM1 = Q @ K^T without any built-in + 1/sqrt(D) scaling. We pre-scale Q by 1/sqrt(D) so that the fused kernel + computes the standard scaled dot-product attention formula. + + Args: + q_4d: Query [B, H, q_seq, D]. + k_4d: Key [B, H, kv_seq, D]. + v_4d: Value [B, H, kv_seq, D]. + causal: Apply causal (autoregressive) mask. Mutually exclusive + with ``mask``. + mask: Optional additive float mask [B, H, q_seq, kv_seq] added + to scaled logits before softmax. Cannot be used with + causal=True. + scale: Optional Q pre-scale factor. Defaults to 1/sqrt(D). + fp32_accumulation: + Cast Q/K/V to FP32 before IAttention, then cast the context + back to the original Q dtype. TRT may still select a + Half-input fused MHA tactic after optimizing the casts, while + keeping the IAttention accumulation/output boundary in FP32. + + Returns: + Context tensor [B, H, q_seq, D]. + """ + output_dtype = q_4d.dtype + if fp32_accumulation and output_dtype != trt.float32: + q_4d = network.add_cast(q_4d, trt.float32).get_output(0) + k_4d = network.add_cast(k_4d, trt.float32).get_output(0) + v_4d = network.add_cast(v_4d, trt.float32).get_output(0) + if mask is not None and mask.dtype != trt.float32: + mask = network.add_cast(mask, trt.float32).get_output(0) + + # Pre-scale Q: TRT IAttention does not apply score scaling itself. + # Match the scale constant's dtype to Q's dtype: in strongly-typed networks + # a FP32 constant mixed with a FP16/BF16 Q causes add_elementwise to emit + # a type-mismatch error and produce a tensor with corrupted dimensions, + # which makes add_attention return None. + if scale is None: + head_dim = q_4d.shape[-1] + scale = float(1.0 / np.sqrt(head_dim)) if head_dim > 0 else 1.0 + # Use FP16 weights directly for FP16; BF16 has no numpy native type so + # create as FP32 and cast; FP32 falls through to the default. + scale_np_dtype = np.float16 if q_4d.dtype == trt.float16 else np.float32 + scale_t = add_constant(network, (1, 1, 1, 1), np.array([[[[scale]]]]), dtype=scale_np_dtype) + if q_4d.dtype == trt.bfloat16: + scale_t = network.add_cast(scale_t, trt.bfloat16).get_output(0) + q_scaled = network.add_elementwise(q_4d, scale_t, trt.ElementWiseOperation.PROD) + + attn = network.add_attention( + q_scaled.get_output(0), + k_4d, + v_4d, + trt.AttentionNormalizationOp.SOFTMAX, + causal, + ) + # Allow TRT to decompose into primitive ops when no fused kernel is + # available (e.g. unsupported head-dim or dtype). This guarantees + # correctness on any configuration at the cost of potential performance. + attn.decomposable = True + if mask is not None and not causal: + attn.mask = mask + return _cast_back_to_trt_dtype(network, attn.get_output(0), output_dtype) + + +# Backward-compatible name used by existing tests and call sites. +_add_attention_core = add_attention_core diff --git a/families/parakeet_tdt/model.py b/families/parakeet_tdt/model.py new file mode 100644 index 0000000000..d526901a29 --- /dev/null +++ b/families/parakeet_tdt/model.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Build the native Transformers Parakeet TDT checkpoint into owned sections.""" + +from pathlib import Path + +from .config import ParakeetTDTConfig + + +def build(request, writer): + if request.dynamic_kv_cache: + raise NotImplementedError("parakeet_tdt does not support dynamic_kv_cache") + required = ( + ("family", request.family, "parakeet_tdt"), + ("task", request.task, "speech_transcription"), + ("backend", request.backend, "trt"), + ("tensor_parallel_size", request.tensor_parallel_size, 1), + ("context_parallel_size", request.context_parallel_size, 1), + ("max_batch_size", request.max_batch_size, 1), + ("max_sequence_length", request.max_sequence_length, None), + ("image_height", request.image_height, None), + ("image_width", request.image_width, None), + ("video_num_frames", request.video_num_frames, None), + ("fp32_layers", request.fp32_layers, ()), + ("graph_transform", request.graph_transform, None), + ) + for name, actual, expected in required: + if actual != expected: + raise ValueError(f"parakeet_tdt requires {name}={expected!r}") + if request.precision not in {"fp16", "fp32"}: + raise ValueError("parakeet_tdt supports precision=fp16 or fp32") + if request.quantization not in {None, "none"}: + raise ValueError("parakeet_tdt does not support quantization") + + model_dir = Path(request.model_dir) + config = ParakeetTDTConfig.from_dir(model_dir) + config.validate_supported_checkpoint() + for filename in ("model.safetensors", "tokenizer.json"): + if not (model_dir / filename).is_file(): + raise FileNotFoundError(f"parakeet_tdt requires {model_dir / filename}") + + from .engines import compile_engines + + plans, runtime = compile_engines( + model_dir, precision=request.precision, verbose=request.verbose, + ) + writer.set_header(family="parakeet_tdt", task=request.task, backend=request.backend) + writer.add_json("runtime.json", runtime) + for name, data in plans.items(): + writer.add_bytes(name, data) + for filename in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"): + path = model_dir / filename + if path.is_file(): + writer.add_bytes(filename, path.read_bytes()) diff --git a/families/parakeet_tdt/model_config.py b/families/parakeet_tdt/model_config.py new file mode 100644 index 0000000000..72d9d60f0a --- /dev/null +++ b/families/parakeet_tdt/model_config.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Family-private build configuration for Parakeet TDT.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from pathlib import Path + + +@dataclass +class ModelConfig: + """Minimal engine-builder view of the Parakeet TDT configuration.""" + + model_type: str = "" + architectures: list[str] = field(default_factory=list) + vocab_size: int = 0 + hidden_size: int = 0 + intermediate_size: int = 0 + num_hidden_layers: int = 0 + num_attention_heads: int = 1 + num_key_value_heads: int = 1 + rms_norm_eps: float = 1e-5 + rope_theta: float = 10000.0 + bos_token_id: int = -1 + eos_token_id: int = -1 + pad_token_id: int = -1 + tie_word_embeddings: bool = False + max_position_embeddings: int = 8192 + hidden_act: str = "" + _head_dim: int = 0 + raw: dict = field(default_factory=dict, repr=False) + + @property + def head_dim(self) -> int: + if self._head_dim > 0: + return self._head_dim + return self.hidden_size // self.num_attention_heads if self.num_attention_heads > 0 else 0 + + @property + def attention_size(self) -> int: + return self.num_attention_heads * self.head_dim + + @staticmethod + def from_json(text: str) -> "ModelConfig": + raw = json.loads(text) + encoder = raw.get("encoder_config") or {} + hidden_size = int(raw.get("hidden_size") or raw.get("decoder_hidden_size") or 0) + num_heads = int(raw.get("num_attention_heads") or 1) + return ModelConfig( + model_type=str(raw.get("model_type", "")), + architectures=[str(value) for value in raw.get("architectures", ())], + vocab_size=int(raw.get("vocab_size", 0)), + hidden_size=hidden_size, + intermediate_size=int(raw.get("intermediate_size") or hidden_size * 4), + num_hidden_layers=int(raw.get("num_hidden_layers") or raw.get("num_decoder_layers") or 0), + num_attention_heads=num_heads, + num_key_value_heads=int(raw.get("num_key_value_heads") or num_heads), + rms_norm_eps=float(raw.get("rms_norm_eps", 1e-5)), + rope_theta=float(raw.get("rope_theta", 10000.0)), + bos_token_id=int(raw.get("bos_token_id", -1) if raw.get("bos_token_id") is not None else -1), + eos_token_id=int(raw.get("eos_token_id", -1) if raw.get("eos_token_id") is not None else -1), + pad_token_id=int(raw.get("pad_token_id", -1) if raw.get("pad_token_id") is not None else -1), + tie_word_embeddings=bool(raw.get("tie_word_embeddings", False)), + max_position_embeddings=int(raw.get("max_position_embeddings") or encoder.get("max_position_embeddings") or 8192), + hidden_act=str(raw.get("hidden_act", "")), + _head_dim=int(raw.get("head_dim", 0)), + raw=raw, + ) + + @classmethod + def create_tiny(cls, model_type: str, **overrides) -> "ModelConfig": + defaults = { + "model_type": model_type, + "vocab_size": 32, + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "rms_norm_eps": 1e-6, + "rope_theta": 10000.0, + "max_position_embeddings": 128, + } + defaults.update(overrides) + return cls.from_json(json.dumps(defaults)) + + @staticmethod + def from_dir(model_dir: str | Path) -> "ModelConfig": + config_path = Path(model_dir) / "config.json" + if not config_path.is_file(): + raise FileNotFoundError( + f"Parakeet TDT model directory is missing config.json: {config_path}" + ) + return ModelConfig.from_json(config_path.read_text(encoding="utf-8")) diff --git a/families/parakeet_tdt/requirements.txt b/families/parakeet_tdt/requirements.txt new file mode 100644 index 0000000000..ebf147fd12 --- /dev/null +++ b/families/parakeet_tdt/requirements.txt @@ -0,0 +1,3 @@ +librosa>=0.11.0 +safetensors +transformers==5.17.0 diff --git a/families/parakeet_tdt/runtime/CMakeLists.txt b/families/parakeet_tdt/runtime/CMakeLists.txt new file mode 100644 index 0000000000..9c7ba6df9f --- /dev/null +++ b/families/parakeet_tdt/runtime/CMakeLists.txt @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +add_library(trtmc_model_parakeet_tdt SHARED + audio_helpers.cpp bpe_tokenizer.cpp pipeline.cpp plugin.cpp resampler.cpp +) +target_include_directories(trtmc_model_parakeet_tdt PRIVATE + ${PROJECT_SOURCE_DIR} ${PROJECT_SOURCE_DIR}/core/runtime/include +) +target_include_directories(trtmc_model_parakeet_tdt SYSTEM PRIVATE ${TRTMC_CUDA_INCLUDE_DIR}) +target_link_libraries(trtmc_model_parakeet_tdt PRIVATE + trtmc_core nlohmann_json::nlohmann_json ${TRTMC_CUDART_LIBRARY} +) +target_compile_options(trtmc_model_parakeet_tdt PRIVATE -Wall -Wextra -Wpedantic) +set_target_properties(trtmc_model_parakeet_tdt PROPERTIES + LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}" BUILD_RPATH "\$ORIGIN" INSTALL_RPATH "\$ORIGIN" +) +install(TARGETS trtmc_model_parakeet_tdt LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}) + +if(TRTMC_BUILD_TESTS) + add_executable(test_parakeet_tdt_sdk_cpp ../tests/sdk_consumer.cpp) + target_link_libraries(test_parakeet_tdt_sdk_cpp PRIVATE trtmc_c nlohmann_json::nlohmann_json) + target_compile_options(test_parakeet_tdt_sdk_cpp PRIVATE -Wall -Wextra -Wpedantic -Werror) + set_target_properties(test_parakeet_tdt_sdk_cpp PROPERTIES + RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}" + ) +endif() diff --git a/families/parakeet_tdt/runtime/audio_helpers.cpp b/families/parakeet_tdt/runtime/audio_helpers.cpp new file mode 100644 index 0000000000..a568d71212 --- /dev/null +++ b/families/parakeet_tdt/runtime/audio_helpers.cpp @@ -0,0 +1,606 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "audio_helpers.h" + +#include "families/parakeet_tdt/runtime/resampler.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace trtmc { +namespace tdt { + +namespace { + +std::vector make_hann_window(int32_t length) { + std::vector window(length); + const double pi2 = 2.0 * 3.14159265358979323846; + for (int32_t i = 0; i < length; ++i) { + window[i] = static_cast( + 0.5 * (1.0 - std::cos(pi2 * static_cast(i) / static_cast(length)))); + } + return window; +} + +std::vector make_symmetric_hann_window(int32_t length) { + std::vector window(length); + if (length <= 1) { + return window; + } + const double pi2 = 2.0 * 3.14159265358979323846; + for (int32_t i = 0; i < length; ++i) { + window[i] = static_cast( + 0.5 * (1.0 - std::cos(pi2 * static_cast(i) / static_cast(length - 1)))); + } + return window; +} + +std::vector make_centered_stft_window(int32_t n_fft, int32_t win_length) { + std::vector window(static_cast(n_fft), 0.0F); + const auto inner = make_symmetric_hann_window(win_length); + const int32_t offset = std::max(0, (n_fft - win_length) / 2); + for (int32_t i = 0; i < win_length && offset + i < n_fft; ++i) { + window[static_cast(offset + i)] = inner[static_cast(i)]; + } + return window; +} + +void rfft_power_direct(const float* x, int32_t n, int32_t n_out, float* power_out) { + const double pi2 = 2.0 * 3.14159265358979323846; + for (int32_t k = 0; k < n_out; ++k) { + double re = 0.0; + double im = 0.0; + const double w = pi2 * static_cast(k) / static_cast(n); + for (int32_t t = 0; t < n; ++t) { + const double angle = w * static_cast(t); + re += static_cast(x[t]) * std::cos(angle); + im -= static_cast(x[t]) * std::sin(angle); + } + power_out[k] = static_cast(re * re + im * im); + } +} + +bool is_power_of_two(int32_t value) { + return value > 0 && (value & (value - 1)) == 0; +} + +void rfft_power_radix2(const float* x, int32_t n, int32_t n_out, float* power_out, + std::vector>& workspace, + const std::vector>& twiddles) { + workspace.resize(static_cast(n)); + for (int32_t i = 0; i < n; ++i) { + workspace[static_cast(i)] = + std::complex(static_cast(x[i]), 0.0); + } + + for (int32_t i = 1, j = 0; i < n; ++i) { + int32_t bit = n >> 1; + for (; (j & bit) != 0; bit >>= 1) { + j ^= bit; + } + j ^= bit; + if (i < j) { + std::swap(workspace[static_cast(i)], + workspace[static_cast(j)]); + } + } + + for (int32_t length = 2; length <= n; length <<= 1) { + const int32_t half = length / 2; + for (int32_t offset = 0; offset < n; offset += length) { + for (int32_t i = 0; i < half; ++i) { + const auto even = workspace[static_cast(offset + i)]; + const auto twiddle = twiddles[static_cast(i * n / length)]; + const auto odd = workspace[static_cast(offset + i + half)] * twiddle; + workspace[static_cast(offset + i)] = even + odd; + workspace[static_cast(offset + i + half)] = even - odd; + } + } + } + + const int32_t bins = std::min(n_out, n / 2 + 1); + for (int32_t k = 0; k < bins; ++k) { + power_out[k] = static_cast(std::norm(workspace[static_cast(k)])); + } + std::fill(power_out + bins, power_out + n_out, 0.0F); +} + +class RfftPowerPlan { + public: + explicit RfftPowerPlan(int32_t n) : n_(n) { + if (is_power_of_two(n_)) { + const double pi2 = 2.0 * 3.14159265358979323846; + twiddles_.reserve(static_cast(n_ / 2)); + for (int32_t k = 0; k < n_ / 2; ++k) { + const double angle = -pi2 * static_cast(k) / static_cast(n_); + twiddles_.emplace_back(std::cos(angle), std::sin(angle)); + } + } + } + + void execute(const float* input, int32_t n_out, float* power_out) { + if (is_power_of_two(n_)) { + rfft_power_radix2(input, n_, n_out, power_out, workspace_, twiddles_); + return; + } + rfft_power_direct(input, n_, n_out, power_out); + } + + private: + int32_t n_{0}; + std::vector> workspace_; + std::vector> twiddles_; +}; + +std::vector build_center_padded_audio(const float* samples, int32_t n_samples, + int32_t chunk_length_s, int32_t sample_rate, + int32_t n_fft, float preemphasis) { + const int32_t audio_length = chunk_length_s * sample_rate; + std::vector audio_padded(audio_length, 0.0F); + const int32_t copy_len = std::min(n_samples, audio_length); + if (copy_len > 0) { + if (preemphasis != 0.0F) { + audio_padded[0] = samples[0]; + for (int32_t i = 1; i < copy_len; ++i) { + audio_padded[static_cast(i)] = + samples[i] - preemphasis * samples[i - 1]; + } + } else { + std::memcpy(audio_padded.data(), samples, copy_len * sizeof(float)); + } + } + + const int32_t pad_size = n_fft / 2; + const int32_t padded_length = pad_size + audio_length + pad_size; + std::vector padded(padded_length, 0.0F); + std::memcpy(padded.data() + pad_size, audio_padded.data(), audio_length * sizeof(float)); + return padded; +} + +std::vector make_stft_window(const MelSpectrogramOptions& options) { + const int32_t win_length = options.win_length > 0 ? options.win_length : options.n_fft; + if (options.center_window_in_fft) { + return make_centered_stft_window(options.n_fft, win_length); + } + if (options.symmetric_window) { + return make_symmetric_hann_window(options.n_fft); + } + return make_hann_window(options.n_fft); +} + +std::vector compute_mel_spectrogram(const std::vector& padded, + const std::vector& window, + const float* mel_filters, int32_t n_fft, + int32_t hop_length, int32_t n_freq_bins, + int32_t n_mel_bins, int32_t frames_to_compute, + int32_t& n_frames_raw) { + n_frames_raw = 1 + (static_cast(padded.size()) - n_fft) / hop_length; + std::vector mel_spec(static_cast(n_mel_bins) * n_frames_raw, 0.0F); + std::vector windowed(n_fft); + std::vector frame_power(n_freq_bins); + std::vector frame_mel(n_mel_bins); + RfftPowerPlan fft_plan(n_fft); + + const int32_t computed_frames = std::clamp(frames_to_compute, 0, n_frames_raw); + for (int32_t t = 0; t < computed_frames; ++t) { + const int32_t start = t * hop_length; + for (int32_t i = 0; i < n_fft; ++i) { + windowed[static_cast(i)] = + padded[static_cast(start + i)] * window[static_cast(i)]; + } + + fft_plan.execute(windowed.data(), n_freq_bins, frame_power.data()); + + std::fill(frame_mel.begin(), frame_mel.end(), 0.0F); + for (int32_t f = 0; f < n_freq_bins; ++f) { + const float p = frame_power[static_cast(f)]; + if (p == 0.0F) { + continue; + } + const float* filter_row = mel_filters + static_cast(f) * n_mel_bins; + for (int32_t m = 0; m < n_mel_bins; ++m) { + frame_mel[static_cast(m)] += + p * filter_row[static_cast(m)]; + } + } + for (int32_t m = 0; m < n_mel_bins; ++m) { + mel_spec[static_cast(m) * n_frames_raw + t] = + frame_mel[static_cast(m)]; + } + } + return mel_spec; +} + +void normalize_log_mel_inplace(std::vector& mel_spec) { + float global_max = -1e10F; + for (float& value : mel_spec) { + value = std::log10(std::max(value, 1e-10F)); + if (value > global_max) { + global_max = value; + } + } + + const float floor = global_max - 8.0F; + for (float& value : mel_spec) { + value = std::max(value, floor); + value = (value + 4.0F) / 4.0F; + } +} + +void natural_log_mel_inplace(std::vector& mel_spec) { + constexpr float kLogGuard = 5.960464477539063e-08F; + for (float& value : mel_spec) { + value = std::log(value + kLogGuard); + } +} + +void normalize_per_feature_inplace(std::vector& mel_spec, int32_t n_mel_bins, + int32_t n_frames_raw, int32_t valid_frames) { + if (n_mel_bins <= 0 || n_frames_raw <= 0) { + return; + } + + valid_frames = std::clamp(valid_frames, 0, n_frames_raw); + if (valid_frames <= 0) { + std::fill(mel_spec.begin(), mel_spec.end(), 0.0F); + return; + } + + constexpr float kStdGuard = 1e-5F; + for (int32_t m = 0; m < n_mel_bins; ++m) { + const std::size_t base = static_cast(m) * n_frames_raw; + + double mean = 0.0; + for (int32_t t = 0; t < valid_frames; ++t) + mean += static_cast(mel_spec[base + static_cast(t)]); + mean /= static_cast(valid_frames); + + if (valid_frames == 1) { + mel_spec[base] = 0.0F; + } else { + double var = 0.0; + for (int32_t t = 0; t < valid_frames; ++t) { + const double diff = + static_cast(mel_spec[base + static_cast(t)]) - mean; + var += diff * diff; + } + const double stddev = + std::sqrt(var / static_cast(valid_frames - 1)) + kStdGuard; + for (int32_t t = 0; t < valid_frames; ++t) { + const std::size_t idx = base + static_cast(t); + mel_spec[idx] = + static_cast((static_cast(mel_spec[idx]) - mean) / stddev); + } + } + + for (int32_t t = valid_frames; t < n_frames_raw; ++t) + mel_spec[base + static_cast(t)] = 0.0F; + } +} + +void apply_log_scale_inplace(std::vector& mel_spec, MelLogScale log_scale) { + switch (log_scale) { + case MelLogScale::kNaturalLog: + natural_log_mel_inplace(mel_spec); + break; + case MelLogScale::kLog10Normalized: + normalize_log_mel_inplace(mel_spec); + break; + } +} + +std::vector trim_last_frame(std::vector mel_spec, int32_t n_mel_bins, + int32_t n_frames_raw, int32_t& n_frames_out) { + n_frames_out = n_frames_raw; + if (n_frames_raw <= 1) { + return mel_spec; + } + + n_frames_out = n_frames_raw - 1; + std::vector trimmed(static_cast(n_mel_bins) * n_frames_out); + for (int32_t m = 0; m < n_mel_bins; ++m) { + std::memcpy(trimmed.data() + static_cast(m) * n_frames_out, + mel_spec.data() + static_cast(m) * n_frames_raw, + n_frames_out * sizeof(float)); + } + return trimmed; +} + +} // namespace + +namespace detail { + +std::vector rfft_power(const std::vector& input) { + const int32_t n = static_cast(input.size()); + const int32_t n_out = n / 2 + 1; + std::vector power(static_cast(n_out), 0.0F); + if (n > 0) { + RfftPowerPlan plan(n); + plan.execute(input.data(), n_out, power.data()); + } + return power; +} + +} // namespace detail + +class IncrementalMelSpectrogram::Impl { + public: + Impl(const float* mel_filters, int32_t n_freq_bins, int32_t n_mel_bins, + MelSpectrogramOptions options, int32_t input_sample_rate) + : options_(std::move(options)), input_sample_rate_(input_sample_rate), + n_freq_bins_(n_freq_bins), n_mel_bins_(n_mel_bins), window_(make_stft_window(options_)), + fft_plan_(options_.n_fft), windowed_(static_cast(options_.n_fft)), + frame_power_(static_cast(n_freq_bins)), + frame_mel_(static_cast(n_mel_bins)) { + if (mel_filters == nullptr || n_freq_bins != options_.n_fft / 2 + 1 || n_mel_bins <= 0) { + throw std::invalid_argument("incremental TDT mel requires a complete filterbank"); + } + if (input_sample_rate_ <= 0 || options_.sample_rate <= 0 || options_.n_fft <= 0 || + options_.hop_length <= 0) { + throw std::invalid_argument("incremental TDT mel requires positive dimensions"); + } + if (options_.log_scale != MelLogScale::kNaturalLog || options_.normalize_per_feature) { + throw std::invalid_argument( + "incremental TDT mel only supports natural-log features without normalization"); + } + mel_filters_.assign(mel_filters, + mel_filters + static_cast(n_freq_bins_) * n_mel_bins_); + } + + void accept_audio(const float* samples, int32_t n_samples) { + if (n_samples < 0 || (n_samples > 0 && samples == nullptr)) { + throw std::invalid_argument("incremental TDT mel received invalid audio"); + } + if (n_samples == 0) { + return; + } + raw_audio_.insert(raw_audio_.end(), samples, samples + n_samples); + } + + int32_t target_sample_count() const { + const int64_t converted = + static_cast(raw_audio_.size()) * options_.sample_rate / input_sample_rate_; + const int64_t chunk_limit = + static_cast(options_.chunk_length_s) * options_.sample_rate; + return static_cast(std::min(converted, chunk_limit)); + } + + int32_t stable_target_sample_count(bool final) const { + if (final || input_sample_rate_ == options_.sample_rate) { + return target_sample_count(); + } + constexpr int32_t kResampleHalfTaps = 16; + const int64_t stable_source = + std::max(0, static_cast(raw_audio_.size()) - kResampleHalfTaps); + const int64_t stable_target = stable_source * options_.sample_rate / input_sample_rate_; + return static_cast(std::min(stable_target, target_sample_count())); + } + + int32_t available_frames() const { return target_sample_count() / options_.hop_length; } + + void ensure_resampled_samples(int32_t required_samples, bool final) { + if (!final && required_samples > stable_target_sample_count(false)) { + throw std::runtime_error( + "incremental TDT mel requested samples before resampler lookahead was ready"); + } + const int32_t bounded_required = std::min(required_samples, target_sample_count()); + const int32_t start = static_cast(resampled_audio_.size()); + if (bounded_required <= start) { + return; + } + auto next = resample_linear_range( + raw_audio_.data(), static_cast(raw_audio_.size()), input_sample_rate_, + options_.sample_rate, start, bounded_required - start); + resampled_audio_.insert(resampled_audio_.end(), next.begin(), next.end()); + } + + float preemphasized_sample(int32_t audio_index) const { + if (audio_index < 0 || audio_index >= static_cast(resampled_audio_.size())) { + return 0.0F; + } + const float current = resampled_audio_[static_cast(audio_index)]; + if (options_.preemphasis == 0.0F || audio_index == 0) { + return current; + } + return current - + options_.preemphasis * resampled_audio_[static_cast(audio_index - 1)]; + } + + void compute_frame(int32_t frame) { + const int32_t pad_size = options_.n_fft / 2; + const int32_t frame_start = frame * options_.hop_length - pad_size; + for (int32_t i = 0; i < options_.n_fft; ++i) { + windowed_[static_cast(i)] = + preemphasized_sample(frame_start + i) * window_[static_cast(i)]; + } + fft_plan_.execute(windowed_.data(), n_freq_bins_, frame_power_.data()); + + std::fill(frame_mel_.begin(), frame_mel_.end(), 0.0F); + for (int32_t f = 0; f < n_freq_bins_; ++f) { + const float power = frame_power_[static_cast(f)]; + const float* filter_row = + mel_filters_.data() + static_cast(f) * n_mel_bins_; + for (int32_t m = 0; m < n_mel_bins_; ++m) { + frame_mel_[static_cast(m)] += + power * filter_row[static_cast(m)]; + } + } + constexpr float kLogGuard = 5.960464477539063e-08F; + for (float value : frame_mel_) { + frames_.push_back(std::log(value + kLogGuard)); + } + } + + void ensure_frames(int32_t end_frame, bool final) { + if (end_frame < 0 || end_frame > available_frames()) { + throw std::out_of_range("incremental TDT mel frame request exceeds available audio"); + } + const int32_t current_frames = frame_count(); + if (end_frame <= current_frames) { + return; + } + const int32_t last_frame = end_frame - 1; + const int32_t required_samples = + last_frame * options_.hop_length + options_.n_fft - options_.n_fft / 2; + ensure_resampled_samples(required_samples, final); + for (int32_t frame = current_frames; frame < end_frame; ++frame) { + compute_frame(frame); + } + } + + void reset() { + raw_audio_.clear(); + resampled_audio_.clear(); + frames_.clear(); + } + + int32_t frame_count() const { + return static_cast(frames_.size() / static_cast(n_mel_bins_)); + } + + int32_t n_mels() const { return n_mel_bins_; } + + float value(int32_t mel_bin, int32_t frame) const { + if (mel_bin < 0 || mel_bin >= n_mel_bins_ || frame < 0 || frame >= frame_count()) { + throw std::out_of_range("incremental TDT mel feature index is out of range"); + } + return frames_[static_cast(frame) * n_mel_bins_ + mel_bin]; + } + + IncrementalMelStats stats() const { + return {static_cast(raw_audio_.size()), + static_cast(resampled_audio_.size()), frame_count()}; + } + + MelSpectrogramOptions options_; + int32_t input_sample_rate_{0}; + int32_t n_freq_bins_{0}; + int32_t n_mel_bins_{0}; + std::vector mel_filters_; + std::vector window_; + RfftPowerPlan fft_plan_; + std::vector raw_audio_; + std::vector resampled_audio_; + std::vector frames_; + std::vector windowed_; + std::vector frame_power_; + std::vector frame_mel_; +}; + +IncrementalMelSpectrogram::IncrementalMelSpectrogram(const float* mel_filters, int32_t n_freq_bins, + int32_t n_mel_bins, + MelSpectrogramOptions options, + int32_t input_sample_rate) + : impl_(std::make_unique(mel_filters, n_freq_bins, n_mel_bins, std::move(options), + input_sample_rate)) {} + +IncrementalMelSpectrogram::~IncrementalMelSpectrogram() = default; +IncrementalMelSpectrogram::IncrementalMelSpectrogram(IncrementalMelSpectrogram&&) noexcept = + default; +IncrementalMelSpectrogram& +IncrementalMelSpectrogram::operator=(IncrementalMelSpectrogram&&) noexcept = default; + +void IncrementalMelSpectrogram::accept_audio(const float* samples, int32_t n_samples) { + impl_->accept_audio(samples, n_samples); +} + +void IncrementalMelSpectrogram::ensure_frames(int32_t end_frame, bool final) { + impl_->ensure_frames(end_frame, final); +} + +void IncrementalMelSpectrogram::reset() { + impl_->reset(); +} + +int32_t IncrementalMelSpectrogram::available_frames() const { + return impl_->available_frames(); +} + +int32_t IncrementalMelSpectrogram::stable_target_sample_count(bool final) const { + return impl_->stable_target_sample_count(final); +} + +int32_t IncrementalMelSpectrogram::frame_count() const { + return impl_->frame_count(); +} + +int32_t IncrementalMelSpectrogram::n_mels() const { + return impl_->n_mels(); +} + +float IncrementalMelSpectrogram::value(int32_t mel_bin, int32_t frame) const { + return impl_->value(mel_bin, frame); +} + +IncrementalMelStats IncrementalMelSpectrogram::stats() const { + return impl_->stats(); +} + +MelResult extract_configured_mel_spectrogram(const float* samples, int32_t n_samples, + const float* mel_filters, int32_t n_freq_bins, + int32_t n_mel_bins, + const MelSpectrogramOptions& options) { + const int32_t expected_freq_bins = options.n_fft / 2 + 1; + const int32_t freq_bins = n_freq_bins == expected_freq_bins ? n_freq_bins : expected_freq_bins; + const int32_t chunk_samples = options.chunk_length_s * options.sample_rate; + const int32_t valid_audio = std::min(std::max(n_samples, 0), chunk_samples); + const int32_t pad_size = options.n_fft / 2; + const int32_t frames_to_compute = valid_audio > 0 && options.hop_length > 0 + ? 1 + (pad_size + valid_audio - 1) / options.hop_length + : 0; + const std::vector padded = + build_center_padded_audio(samples, n_samples, options.chunk_length_s, options.sample_rate, + options.n_fft, options.preemphasis); + + int32_t n_frames_raw = 0; + std::vector mel_spec = compute_mel_spectrogram( + padded, make_stft_window(options), mel_filters, options.n_fft, options.hop_length, + freq_bins, n_mel_bins, frames_to_compute, n_frames_raw); + apply_log_scale_inplace(mel_spec, options.log_scale); + if (options.normalize_per_feature) { + const int32_t valid_frames = + (options.hop_length > 0) ? (n_samples / options.hop_length) : 0; + normalize_per_feature_inplace(mel_spec, n_mel_bins, n_frames_raw, valid_frames); + } + + int32_t n_frames_out = 0; + mel_spec = trim_last_frame(std::move(mel_spec), n_mel_bins, n_frames_raw, n_frames_out); + + MelResult result; + result.data = std::move(mel_spec); + result.n_mels = n_mel_bins; + result.n_frames = n_frames_out; + return result; +} + +MelResult extract_tdt_mel_spectrogram(const float* samples, int32_t n_samples, + const float* mel_filters, int32_t n_freq_bins, + int32_t n_mel_bins, int32_t n_fft, int32_t win_length, + int32_t hop_length, int32_t chunk_length_s, + int32_t sample_rate, float preemphasis) { + MelSpectrogramOptions options; + options.n_fft = n_fft; + options.win_length = win_length; + options.hop_length = hop_length; + options.chunk_length_s = chunk_length_s; + options.sample_rate = sample_rate; + options.symmetric_window = true; + options.center_window_in_fft = true; + options.preemphasis = preemphasis; + options.log_scale = MelLogScale::kNaturalLog; + options.normalize_per_feature = true; + return extract_configured_mel_spectrogram(samples, n_samples, mel_filters, n_freq_bins, + n_mel_bins, options); +} + +} // namespace tdt +} // namespace trtmc diff --git a/families/parakeet_tdt/runtime/audio_helpers.h b/families/parakeet_tdt/runtime/audio_helpers.h new file mode 100644 index 0000000000..40958b739a --- /dev/null +++ b/families/parakeet_tdt/runtime/audio_helpers.h @@ -0,0 +1,90 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace trtmc { +namespace tdt { + +struct MelResult { + std::vector data; // [n_mels, n_frames] row-major + int32_t n_mels{0}; + int32_t n_frames{0}; +}; + +enum class MelLogScale { + kLog10Normalized, + kNaturalLog, +}; + +struct MelSpectrogramOptions { + int32_t n_fft{400}; + int32_t win_length{0}; + int32_t hop_length{160}; + int32_t chunk_length_s{30}; + int32_t sample_rate{16000}; + bool symmetric_window{false}; + bool center_window_in_fft{false}; + float preemphasis{0.0F}; + MelLogScale log_scale{MelLogScale::kLog10Normalized}; + bool normalize_per_feature{false}; +}; + +namespace detail { + +std::vector rfft_power(const std::vector& input); + +} // namespace detail + +struct IncrementalMelStats { + int64_t accepted_source_samples{0}; + int64_t generated_resampled_samples{0}; + int64_t computed_mel_frames{0}; +}; + +class IncrementalMelSpectrogram { + public: + IncrementalMelSpectrogram(const float* mel_filters, int32_t n_freq_bins, int32_t n_mel_bins, + MelSpectrogramOptions options, int32_t input_sample_rate); + ~IncrementalMelSpectrogram(); + + IncrementalMelSpectrogram(IncrementalMelSpectrogram&&) noexcept; + IncrementalMelSpectrogram& operator=(IncrementalMelSpectrogram&&) noexcept; + IncrementalMelSpectrogram(const IncrementalMelSpectrogram&) = delete; + IncrementalMelSpectrogram& operator=(const IncrementalMelSpectrogram&) = delete; + + void accept_audio(const float* samples, int32_t n_samples); + void ensure_frames(int32_t end_frame, bool final); + void reset(); + + int32_t available_frames() const; + int32_t stable_target_sample_count(bool final = false) const; + int32_t frame_count() const; + int32_t n_mels() const; + float value(int32_t mel_bin, int32_t frame) const; + IncrementalMelStats stats() const; + + private: + class Impl; + std::unique_ptr impl_; +}; + +MelResult extract_configured_mel_spectrogram(const float* samples, int32_t n_samples, + const float* mel_filters, int32_t n_freq_bins, + int32_t n_mel_bins, + const MelSpectrogramOptions& options); + +MelResult extract_tdt_mel_spectrogram(const float* samples, int32_t n_samples, + const float* mel_filters, int32_t n_freq_bins, + int32_t n_mel_bins, int32_t n_fft, int32_t win_length, + int32_t hop_length, int32_t chunk_length_s, + int32_t sample_rate, float preemphasis); + +} // namespace tdt +} // namespace trtmc diff --git a/families/parakeet_tdt/runtime/audio_input.h b/families/parakeet_tdt/runtime/audio_input.h new file mode 100644 index 0000000000..35963a8f50 --- /dev/null +++ b/families/parakeet_tdt/runtime/audio_input.h @@ -0,0 +1,55 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "families/parakeet_tdt/runtime/resampler.h" +#include "trtmc/internal/audio.h" + +#include +#include +#include +#include +#include + +namespace trtmc::tdt { + +inline std::vector prepare_audio(const internal::SpeechTranscriptionRequest& request) { + const auto& audio = request.audio; + if (request.source_language.has_value()) + throw std::invalid_argument("Parakeet TDT does not support source_language control"); + if (audio.channels == 0 || audio.samples.empty() || audio.samples.data() == nullptr || + audio.samples.size() % audio.channels != 0) + throw std::invalid_argument( + "Parakeet TDT requires complete nonempty interleaved PCM frames"); + const auto rate = audio.sample_rate.value_or(16000); + if (rate == 0 || rate > static_cast(std::numeric_limits::max())) + throw std::invalid_argument("Parakeet TDT input sample rate is out of range"); + const auto frames = audio.samples.size() / audio.channels; + if (frames > static_cast(std::numeric_limits::max()) || + frames > static_cast(rate) * 30) + throw std::invalid_argument("Parakeet TDT supports at most 30 seconds per request"); + + std::vector mono(frames); + for (std::size_t frame = 0; frame < frames; ++frame) { + double sum = 0; + for (std::uint32_t channel = 0; channel < audio.channels; ++channel) { + const auto value = audio.samples[frame * audio.channels + channel]; + if (!std::isfinite(value)) + throw std::invalid_argument("Parakeet TDT PCM must contain only finite samples"); + sum += value; + } + mono[frame] = static_cast(sum / audio.channels); + } + if (rate == 16000) + return mono; + auto result = resample_linear(mono.data(), static_cast(frames), + static_cast(rate), 16000); + if (result.empty()) + throw std::invalid_argument("Parakeet TDT input is shorter than one model-rate sample"); + return result; +} + +} // namespace trtmc::tdt diff --git a/families/parakeet_tdt/runtime/bpe_tokenizer.cpp b/families/parakeet_tdt/runtime/bpe_tokenizer.cpp new file mode 100644 index 0000000000..d6973e416f --- /dev/null +++ b/families/parakeet_tdt/runtime/bpe_tokenizer.cpp @@ -0,0 +1,1560 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "families/parakeet_tdt/runtime/tokenizer.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace trtmc::parakeet_tdt { +namespace { + +// ─── UTF-8 helpers ─── + +// Decode one UTF-8 codepoint from s starting at pos, advance pos. +inline char32_t utf8_to_char32(const std::string& s, size_t& pos) { + unsigned char c = static_cast(s[pos]); + if (c < 0x80) { + ++pos; + return static_cast(c); + } + if ((c & 0xE0) == 0xC0 && pos + 1 < s.size()) { + char32_t cp = (static_cast(c & 0x1F) << 6) | + static_cast(static_cast(s[pos + 1]) & 0x3F); + pos += 2; + return cp; + } + if ((c & 0xF0) == 0xE0 && pos + 2 < s.size()) { + char32_t cp = (static_cast(c & 0x0F) << 12) | + (static_cast(static_cast(s[pos + 1]) & 0x3F) << 6) | + static_cast(static_cast(s[pos + 2]) & 0x3F); + pos += 3; + return cp; + } + if ((c & 0xF8) == 0xF0 && pos + 3 < s.size()) { + char32_t cp = (static_cast(c & 0x07) << 18) | + (static_cast(static_cast(s[pos + 1]) & 0x3F) << 12) | + (static_cast(static_cast(s[pos + 2]) & 0x3F) << 6) | + static_cast(static_cast(s[pos + 3]) & 0x3F); + pos += 4; + return cp; + } + ++pos; + return 0xFFFD; +} + +inline std::string utf32_to_utf8(char32_t cp) { + std::string r; + if (cp <= 0x7F) { + r.push_back(static_cast(cp)); + } else if (cp <= 0x7FF) { + r.push_back(static_cast(0xC0 | ((cp >> 6) & 0x1F))); + r.push_back(static_cast(0x80 | (cp & 0x3F))); + } else if (cp <= 0xFFFF) { + r.push_back(static_cast(0xE0 | ((cp >> 12) & 0x0F))); + r.push_back(static_cast(0x80 | ((cp >> 6) & 0x3F))); + r.push_back(static_cast(0x80 | (cp & 0x3F))); + } else if (cp <= 0x10FFFF) { + r.push_back(static_cast(0xF0 | ((cp >> 18) & 0x07))); + r.push_back(static_cast(0x80 | ((cp >> 12) & 0x3F))); + r.push_back(static_cast(0x80 | ((cp >> 6) & 0x3F))); + r.push_back(static_cast(0x80 | (cp & 0x3F))); + } + return r; +} + +// Read one UTF-8 codepoint from raw bytes, advance ptr. Returns 0xFFFD on error. +inline char32_t read_utf8(const char*& p, const char* end) { + if (p >= end) + return 0xFFFD; + unsigned char c = static_cast(*p); + if (c < 0x80) { + ++p; + return c; + } + if ((c & 0xE0) == 0xC0 && p + 1 < end) { + char32_t cp = + (static_cast(c & 0x1F) << 6) | (static_cast(p[1]) & 0x3F); + p += 2; + return cp; + } + if ((c & 0xF0) == 0xE0 && p + 2 < end) { + char32_t cp = (static_cast(c & 0x0F) << 12) | + (static_cast(static_cast(p[1]) & 0x3F) << 6) | + (static_cast(p[2]) & 0x3F); + p += 3; + return cp; + } + if ((c & 0xF8) == 0xF0 && p + 3 < end) { + char32_t cp = (static_cast(c & 0x07) << 18) | + (static_cast(static_cast(p[1]) & 0x3F) << 12) | + (static_cast(static_cast(p[2]) & 0x3F) << 6) | + (static_cast(p[3]) & 0x3F); + p += 4; + return cp; + } + ++p; + return 0xFFFD; +} + +// ─── GPT-2 byte encoder: byte value <-> Unicode codepoint ─── + +struct ByteEncoderTables { + // byte -> UTF-8 encoded string (precomputed for speed) + std::string byte_to_utf8[256]; + // Unicode codepoint -> byte value + std::unordered_map cp_to_byte; + + ByteEncoderTables() { + // GPT-2 byte encoder: printable bytes map to themselves, + // others map to 256+ to avoid control chars. + bool direct[256] = {}; + for (int b = 33; b <= 126; ++b) + direct[b] = true; // !"#$...~ + for (int b = 161; b <= 172; ++b) + direct[b] = true; // non-breaking space area + for (int b = 174; b <= 255; ++b) + direct[b] = true; // extended latin + + int n = 0; + for (int b = 0; b < 256; ++b) { + char32_t cp = direct[b] ? static_cast(b) : static_cast(256 + n++); + byte_to_utf8[b] = utf32_to_utf8(cp); + cp_to_byte[cp] = static_cast(b); + } + } +}; + +static const ByteEncoderTables& byte_tables() { + static ByteEncoderTables tables; + return tables; +} + +// ─── Hand-written BPE pre-tokenizer ─── +// +// Supports two regex variants: +// +// GPT-2 (used by GPT-2, Falcon, OPT): +// 'contractions| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+ +// +// Qwen3 (used by Qwen3, LLaMA-3, Mistral, Phi): +// (?i:contractions)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| +// ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+ +// +// Key differences: Qwen3 allows any non-CR/LF/letter/digit as optional prefix before +// letters, matches single digits, and has explicit newline handling. + +namespace pretok { + +enum class Variant { kGpt2, kQwen3, kBloom, kDeepSeek, kClip }; + +// ── Character classification ── + +struct UnicodeRange { + char32_t lo, hi; +}; + +constexpr UnicodeRange kLetterRanges[] = { + {'A', 'Z'}, {'a', 'z'}, {0xB5, 0xB5}, // µ (micro sign, treated as letter) + {0xC0, 0xD6}, {0xD8, 0xF6}, {0xF8, 0x1BA}, // Latin-1 + Extended A/B + {0x1BC, 0x1BF}, {0x1C4, 0x293}, {0x295, 0x2AF}, // Latin Extended B cont. + {0x370, 0x373}, {0x376, 0x377}, {0x37B, 0x37D}, {0x37F, 0x37F}, // Greek + {0x386, 0x386}, {0x388, 0x38A}, {0x38C, 0x38C}, {0x38E, 0x3A1}, + {0x3A3, 0x3F5}, {0x3F7, 0x481}, {0x48A, 0x52F}, // Greek + Cyrillic + {0x531, 0x556}, {0x559, 0x559}, // Armenian + {0x560, 0x588}, // Armenian lowercase + {0x600, 0x6FF}, // Arabic + {0x900, 0x97F}, // Devanagari + {0xE00, 0xE7F}, // Thai + {0x10A0, 0x10C5}, // Georgian + {0x13A0, 0x13F5}, // Cherokee + {0x1C90, 0x1CBA}, {0x1CBD, 0x1CBF}, // Georgian Extended + {0x1D00, 0x1D2B}, {0x1D6B, 0x1D77}, {0x1D79, 0x1D9A}, // Phonetic Extensions + {0x1E00, 0x1F15}, {0x1F18, 0x1F1D}, {0x1F20, 0x1F45}, // Latin Ext. Additional + Greek Ext. + {0x2C00, 0x2C5F}, // Glagolitic + {0x3040, 0x309F}, // Hiragana + {0x30A0, 0x30FF}, // Katakana + {0x3400, 0x4DBF}, // CJK Extension A + {0x4E00, 0x9FFF}, // CJK Unified Ideographs + {0xAC00, 0xD7AF}, // Hangul Syllables + {0xFB00, 0xFDFF}, // Alphabetic Presentation Forms + Arabic Forms + {0x10000, 0x1007F}, // Linear B Syllabary +}; + +inline bool is_letter(char32_t cp) { + for (const auto& r : kLetterRanges) { + if (cp >= r.lo && cp <= r.hi) + return true; + } + return false; +} + +inline bool is_digit(char32_t cp) { + return (cp >= '0' && cp <= '9'); +} + +inline bool is_whitespace(char32_t cp) { + return cp == ' ' || cp == '\t' || cp == '\n' || cp == '\r' || cp == 0x0B || cp == 0x0C // VT, FF + || cp == 0xA0 // non-breaking space + || cp == 0x2000 || cp == 0x200A // en space through hair space + || cp == 0x3000; // ideographic space +} + +// BLOOM punctuation set: .,!?... +constexpr char32_t kBloomPunct[] = { + '.', ',', '!', '?', + 0x2026, // ellipsis + 0x3002, // ideographic full stop + 0xFF0C, // fullwidth comma + 0x3001, // ideographic comma + 0x0964, // Devanagari danda + 0x06D4, // Arabic full stop + 0x060C, // Arabic comma +}; + +inline bool is_bloom_punct(char32_t cp) { + for (auto c : kBloomPunct) { + if (cp == c) + return true; + } + return false; +} + +inline bool is_newline(char32_t cp) { + return cp == '\n' || cp == '\r'; +} + +// Is this char a valid optional prefix before a word? +// GPT-2/BLOOM: only space (0x20) +// Qwen3: any char except CR, LF, letter, digit +// DeepSeek: any whitespace (\s?) +inline bool is_prefix(char32_t cp, Variant v) { + if (v == Variant::kClip) { + return false; + } + if (v == Variant::kQwen3) { + return !is_newline(cp) && !is_letter(cp) && !is_digit(cp); + } + if (v == Variant::kDeepSeek) { + return is_whitespace(cp); + } + return cp == ' '; +} + +// BLOOM word char: not whitespace and not BLOOM punctuation +inline bool is_bloom_word_char(char32_t cp) { + // The serialized regex is `[^(\s|[.,!?…。，、।۔،])]`. Parentheses and + // the pipe are literal members of that negated character class. + return !is_whitespace(cp) && !is_bloom_punct(cp) && cp != '(' && cp != ')' && cp != '|'; +} + +// ── Scanning helpers (advance pointer past matching chars) ── + +inline void scan_letters(const char*& p, const char* end) { + while (p < end) { + const char* peek = p; + if (!is_letter(read_utf8(peek, end))) + break; + p = peek; + } +} + +inline void scan_digits(const char*& p, const char* end, int max_group = 0) { + int count = 0; + while (p < end) { + if (max_group > 0 && count >= max_group) + break; + const char* peek = p; + if (!is_digit(read_utf8(peek, end))) + break; + p = peek; + ++count; + } +} + +inline void scan_others(const char*& p, const char* end) { + while (p < end) { + const char* peek = p; + char32_t nc = read_utf8(peek, end); + if (is_whitespace(nc) || is_letter(nc) || is_digit(nc)) + break; + p = peek; + } +} + +inline void scan_newlines(const char*& p, const char* end) { + while (p < end) { + const char* peek = p; + if (!is_newline(read_utf8(peek, end))) + break; + p = peek; + } +} + +inline void scan_all_whitespace(const char*& p, const char* end) { + while (p < end) { + const char* peek = p; + if (!is_whitespace(read_utf8(peek, end))) + break; + p = peek; + } +} + +// Qwen3 regex branch: \s*[\r\n]+ +// Consume any leading whitespace up to the first newline, then consume the +// contiguous newline run itself, but stop before whitespace that follows the +// last newline. +inline void scan_qwen3_newline_chunk(char32_t first_cp, const char*& p, const char* end) { + bool saw_newline = is_newline(first_cp); + while (p < end) { + const char* peek = p; + char32_t nc = read_utf8(peek, end); + if (!saw_newline) { + if (is_newline(nc)) { + saw_newline = true; + p = peek; + continue; + } + if (is_whitespace(nc)) { + p = peek; + continue; + } + break; + } + if (!is_newline(nc)) + break; + p = peek; + } +} + +inline void scan_bloom_words(const char*& p, const char* end) { + while (p < end) { + const char* peek = p; + if (!is_bloom_word_char(read_utf8(peek, end))) + break; + p = peek; + } +} + +// Emit whitespace run, leaving last ws char for next token's optional prefix +inline void emit_whitespace_leave_last(const char*& p, const char* end, const char* start, + std::vector& result) { + const char* last_ws_start = start; + while (p < end) { + const char* peek = p; + char32_t nc = read_utf8(peek, end); + if (!is_whitespace(nc)) + break; + last_ws_start = p; + p = peek; + } + if (p < end && last_ws_start > start) { + p = last_ws_start; + } + result.emplace_back(start, p); +} + +// ── Contraction matching ── + +inline bool is_two_char_contraction(char c1, char c2) { + return (c1 == 'r' && c2 == 'e') || (c1 == 'v' && c2 == 'e') || (c1 == 'l' && c2 == 'l'); +} + +// Returns length of contraction suffix after apostrophe ('s 't 'm 'd 're 've 'll) +inline int match_contraction_suffix(const char* p, const char* end) { + if (p >= end) + return 0; + char c = *p; + if (c == 's' || c == 't' || c == 'm' || c == 'd') + return 1; + if (p + 1 < end && is_two_char_contraction(c, p[1])) + return 2; + return 0; +} + +// ── GPT-2/Qwen3 pre-tokenize dispatch helpers ── + +// Handle apostrophe: either contraction or "other" chars run +inline bool try_contraction(char32_t cp, const char*& p, const char* end, const char* start, + std::vector& result) { + if (cp != '\'') + return false; + int suffix = match_contraction_suffix(p, end); + if (suffix > 0) { + p += suffix; + } else { + scan_others(p, end); + } + result.emplace_back(start, p); + return true; +} + +// Handle optional prefix + letter/digit/other run +inline bool try_prefix_run(char32_t cp, const char*& p, const char* end, const char* start, + Variant variant, std::vector& result) { + if (!is_prefix(cp, variant) || p >= end) + return false; + const char* after_prefix = p; + char32_t next_cp = read_utf8(p, end); + + if (is_letter(next_cp)) { + scan_letters(p, end); + result.emplace_back(start, p); + return true; + } + if (is_digit(next_cp)) { + // The regex prefix [^\r\n\p{L}\p{N}]? only applies before \p{L}+ (letters). + // Digits are matched by standalone \p{N} — no prefix. Back up so the + // digit is handled by try_simple_run instead. + if (variant == Variant::kQwen3) { + p = after_prefix; + return false; + } + scan_digits(p, end); + result.emplace_back(start, p); + return true; + } + if (!is_whitespace(next_cp)) { + scan_others(p, end); + if (variant == Variant::kQwen3) + scan_newlines(p, end); + result.emplace_back(start, p); + return true; + } + // Prefix followed by whitespace — back up, let whitespace handler deal with it + p = after_prefix; + return false; +} + +// Check if a whitespace run contains any newline character +inline bool has_newline_in_ws(char32_t first_cp, const char* p, const char* end) { + if (is_newline(first_cp)) + return true; + const char* scan = p; + while (scan < end) { + const char* peek = scan; + char32_t nc = read_utf8(peek, end); + if (!is_whitespace(nc)) + break; + if (is_newline(nc)) + return true; + scan = peek; + } + return false; +} + +// Handle whitespace (Qwen3 newline sequences + general whitespace) +inline bool try_whitespace_run(char32_t cp, const char*& p, const char* end, const char* start, + Variant variant, std::vector& result) { + if (!is_whitespace(cp)) + return false; + + // kClip is selected only for CLIP's Removed + inverted Split contract, + // which retains regex matches and removes the gaps. Whitespace therefore + // never reaches ByteLevel. + if (variant == Variant::kClip) { + scan_all_whitespace(p, end); + return true; + } + + // Qwen3: \s*[\r\n]+ — newline sequences take priority + if (variant == Variant::kQwen3 && has_newline_in_ws(cp, p, end)) { + scan_qwen3_newline_chunk(cp, p, end); + result.emplace_back(start, p); + return true; + } + + // General whitespace: leave last ws char for next token's prefix + emit_whitespace_leave_last(p, end, start, result); + return true; +} + +// Handle letter or digit run (no prefix) +inline bool try_simple_run(char32_t cp, const char*& p, const char* end, const char* start, + Variant variant, std::vector& result, int digit_group = 0) { + if (is_letter(cp)) { + scan_letters(p, end); + result.emplace_back(start, p); + return true; + } + if (is_digit(cp)) { + if (variant == Variant::kQwen3) { + if (digit_group > 1) + scan_digits(p, end, digit_group - 1); + } else if (variant != Variant::kClip) { + scan_digits(p, end); + } + result.emplace_back(start, p); + return true; + } + return false; +} + +// ── Main pre-tokenize functions ── + +std::vector pre_tokenize(const std::string& text, Variant variant, + int digit_group = 0) { + std::vector result; + if (text.empty()) + return result; + + const char* p = text.data(); + const char* end = p + text.size(); + + while (p < end) { + const char* start = p; + char32_t cp = read_utf8(p, end); + + if (try_contraction(cp, p, end, start, result)) + continue; + if (try_prefix_run(cp, p, end, start, variant, result)) + continue; + if (try_whitespace_run(cp, p, end, start, variant, result)) + continue; + if (try_simple_run(cp, p, end, start, variant, result, digit_group)) + continue; + + // Other chars (punctuation/symbols) + scan_others(p, end); + if (variant == pretok::Variant::kQwen3) { + // Qwen3's ` ?[^\s\p{L}\p{N}]+[\r\n]*` keeps trailing newlines + // attached to punctuation/symbol runs even without an optional prefix. + scan_newlines(p, end); + } + result.emplace_back(start, p); + } + + return result; +} + +// Return the end of a BLOOM Split-regex match beginning at `start`. +// The regex is " ?[^(\s|[.,!?...])]+": an optional ASCII space followed by +// one or more non-whitespace, non-punctuation characters. +const char* bloom_match_end(const char* start, const char* end) { + const char* p = start; + char32_t cp = read_utf8(p, end); + if (cp == ' ') { + if (p >= end) + return nullptr; + cp = read_utf8(p, end); + if (!is_bloom_word_char(cp)) + return nullptr; + } else if (!is_bloom_word_char(cp)) { + return nullptr; + } + scan_bloom_words(p, end); + return p; +} + +// BLOOM uses Split(..., behavior="Isolated", invert=false) followed by +// ByteLevel(use_regex=false). Regex matches are isolated, while every +// contiguous unmatched span remains a single pre-token. In particular, +// punctuation and adjacent newlines must stay together so BPE can merge +// strings such as ".\n\n". +std::vector bloom_pre_tokenize(const std::string& text) { + std::vector result; + if (text.empty()) + return result; + + const char* p = text.data(); + const char* end = p + text.size(); + + while (p < end) { + const char* start = p; + if (const char* match_end = bloom_match_end(start, end)) { + result.emplace_back(start, match_end); + p = match_end; + continue; + } + + // Preserve one contiguous unmatched span until the next regex match. + read_utf8(p, end); + while (p < end && bloom_match_end(p, end) == nullptr) + read_utf8(p, end); + result.emplace_back(start, p); + } + + return result; +} + +} // namespace pretok + +// ─── BpeTokenizer implementation ─── + +class BpeTokenizer final : public ITokenizer { + public: + static std::unique_ptr Create(const char* tokenizer_json_data, + std::size_t tokenizer_json_size, + bool add_special_tokens = false) { + auto tokenizer = std::unique_ptr(new BpeTokenizer()); + tokenizer->mAddSpecialTokens = add_special_tokens; + tokenizer->parse_tokenizer_json(tokenizer_json_data, tokenizer_json_size); + return tokenizer; + } + + void encode_segment(const std::string& text, std::vector& result) const { + const auto normalized = normalize_text(text); + if (mIsSentencePiece) { + encode_sentencepiece(normalized, result); + } else if (mIsMetaspace) { + encode_metaspace(normalized, result); + } else { + encode_bytelevel(normalized, result); + } + } + + std::vector encode(const std::string& text) const override { + std::vector result; + if (text.empty()) + return result; + + if (mAddSpecialTokens) { + for (int32_t bos_id : mPostBosIds) + result.push_back(bos_id); + } + + auto segments = split_added_tokens(text); + for (const auto& seg : segments) { + if (seg.added_id >= 0) + result.push_back(seg.added_id); + else + encode_segment(seg.text, result); + } + + if (mAddSpecialTokens) { + for (int32_t eos_id : mPostEosIds) + result.push_back(eos_id); + } + + return result; + } + + std::string decode(const std::vector& ids) const override { + std::string joined = join_vocab_tokens(ids); + switch (mDecoderType) { + case DecoderType::kByteLevel: + return byte_decode(joined); + case DecoderType::kMetaspace: + return decode_metaspace(joined); + case DecoderType::kSequence: + return decode_sequence(joined); + } + return byte_decode(joined); + } + + int32_t id_for_token(std::string_view token) const override { + auto it = mTokenToId.find(std::string(token)); + return it != mTokenToId.end() ? it->second : -1; + } + + std::string token_for_id(int32_t id) const override { + if (id >= 0 && static_cast(id) < mVocab.size()) { + return mVocab[id]; + } + return ""; + } + + private: + BpeTokenizer() = default; + + enum class DecoderType { kByteLevel, kMetaspace, kSequence }; + + std::string normalize_text(const std::string& text) const { + if (!mNormalizeLowercase && !mNormalizeWhitespace) { + return text; + } + + std::string normalized; + normalized.reserve(text.size()); + const char* cursor = text.data(); + const char* end = cursor + text.size(); + bool previous_was_whitespace = false; + while (cursor < end) { + const char* char_start = cursor; + const char32_t cp = read_utf8(cursor, end); + if (mNormalizeWhitespace && pretok::is_whitespace(cp)) { + if (!previous_was_whitespace) { + normalized.push_back(' '); + } + previous_was_whitespace = true; + continue; + } + previous_was_whitespace = false; + if (mNormalizeLowercase && cp >= 'A' && cp <= 'Z') { + normalized.push_back(static_cast(cp - 'A' + 'a')); + } else { + normalized.append(char_start, cursor); + } + } + return normalized; + } + + // ─── Added token segmentation ─── + + struct Segment { + std::string text; + int32_t added_id; /* -1 = normal */ + }; + + std::pair find_longest_added_token(const std::string& text, size_t pos) const { + int32_t best_id = -1; + size_t best_len = 0; + for (const auto& [content, id] : mAddedTokenPatterns) { + if (pos + content.size() <= text.size() && content.size() > best_len && + text.compare(pos, content.size(), content) == 0) { + best_id = id; + best_len = content.size(); + } + } + return {best_id, best_len}; + } + + std::vector split_added_tokens(const std::string& text) const { + std::vector segments; + if (mAddedTokenPatterns.empty()) { + segments.push_back({text, -1}); + return segments; + } + size_t pos = 0; + while (pos < text.size()) { + auto [best_id, best_len] = find_longest_added_token(text, pos); + if (best_id >= 0) { + segments.push_back({text.substr(pos, best_len), best_id}); + pos += best_len; + } else { + if (segments.empty() || segments.back().added_id >= 0) { + segments.push_back({"", -1}); + } + segments.back().text.push_back(text[pos]); + ++pos; + } + } + return segments; + } + + // ─── Encoding helpers ─── + + // Metaspace encode (DeepSeek style): raw UTF-8 char split, BPE merge. + // Spaces are dropped (not in vocab), Ġ (U+0120) used as separator. + void encode_metaspace(const std::string& text, std::vector& result) const { + std::vector chars; + const char* cp = text.data(); + const char* ce = cp + text.size(); + while (cp < ce) { + const char* cs = cp; + read_utf8(cp, ce); + std::string ch(cs, cp); + if (mTokenToId.count(ch)) { + chars.push_back(std::move(ch)); + } + } + auto tokens = apply_merges(std::move(chars)); + for (const auto& token : tokens) { + auto it = mTokenToId.find(token); + if (it != mTokenToId.end()) { + result.push_back(it->second); + } + } + } + + // SentencePiece-style encode: replace spaces with ▁, split to chars, BPE merge. + // Used for Metaspace pre-tokenizer and Sequence decoder models (LLaMA, Mistral, Phi-3). + // Normalize text for SentencePiece: replace spaces with ▁, handle prepend + std::string normalize_sentencepiece(const std::string& text) const { + static const std::string sp = "\xe2\x96\x81"; // U+2581 + std::string out; + if (mSentencePiecePrependAlways) { + out = sp; + } + for (char c : text) { + out += (c == ' ') ? sp : std::string(1, c); + } + // Metaspace prepend_scheme=first: prepend if not already starting with ▁ + if (!mSentencePiecePrependAlways && mSentencePiecePrependIfMissing && + (out.empty() || out.compare(0, sp.size(), sp) != 0)) { + out = sp + out; + } + return out; + } + + void encode_sentencepiece(const std::string& text, std::vector& result) const { + std::string normalized = normalize_sentencepiece(text); + + // Split into UTF-8 characters + std::vector chars; + const char* cp = normalized.data(); + const char* ce = cp + normalized.size(); + while (cp < ce) { + const char* cs = cp; + read_utf8(cp, ce); + chars.emplace_back(cs, cp); + } + + // byte_fallback: replace unknown chars with <0xXX> + if (mByteFallback) { + chars = apply_byte_fallback_encode(std::move(chars)); + } + + // BPE merge and lookup + auto tokens = apply_merges(std::move(chars)); + for (const auto& token : tokens) { + auto it = mTokenToId.find(token); + if (it != mTokenToId.end()) { + result.push_back(it->second); + } + } + } + + // For byte_fallback: replace chars not in vocab with <0xXX> byte tokens + std::vector apply_byte_fallback_encode(std::vector chars) const { + std::vector result; + for (auto& ch : chars) { + if (mTokenToId.count(ch)) { + result.push_back(std::move(ch)); + } else { + // Split into individual bytes as <0xXX> + for (unsigned char byte : ch) { + char buf[8]; + std::snprintf(buf, sizeof(buf), "<0x%02X>", byte); + result.push_back(std::string(buf)); + } + } + } + return result; + } + + void encode_bytelevel(const std::string& text, std::vector& result) const { + std::vector words; + if (!mUsePreTokenizer) { + words = fallback_pre_tokenize(text); + } else if (mPreTokenizerVariant == pretok::Variant::kBloom) { + words = pretok::bloom_pre_tokenize(text); + } else { + words = pretok::pre_tokenize(text, mPreTokenizerVariant, mPreTokenizerDigitGroup); + } + + for (const auto& word : words) { + auto chars = byte_encode(word); + if (!chars.empty() && !mEndOfWordSuffix.empty()) { + chars.back() += mEndOfWordSuffix; + } + auto tokens = apply_merges(std::move(chars)); + for (const auto& token : tokens) { + auto it = mTokenToId.find(token); + if (it != mTokenToId.end()) { + result.push_back(it->second); + } + } + } + } + + // ─── Decoding helpers ─── + + std::string join_vocab_tokens(const std::vector& ids) const { + std::string joined; + for (int32_t id : ids) { + if (mSpecialIds.count(id)) + continue; + if (id >= 0 && static_cast(id) < mVocab.size()) { + joined += mVocab[id]; + } + } + return joined; + } + + std::string decode_metaspace(const std::string& joined) const { + std::string result; + const std::string spiece_char = "\xe2\x96\x81"; // U+2581 + size_t pos = 0; + while (pos < joined.size()) { + if (pos + spiece_char.size() <= joined.size() && + joined.compare(pos, spiece_char.size(), spiece_char) == 0) { + result.push_back(' '); + pos += spiece_char.size(); + } else { + result.push_back(joined[pos]); + ++pos; + } + } + if (!result.empty() && result[0] == ' ') { + result.erase(0, 1); + } + return result; + } + + // ─── Sequence decoder (SentencePiece BPE models: LLaMA, Mistral, Phi-3) ─── + + std::string decode_sequence(const std::string& joined) const { + std::string text = joined; + // Step 1: Apply Replace operations (e.g. ▁ → space) + for (const auto& rep : mSeqDecoderReplaces) { + text = string_replace_all(text, rep.pattern, rep.content); + } + // Step 2: ByteFallback — convert <0xXX> tokens to raw bytes + if (mSeqDecoderByteFallback) { + text = apply_byte_fallback(text); + } + // Step 3: Fuse is implicit (tokens already joined) + // Step 4: Strip leading space + if (mSeqDecoderStripLeft && !text.empty() && text[0] == ' ') { + text.erase(0, 1); + } + return text; + } + + static std::string string_replace_all(const std::string& input, const std::string& from, + const std::string& to) { + if (from.empty()) + return input; + std::string result; + result.reserve(input.size()); + size_t pos = 0; + while (pos < input.size()) { + if (pos + from.size() <= input.size() && input.compare(pos, from.size(), from) == 0) { + result += to; + pos += from.size(); + } else { + result.push_back(input[pos]); + ++pos; + } + } + return result; + } + + static int hex_char_value(char c) { + if (c >= '0' && c <= '9') + return c - '0'; + if (c >= 'a' && c <= 'f') + return c - 'a' + 10; + if (c >= 'A' && c <= 'F') + return c - 'A' + 10; + return -1; + } + + // Try to parse <0xXX> at position pos. Returns parsed byte or -1. + static int try_parse_byte_token(const std::string& text, size_t pos) { + if (pos + 6 > text.size()) + return -1; + if (text[pos] != '<' || text[pos + 1] != '0' || text[pos + 2] != 'x' || + text[pos + 5] != '>') + return -1; + int h = hex_char_value(text[pos + 3]); + int l = hex_char_value(text[pos + 4]); + if (h < 0 || l < 0) + return -1; + return (h << 4) | l; + } + + static std::string apply_byte_fallback(const std::string& text) { + std::string result; + result.reserve(text.size()); + size_t pos = 0; + while (pos < text.size()) { + int byte_val = try_parse_byte_token(text, pos); + if (byte_val >= 0) { + result.push_back(static_cast(byte_val)); + pos += 6; + } else { + result.push_back(text[pos]); + ++pos; + } + } + return result; + } + + // ─── Byte-level encoding ─── + + static std::vector byte_encode(const std::string& text) { + const auto& tables = byte_tables(); + std::vector result; + result.reserve(text.size()); + for (unsigned char byte : text) { + result.push_back(tables.byte_to_utf8[byte]); + } + return result; + } + + static std::string byte_decode(const std::string& text) { + const auto& tables = byte_tables(); + std::string result; + result.reserve(text.size()); + size_t pos = 0; + while (pos < text.size()) { + size_t start = pos; + char32_t cp = utf8_to_char32(text, pos); + auto it = tables.cp_to_byte.find(cp); + if (it != tables.cp_to_byte.end()) { + result.push_back(static_cast(it->second)); + } else { + // Not a GPT-2 byte-encoded codepoint — pass through raw UTF-8 + result.append(text, start, pos - start); + } + } + return result; + } + + // Generic pre-tokenizer used when no GPT-2 pattern is declared. + + static std::vector fallback_pre_tokenize(const std::string& text) { + std::vector result; + if (text.empty()) + return result; + result.push_back(text); + return result; + } + + // ─── BPE merge algorithm ─── + + struct MergeCandidate { + int rank; + std::string first; + std::string second; + }; + + MergeCandidate find_best_merge(const std::vector& tokens) const { + MergeCandidate best{INT_MAX, "", ""}; + for (size_t i = 0; i + 1 < tokens.size(); ++i) { + auto it = + mMergeRank.find(std::make_pair(std::cref(tokens[i]), std::cref(tokens[i + 1]))); + if (it != mMergeRank.end() && it->second < best.rank) { + best = {it->second, tokens[i], tokens[i + 1]}; + } + } + return best; + } + + static std::vector merge_all_pairs(std::vector tokens, + const std::string& first, + const std::string& second) { + std::string merged = first + second; + std::vector result; + result.reserve(tokens.size()); + for (size_t i = 0; i < tokens.size(); ++i) { + if (i + 1 < tokens.size() && tokens[i] == first && tokens[i + 1] == second) { + result.push_back(merged); + ++i; // skip next + } else { + result.push_back(std::move(tokens[i])); + } + } + return result; + } + + // Optimized: merge ALL occurrences of the best pair per pass. + std::vector apply_merges(std::vector tokens) const { + if (tokens.size() <= 1) + return tokens; + + while (true) { + auto best = find_best_merge(tokens); + if (best.rank == INT_MAX) + break; + tokens = merge_all_pairs(std::move(tokens), best.first, best.second); + if (tokens.size() <= 1) + break; + } + + return tokens; + } + + // ─── JSON parsing helpers ─── + + static bool is_eos_content(const std::string& s) { + return s == "<|endoftext|>" || s == "" || s == "<|end_of_text|>"; + } + + void parse_vocab(const nlohmann::json& j) { + auto& vocab_obj = j["model"]["vocab"]; + size_t vocab_size = vocab_obj.size(); + mVocab.resize(vocab_size); + + for (auto& [token, id] : vocab_obj.items()) { + int32_t token_id = id.get(); + if (token_id >= 0 && token_id < static_cast(vocab_size)) { + mVocab[token_id] = token; + mTokenToId[token] = token_id; + } + } + } + + void parse_merges(const nlohmann::json& j) { + if (!j["model"].contains("merges")) + throw std::runtime_error("Invalid tokenizer.json: missing model.merges"); + + auto& merges_arr = j["model"]["merges"]; + mMergeRank.reserve(merges_arr.size()); + + for (size_t i = 0; i < merges_arr.size(); ++i) { + std::string first, second; + + if (merges_arr[i].is_array()) { + auto arr = merges_arr[i].get>(); + if (arr.size() != 2) + continue; + first = std::move(arr[0]); + second = std::move(arr[1]); + } else if (merges_arr[i].is_string()) { + std::string merge_str = merges_arr[i].get(); + auto space_pos = merge_str.find(' '); + if (space_pos == std::string::npos) + continue; + first = merge_str.substr(0, space_pos); + second = merge_str.substr(space_pos + 1); + } else { + continue; + } + + auto pair = std::make_pair(std::move(first), std::move(second)); + mMergeRank[pair] = static_cast(i); + } + } + + void parse_added_tokens(const nlohmann::json& j) { + if (!j.contains("added_tokens")) + return; + for (auto& tok : j["added_tokens"]) { + std::string content = tok["content"].get(); + int32_t id = tok["id"].get(); + + if (id >= 0 && static_cast(id) >= mVocab.size()) { + mVocab.resize(static_cast(id) + 1); + } + if (id >= 0) { + mVocab[id] = content; + mTokenToId[content] = id; + } + + if (tok.value("special", false)) { + mSpecialTokens[content] = id; + mSpecialIds.insert(id); + if (is_eos_content(content)) + mEosId = id; + } + // All added tokens (special and non-special) participate in pre-split + // matching during encode, matching HuggingFace's AddedToken behavior. + // The special flag only controls decode filtering (mSpecialIds) and + // post_processor BOS/EOS insertion. + mAddedTokenPatterns.push_back({content, id}); + } + // Sort by length descending for longest-match-first + std::sort(mAddedTokenPatterns.begin(), mAddedTokenPatterns.end(), + [](const auto& a, const auto& b) { return a.first.size() > b.first.size(); }); + } + + // Parse digit group size from regex pattern like \p{N}{1,3} → 3. + // Returns 0 if no digit grouping found (single digit or unlimited). + static int parse_digit_group(const std::string& regex) { + // Search for "\p{N}{" — the standalone grouped variant, not [^\p{N}] + auto pos = regex.find("\\p{N}{"); + if (pos == std::string::npos) + return 0; + auto open = pos + 6; // after "\\p{N}{" + auto close = regex.find('}', open); + if (close == std::string::npos) + return 0; + auto comma = regex.find(',', open); + if (comma == std::string::npos || comma >= close) + return 0; + return std::stoi(regex.substr(comma + 1, close - comma - 1)); + } + + static bool is_clip_split_regex(const std::string& regex) { + const bool has_boundary_token = regex.find("startoftext") != std::string::npos || + regex.find("endoftext") != std::string::npos; + return has_boundary_token && regex.find("\\p{L}") != std::string::npos && + regex.find("\\p{N}") != std::string::npos; + } + + static pretok::Variant classify_clip_split(const nlohmann::json& split) { + const auto behavior = split.value("behavior", ""); + const bool invert = split.value("invert", false); + if (behavior == "Removed" && invert) + return pretok::Variant::kClip; + if (behavior == "Isolated" && invert) + return pretok::Variant::kGpt2; + throw std::runtime_error("Unsupported CLIP Split contract: behavior=" + behavior + + ", invert=" + (invert ? "true" : "false")); + } + + static bool is_qwen3_split_regex(const std::string& regex) { + return regex.find("[^\r\n") != std::string::npos || + regex.find("[^\\r\\n") != std::string::npos; + } + + static bool is_bloom_split_regex(const std::string& regex) { + return regex.find("[^(\\s") != std::string::npos || + regex.find("[^(\\\\s") != std::string::npos; + } + + // Classify one Split pre-tokenizer from both its pattern and contract. + static pretok::Variant classify_split(const nlohmann::json& split, int& digit_group_out) { + const auto regex = split["pattern"]["Regex"].get(); + if (is_clip_split_regex(regex)) + return classify_clip_split(split); + if (is_qwen3_split_regex(regex)) { + digit_group_out = parse_digit_group(regex); + return pretok::Variant::kQwen3; + } + if (is_bloom_split_regex(regex)) { + return pretok::Variant::kBloom; + } + if (regex.find("\\s?[A-Za-z") != std::string::npos) { + return pretok::Variant::kDeepSeek; + } + return pretok::Variant::kGpt2; + } + + // Detect variant from the Split inside a Sequence pre_tokenizer. + static pretok::Variant detect_split_variant(const nlohmann::json& pt, int& digit_group_out) { + digit_group_out = 0; + if (!pt.contains("pretokenizers")) + return pretok::Variant::kGpt2; + for (auto& sub : pt["pretokenizers"]) { + if (sub.value("type", "") != "Split") + continue; + if (!sub.contains("pattern") || !sub["pattern"].contains("Regex")) + continue; + return classify_split(sub, digit_group_out); + } + return pretok::Variant::kGpt2; + } + + static bool is_space_split_pre_tokenizer(const nlohmann::json& pt, const std::string& pt_type) { + if (pt_type != "Split") + return false; + if (!pt.contains("pattern") || !pt["pattern"].contains("String")) + return false; + return pt["pattern"]["String"].get() == " "; + } + + void detect_pre_tokenizer(const nlohmann::json& j) { + mUsePreTokenizer = true; + mPreTokenizerVariant = pretok::Variant::kGpt2; + + if (!j.contains("pre_tokenizer") || j["pre_tokenizer"].is_null()) + return; + auto& pt = j["pre_tokenizer"]; + std::string pt_type = pt.value("type", ""); + + if (pt_type == "ByteLevel") { + mPreTokenizerVariant = pretok::Variant::kGpt2; + } else if (pt_type == "Sequence") { + int digit_group = 0; + mPreTokenizerVariant = detect_split_variant(pt, digit_group); + mPreTokenizerDigitGroup = digit_group; + } else if (is_space_split_pre_tokenizer(pt, pt_type)) { + // Gemma SentencePiece-BPE tokenizers use a direct Split(" ") + // pre-tokenizer with spaces already normalized to U+2581. + mUsePreTokenizer = false; + } else if (pt_type == "Metaspace") { + mIsMetaspace = true; + mUsePreTokenizer = false; + } else if (pt_type.empty()) { + mUsePreTokenizer = false; + } else { + throw std::runtime_error("Unsupported pre_tokenizer type: " + pt_type); + } + } + + DecoderType classify_decoder_type(const nlohmann::json& j) const { + if (!j.contains("decoder") || j["decoder"].is_null()) { + return mIsMetaspace ? DecoderType::kMetaspace : DecoderType::kByteLevel; + } + std::string dt = j["decoder"].value("type", ""); + if (dt == "ByteLevel") + return DecoderType::kByteLevel; + if (dt == "Metaspace") + return DecoderType::kMetaspace; + if (dt == "Sequence") + return DecoderType::kSequence; + return mIsMetaspace ? DecoderType::kMetaspace : DecoderType::kByteLevel; + } + + void detect_decoder(const nlohmann::json& j) { + mDecoderType = classify_decoder_type(j); + if (mDecoderType == DecoderType::kSequence) { + parse_sequence_decoder(j["decoder"]); + } + // SentencePiece encode: vocab uses ▁ (U+2581) for spaces. + // Detect by: (1) ▁ in vocabulary, or (2) normalizer prepends ▁. + static const std::string spiece_marker = "\xe2\x96\x81"; // U+2581 + mIsSentencePiece = mTokenToId.count(spiece_marker) > 0 || mSentencePiecePrependAlways; + } + + void parse_seq_decoder_replace(const nlohmann::json& sub) { + SeqDecoderReplace rep; + if (sub.contains("pattern") && sub["pattern"].contains("String")) { + rep.pattern = sub["pattern"]["String"].get(); + } + rep.content = sub.value("content", ""); + if (!rep.pattern.empty()) { + mSeqDecoderReplaces.push_back(std::move(rep)); + } + } + + void parse_sequence_decoder(const nlohmann::json& dec) { + if (!dec.contains("decoders")) + return; + for (auto& sub : dec["decoders"]) { + std::string sub_type = sub.value("type", ""); + if (sub_type == "Replace") { + parse_seq_decoder_replace(sub); + } else if (sub_type == "ByteFallback") { + mSeqDecoderByteFallback = true; + } else if (sub_type == "Strip") { + if (sub.value("content", " ") == " " && sub.value("start", 0) > 0) { + mSeqDecoderStripLeft = true; + } + } + // Fuse is implicit (tokens already joined) + } + } + + static std::string optional_model_string(const nlohmann::json& model, const char* key) { + auto it = model.find(key); + if (it != model.end() && it->is_string()) + return it->get(); + return {}; + } + + void parse_tokenizer_json(const char* json_data, std::size_t json_size) { + nlohmann::json j; + try { + j = nlohmann::json::parse(json_data, json_data + json_size); + } catch (const std::exception& e) { + throw std::runtime_error(std::string("Failed to parse tokenizer.json: ") + e.what()); + } + + if (!j.contains("model")) + throw std::runtime_error("Invalid tokenizer.json: missing model"); + auto& model = j["model"]; + if (!model.contains("vocab") || !model["vocab"].is_object()) + throw std::runtime_error("Invalid tokenizer.json: model.vocab must be an object"); + if (!model.contains("merges") || !model["merges"].is_array()) + throw std::runtime_error("Invalid tokenizer.json: model.merges must be an array"); + + mByteFallback = model.value("byte_fallback", false); + mEndOfWordSuffix = optional_model_string(model, "end_of_word_suffix"); + parse_vocab(j); + parse_merges(j); + parse_added_tokens(j); + parse_post_processor(j); + detect_normalizer(j); + detect_pre_tokenizer(j); + detect_decoder(j); + } + + // Detect normalizer: check for Prepend (always prepend ▁) vs none + static bool normalizer_sequence_prepends(const nlohmann::json& norm, + const std::string& norm_type) { + if (norm_type != "Sequence" || !norm.contains("normalizers")) + return false; + for (auto& sub : norm["normalizers"]) { + if (sub.value("type", "") == "Prepend") + return true; + } + return false; + } + + static bool normalizer_replaces_space_with_sentence_piece(const nlohmann::json& norm, + const std::string& norm_type) { + if (norm_type == "Sequence" && norm.contains("normalizers")) { + for (const auto& sub : norm["normalizers"]) { + if (normalizer_replaces_space_with_sentence_piece(sub, sub.value("type", ""))) { + return true; + } + } + return false; + } + if (norm_type != "Replace") + return false; + if (!norm.contains("pattern") || !norm["pattern"].contains("String")) + return false; + if (norm["pattern"]["String"].get() != " ") + return false; + return norm.value("content", "") == "\xe2\x96\x81"; + } + + void detect_normalizer(const nlohmann::json& j) { + if (!j.contains("normalizer") || j["normalizer"].is_null()) + return; + auto& norm = j["normalizer"]; + std::string norm_type = norm.value("type", ""); + detect_text_normalizers(norm); + if (normalizer_sequence_prepends(norm, norm_type)) { + mSentencePiecePrependAlways = true; + } else if (normalizer_replaces_space_with_sentence_piece(norm, norm_type)) { + // Gemma replaces spaces with ▁ but does not prepend ▁ to the + // first token. Preserve the old prepend-if-missing behavior for + // Metaspace-style SentencePiece models. + mSentencePiecePrependIfMissing = false; + } + } + + void detect_text_normalizers(const nlohmann::json& normalizer) { + const std::string type = normalizer.value("type", ""); + if (type == "Sequence" && normalizer.contains("normalizers")) { + for (const auto& child : normalizer["normalizers"]) { + detect_text_normalizers(child); + } + return; + } + if (type == "Lowercase") { + mNormalizeLowercase = true; + return; + } + if (type == "Replace" && normalizer.contains("pattern") && + normalizer["pattern"].contains("Regex") && + normalizer["pattern"]["Regex"].get() == "\\s+" && + normalizer.value("content", "") == " ") { + mNormalizeWhitespace = true; + } + } + + // Parse post_processor to extract BOS/EOS tokens for add_special_tokens + void parse_post_processor(const nlohmann::json& j) { + if (!j.contains("post_processor") || j["post_processor"].is_null()) + return; + auto& pp = j["post_processor"]; + std::string pp_type = pp.value("type", ""); + + if (pp_type == "TemplateProcessing") { + parse_template_post_processor(pp); + } else if (pp_type == "Sequence") { + parse_sequence_post_processor(pp); + } else if (pp_type == "RobertaProcessing") { + parse_roberta_post_processor(pp); + } + // ByteLevel post_processor doesn't add tokens, nothing to do + } + + // Iterate Sequence post_processor's processors array to find TemplateProcessing + void parse_sequence_post_processor(const nlohmann::json& pp) { + if (!pp.contains("processors")) + return; + for (auto& sub : pp["processors"]) { + if (sub.value("type", "") == "TemplateProcessing") { + parse_template_post_processor(sub); + return; // first TemplateProcessing wins + } + } + } + + // RobertaProcessing: {"type":"RobertaProcessing","cls":["",0],"sep":["",2],...} + void parse_roberta_post_processor(const nlohmann::json& pp) { + if (pp.contains("cls") && pp["cls"].is_array() && pp["cls"].size() >= 2) { + int32_t cls_id = pp["cls"][1].get(); + mPostBosIds.push_back(cls_id); + } + if (pp.contains("sep") && pp["sep"].is_array() && pp["sep"].size() >= 2) { + int32_t sep_id = pp["sep"][1].get(); + mPostEosIds.push_back(sep_id); + } + } + + int32_t resolve_special_token_id(const nlohmann::json& entry) const { + if (!entry.contains("SpecialToken")) + return -1; + std::string token_id = entry["SpecialToken"].value("id", ""); + if (token_id.empty()) + return -1; + auto it = mSpecialTokens.find(token_id); + return it != mSpecialTokens.end() ? it->second : -1; + } + + void parse_template_post_processor(const nlohmann::json& pp) { + if (!pp.contains("single")) + return; + bool seen_sequence = false; + for (auto& entry : pp["single"]) { + if (entry.contains("Sequence")) { + seen_sequence = true; + continue; + } + int32_t id = resolve_special_token_id(entry); + if (id < 0) + continue; + if (seen_sequence) { + mPostEosIds.push_back(id); + } else { + mPostBosIds.push_back(id); + } + } + } + + // ─── Data members ─── + + std::vector mVocab; + std::unordered_map mTokenToId; + + struct PairHash { + size_t operator()(const std::pair& p) const { + size_t h1 = std::hash{}(p.first); + size_t h2 = std::hash{}(p.second); + return h1 ^ (h2 * 0x9e3779b97f4a7c15ULL + 0x9e3779b9 + (h1 << 6) + (h1 >> 2)); + } + }; + + std::unordered_map, int, PairHash> mMergeRank; + + std::unordered_map mSpecialTokens; + std::unordered_set mSpecialIds; // O(1) lookup in decode() + int32_t mEosId = -1; + + // Non-special added tokens: matched before pre-tokenization (longest first) + std::vector> mAddedTokenPatterns; + + bool mAddSpecialTokens = false; + bool mUsePreTokenizer = true; + bool mIsMetaspace = false; // Set by detect_pre_tokenizer. + bool mIsSentencePiece = false; + bool mSentencePiecePrependAlways = + false; // true for Normalizer Prepend, false for Metaspace first + bool mSentencePiecePrependIfMissing = true; + bool mByteFallback = false; + bool mNormalizeLowercase = false; + bool mNormalizeWhitespace = false; + std::string mEndOfWordSuffix; + + // Post-processor: BOS/EOS token IDs to add when add_special_tokens=true + // Vectors to support multiple BOS/EOS tokens (e.g. GLM-4: [gMASK] + ) + std::vector mPostBosIds; + std::vector mPostEosIds; + pretok::Variant mPreTokenizerVariant = pretok::Variant::kGpt2; + int mPreTokenizerDigitGroup = 0; // 0=unlimited, 3=\p{N}{1,3} (LLaMA 3.1/GLM-4) + + DecoderType mDecoderType = DecoderType::kByteLevel; + + // Sequence decoder config (parsed from tokenizer.json decoder field) + struct SeqDecoderReplace { + std::string pattern; + std::string content; + }; + std::vector mSeqDecoderReplaces; + bool mSeqDecoderByteFallback = false; + bool mSeqDecoderStripLeft = false; +}; + +} // namespace + +std::unique_ptr CreateBpeTokenizer(const char* tokenizer_json_data, + std::size_t tokenizer_json_size, + bool add_special_tokens) { + return BpeTokenizer::Create(tokenizer_json_data, tokenizer_json_size, add_special_tokens); +} + +} // namespace trtmc::parakeet_tdt diff --git a/families/parakeet_tdt/runtime/pipeline.cpp b/families/parakeet_tdt/runtime/pipeline.cpp new file mode 100644 index 0000000000..480590c786 --- /dev/null +++ b/families/parakeet_tdt/runtime/pipeline.cpp @@ -0,0 +1,299 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#include "families/parakeet_tdt/runtime/pipeline.h" + +#include "families/parakeet_tdt/runtime/audio_helpers.h" +#include "families/parakeet_tdt/runtime/audio_input.h" + +#include +#include +#include +#include +#include +#include + +namespace trtmc::parakeet_tdt { +namespace { +const internal::ConfigField fields[] = { + {"max_new_tokens", internal::ConfigKind::I64, internal::ConfigValue{int64_t{256}}, + "Maximum emitted transcript tokens"}, +}; +void validate_float_outputs(const TensorMap& outputs) { + for (const auto& item : outputs) { + const auto& tensor = item.second; + if (tensor.dtype != DType::kFloat32 || tensor.data == nullptr || tensor.shape.empty() || + std::any_of(tensor.shape.begin(), tensor.shape.end(), [](auto d) { return d <= 0; })) + throw std::runtime_error("Parakeet TDT requires nonempty FP32 host outputs: " + + item.first); + } +} +Tensor make_tensor(void* data, std::vector shape, DType dtype) { + Tensor t; + t.data = data; + t.shape = std::move(shape); + t.dtype = dtype; + return t; +} + +int32_t infer_encoder_frames(const Tensor& encoder_output, int32_t hidden_size) { + if (hidden_size <= 0) + return 0; + const auto elems = static_cast(encoder_output.numel()); + return elems > 0 ? elems / hidden_size : 0; +} + +int32_t argmax_token(const std::vector& logits) { + if (logits.empty()) + return -1; + return static_cast( + std::distance(logits.begin(), std::max_element(logits.begin(), logits.end()))); +} + +int32_t subsampled_frame_count(int32_t frames, bool causal) { + if (frames <= 0) + return 0; + for (int i = 0; i < 3; ++i) + frames = causal ? (frames / 2 + 1) : ((frames + 2 - 3) / 2 + 1); + return frames; +} +} // namespace +TdtPipeline::TdtPipeline(std::unique_ptr encoder, std::unique_ptr predictor, + std::unique_ptr joint, TdtConfig config, MelFilterbank mel, + std::shared_ptr tokenizer) + : encoder_(std::move(encoder)), predictor_(std::move(predictor)), joint_(std::move(joint)), + config_(std::move(config)), mel_fb_(std::make_unique(std::move(mel))), + tokenizer_(std::move(tokenizer)) { + if (!encoder_ || !predictor_ || !joint_ || !encoder_->ok() || !predictor_->ok() || + !joint_->ok() || !tokenizer_) + throw std::invalid_argument("Parakeet TDT requires three valid engines and a tokenizer"); + validate_tdt_mel_geometry(config_, mel_fb_->n_freq_bins, mel_fb_->n_mel_bins); + if (mel_fb_->data.size() != static_cast(mel_fb_->n_freq_bins) * mel_fb_->n_mel_bins || + config_.sample_rate != 16000 || config_.encoder_hidden_size <= 0 || + config_.pred_hidden_size <= 0 || config_.pred_num_layers <= 0 || config_.blank_id < 0 || + config_.max_symbols_per_step <= 0 || config_.duration_values.empty() || + std::any_of( + config_.duration_values.begin(), config_.duration_values.end(), + [](auto d) { return d < 0; })) + throw std::invalid_argument("invalid Parakeet TDT runtime dimensions"); +} +std::vector TdtPipeline::task_bindings() { + return {internal::bind(*this, fields)}; +} +TextResult TdtPipeline::run(const internal::SpeechTranscriptionRequest& request, + internal::ConfigView options) { + internal::validate_config(fields, options); + const auto limit = internal::config_get(options, fields, "max_new_tokens").value(); + if (limit <= 0 || limit > std::numeric_limits::max()) + throw std::invalid_argument("Parakeet TDT max_new_tokens must be a positive int32"); + const auto audio = tdt::prepare_audio(request); + std::lock_guard lock(mutex_); + const auto start = std::chrono::steady_clock::now(); + int32_t actual_frames = 0; + auto mel = + extract_padded_mel(audio.data(), static_cast(audio.size()), 16000, actual_frames); + if (mel.empty()) + throw std::runtime_error("Parakeet TDT mel extraction returned no features"); + auto encoded = run_encoder(mel, actual_frames); + const auto after_encoder = std::chrono::steady_clock::now(); + std::vector state_h( + static_cast(config_.pred_num_layers) * config_.pred_hidden_size, 0); + auto state_c = state_h; + auto predicted = run_predictor(config_.blank_id, state_h, state_c); + TextResult result; + decode_encoder_frames( + encoded, static_cast(encoded.size() / config_.encoder_hidden_size), + static_cast(limit), predicted, state_h, state_c, result.token_ids); + result.text = tokenizer_->decode(result.token_ids); + result.prefill_ms = std::chrono::duration(after_encoder - start).count(); + result.decode_ms = + std::chrono::duration(std::chrono::steady_clock::now() - after_encoder) + .count(); + return result; +} +std::vector TdtPipeline::extract_padded_mel(const float* audio_data, int32_t num_samples, + int32_t input_sample_rate, + int32_t& actual_frames) const { + const float* samples_ptr = audio_data; + int32_t samples_count = num_samples; + std::vector resampled; + if (input_sample_rate > 0 && input_sample_rate != config_.sample_rate) { + std::cerr << "[tdt] Resampling audio from " << input_sample_rate << " Hz to " + << config_.sample_rate << " Hz" << std::endl; + resampled = + tdt::resample_linear(audio_data, num_samples, input_sample_rate, config_.sample_rate); + samples_ptr = resampled.data(); + samples_count = static_cast(resampled.size()); + } + + tdt::MelResult mel = tdt::extract_tdt_mel_spectrogram( + samples_ptr, samples_count, mel_fb_->data.data(), mel_fb_->n_freq_bins, mel_fb_->n_mel_bins, + config_.mel_n_fft, config_.mel_win_length, config_.mel_hop_length, config_.mel_chunk_length, + config_.sample_rate, config_.mel_preemph); + actual_frames = std::min(mel.n_frames, std::max(0, samples_count / config_.mel_hop_length)); + if (mel.data.empty()) + return {}; + + const int32_t target_frames = config_.mel_length > 0 ? config_.mel_length : mel.n_frames; + std::vector padded(static_cast(mel.n_mels) * target_frames, 0.0F); + const int32_t copy_frames = std::min(mel.n_frames, target_frames); + for (int32_t m = 0; m < mel.n_mels; ++m) { + std::memcpy(padded.data() + static_cast(m) * target_frames, + mel.data.data() + static_cast(m) * mel.n_frames, + static_cast(copy_frames) * sizeof(float)); + } + return padded; +} + +std::vector TdtPipeline::run_encoder(const std::vector& mel, int32_t actual_frames) { + TensorMap inputs; + inputs["mel_features"] = + make_tensor(const_cast(mel.data()), {config_.num_mel_bins, config_.mel_length}, + DType::kFloat32); + + std::vector encoder_mask; + if (encoder_->has_input("encoder_mask")) { + const int32_t max_frames = std::max( + 1, config_.encoder_seq_len > 0 + ? config_.encoder_seq_len + : subsampled_frame_count(config_.mel_length, config_.causal_downsampling)); + int32_t actual_encoder_frames = + std::max(1, subsampled_frame_count(actual_frames, config_.causal_downsampling)); + actual_encoder_frames = std::min(actual_encoder_frames, max_frames); + encoder_mask.assign(static_cast(max_frames) * max_frames, -10000.0F); + for (int32_t q = 0; q < actual_encoder_frames; ++q) { + const int32_t k_begin = + config_.att_context_left < 0 ? 0 : std::max(0, q - config_.att_context_left); + const int32_t k_end = + config_.att_context_right < 0 + ? actual_encoder_frames - 1 + : std::min(actual_encoder_frames - 1, q + config_.att_context_right); + for (int32_t k = k_begin; k <= k_end; ++k) + encoder_mask[static_cast(q) * max_frames + k] = 0.0F; + } + inputs["encoder_mask"] = + make_tensor(encoder_mask.data(), {1, max_frames, max_frames}, DType::kFloat32); + } + + auto outputs = encoder_->forward(inputs); + validate_float_outputs(outputs); + auto it = outputs.find("encoder_output"); + if (it == outputs.end()) + throw std::runtime_error("TdtPipeline: encoder missing 'encoder_output'"); + if (it->second.numel() <= 0 || it->second.numel() % config_.encoder_hidden_size != 0) + throw std::runtime_error("TdtPipeline: encoder output size does not match config"); + const auto frames = infer_encoder_frames(it->second, config_.encoder_hidden_size); + const auto valid_frames = std::min( + frames, std::max(1, subsampled_frame_count(actual_frames, config_.causal_downsampling))); + const auto count = static_cast(valid_frames) * config_.encoder_hidden_size; + const auto* src = static_cast(it->second.data); + return std::vector(src, src + count); +} + +void TdtPipeline::decode_encoder_frames(const std::vector& encoder_output, + int32_t frame_count, int32_t token_limit, + std::vector& pred_output, + std::vector& state_h, std::vector& state_c, + std::vector& emitted) { + int32_t frame = 0; + while (frame < frame_count && static_cast(emitted.size()) < token_limit) { + const float* enc_ptr = + encoder_output.data() + static_cast(frame) * config_.encoder_hidden_size; + const float* enc_frame = enc_ptr; + int32_t symbols_on_frame = 0; + while (static_cast(emitted.size()) < token_limit) { + const auto logits = run_joint(enc_frame, pred_output.data()); + const auto token_count = static_cast(config_.blank_id + 1); + if (logits.size() != token_count + config_.duration_values.size()) + throw std::runtime_error("TdtPipeline: joint output sizes do not match config"); + const std::vector token_logits(logits.begin(), logits.begin() + token_count); + const std::vector duration_logits(logits.begin() + token_count, logits.end()); + const int32_t token = argmax_token(token_logits); + const int32_t duration_index = argmax_token(duration_logits); + if (token < 0) + break; + + const auto decision = make_tdt_greedy_decision( + token, duration_index, config_.duration_values, config_.blank_id); + if (decision.emit_token) { + emitted.push_back(token); + ++symbols_on_frame; + pred_output = run_predictor(token, state_h, state_c); + } + frame += decision.frame_advance; + if (decision.frame_advance > 0) + break; + if (symbols_on_frame >= config_.max_symbols_per_step) { + ++frame; + break; + } + } + } +} + +std::vector TdtPipeline::run_predictor(int32_t token_id, std::vector& state_h, + std::vector& state_c) { + TensorMap inputs; + inputs["token_id"] = make_tensor(&token_id, {1}, DType::kInt32); + const auto layer_stride = static_cast(config_.pred_hidden_size); + for (int32_t layer = 0; layer < config_.pred_num_layers; ++layer) { + const std::string suffix = "_" + std::to_string(layer); + inputs["state_h" + suffix] = + make_tensor(state_h.data() + static_cast(layer) * layer_stride, + {1, config_.pred_hidden_size}, DType::kFloat32); + inputs["state_c" + suffix] = + make_tensor(state_c.data() + static_cast(layer) * layer_stride, + {1, config_.pred_hidden_size}, DType::kFloat32); + } + + auto outputs = predictor_->forward(inputs); + validate_float_outputs(outputs); + auto pred_it = outputs.find("pred_output"); + if (pred_it == outputs.end()) + throw std::runtime_error("TdtPipeline: predictor missing 'pred_output'"); + + if (pred_it->second.numel() != static_cast(config_.pred_hidden_size)) + throw std::runtime_error("TdtPipeline: predictor output size does not match config"); + for (int32_t layer = 0; layer < config_.pred_num_layers; ++layer) { + const std::string suffix = "_" + std::to_string(layer); + auto h_it = outputs.find("next_h" + suffix); + auto c_it = outputs.find("next_c" + suffix); + if (h_it == outputs.end() || c_it == outputs.end()) + throw std::runtime_error("TdtPipeline: predictor missing next state outputs"); + if (h_it->second.numel() != static_cast(config_.pred_hidden_size) || + c_it->second.numel() != static_cast(config_.pred_hidden_size)) + throw std::runtime_error( + "TdtPipeline: predictor state output size does not match config"); + std::memcpy(state_h.data() + static_cast(layer) * layer_stride, + h_it->second.data, layer_stride * sizeof(float)); + std::memcpy(state_c.data() + static_cast(layer) * layer_stride, + c_it->second.data, layer_stride * sizeof(float)); + } + + const auto* pred = static_cast(pred_it->second.data); + return std::vector(pred, pred + config_.pred_hidden_size); +} + +std::vector TdtPipeline::run_joint(const float* encoder_frame, const float* pred_output) { + TensorMap inputs; + inputs["encoder_frame"] = make_tensor(const_cast(encoder_frame), + {1, config_.encoder_hidden_size}, DType::kFloat32); + inputs["pred_output"] = make_tensor(const_cast(pred_output), + {1, config_.pred_hidden_size}, DType::kFloat32); + auto outputs = joint_->forward(inputs); + validate_float_outputs(outputs); + auto token_it = outputs.find("token_logits"); + auto duration_it = outputs.find("duration_logits"); + if (token_it == outputs.end() || duration_it == outputs.end()) + throw std::runtime_error("TdtPipeline: joint missing token_logits or duration_logits"); + if (token_it->second.numel() != static_cast(config_.blank_id) + 1 || + duration_it->second.numel() != config_.duration_values.size()) + throw std::runtime_error("TdtPipeline: joint output sizes do not match config"); + const auto* tokens = static_cast(token_it->second.data); + const auto* durations = static_cast(duration_it->second.data); + std::vector logits(tokens, tokens + token_it->second.numel()); + logits.insert(logits.end(), durations, durations + duration_it->second.numel()); + return logits; +} +} // namespace trtmc::parakeet_tdt diff --git a/families/parakeet_tdt/runtime/pipeline.h b/families/parakeet_tdt/runtime/pipeline.h new file mode 100644 index 0000000000..ff04595000 --- /dev/null +++ b/families/parakeet_tdt/runtime/pipeline.h @@ -0,0 +1,43 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once +#include "families/parakeet_tdt/runtime/tdt_config.h" +#include "families/parakeet_tdt/runtime/tokenizer.h" +#include "trtmc/internal/audio.h" +#include "trtmc/internal/model.h" +#include "trtmc/runtime/trt_module.h" + +#include + +namespace trtmc::parakeet_tdt { +struct MelFilterbank { + std::vector data; + int32_t n_freq_bins{0}; + int32_t n_mel_bins{0}; +}; + +class TdtPipeline final : public internal::IModel, public internal::ISpeechTranscription { + public: + TdtPipeline(std::unique_ptr encoder, std::unique_ptr predictor, + std::unique_ptr joint, TdtConfig config, MelFilterbank mel, + std::shared_ptr tokenizer); + const char* task() const noexcept override { return "speech_transcription"; } + std::vector task_bindings() override; + TextResult run(const internal::SpeechTranscriptionRequest&, internal::ConfigView) override; + + private: + std::vector extract_padded_mel(const float*, int32_t, int32_t, int32_t&) const; + std::vector run_encoder(const std::vector&, int32_t); + std::vector run_predictor(int32_t, std::vector&, std::vector&); + std::vector run_joint(const float*, const float*); + void decode_encoder_frames(const std::vector&, int32_t, int32_t, std::vector&, + std::vector&, std::vector&, std::vector&); + std::unique_ptr encoder_, predictor_, joint_; + TdtConfig config_; + std::unique_ptr mel_fb_; + std::shared_ptr tokenizer_; + std::mutex mutex_; +}; +} // namespace trtmc::parakeet_tdt diff --git a/families/parakeet_tdt/runtime/plugin.cpp b/families/parakeet_tdt/runtime/plugin.cpp new file mode 100644 index 0000000000..ae7d277f14 --- /dev/null +++ b/families/parakeet_tdt/runtime/plugin.cpp @@ -0,0 +1,103 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#include "families/parakeet_tdt/runtime/pipeline.h" +#include "trtmc/runtime/family_factory.h" +#include "trtmc/runtime/trt_backend.h" + +#include +#include +#include +#include + +namespace trtmc::parakeet_tdt { +namespace { +std::vector require_section(const BundleReader& reader, const char* name) { + const auto* section = reader.find_section(name); + if (!section || section->length == 0) + throw std::runtime_error(std::string("Parakeet TDT missing section: ") + name); + return reader.read_section(name); +} +std::unique_ptr load_engine(const FamilyContext& context, const char* name) { + auto plan = require_section(context.reader, name); + auto module = context.backend.create_module(plan.data(), plan.size(), {}); + if (!module || !module->ok()) + throw std::runtime_error(std::string("Parakeet TDT failed to load ") + name); + return module; +} +} // namespace +} // namespace trtmc::parakeet_tdt + +extern "C" trtmc::ITask* trtmc_create_family(const trtmc::FamilyContext& context) { + using namespace trtmc; + using namespace trtmc::parakeet_tdt; + if (context.kv_cache_size_bytes != 0) + throw std::invalid_argument("parakeet_tdt does not support --kv-cache-size"); + if (context.reader.info().family != "parakeet_tdt" || + context.reader.info().task != "speech_transcription" || + context.reader.info().backend != "trt" || std::string(context.backend.name()) != "trt") + throw std::invalid_argument("Parakeet TDT bundle identity does not match runtime"); + const auto data = require_section(context.reader, "runtime.json"); + const auto raw = nlohmann::json::parse(data.begin(), data.end()); + TdtConfig config; + const std::pair dimensions[] = { + {"tensor_parallel_size", 1}, + {"mel_sampling_rate", 16000}, + {"num_mel_bins", 128}, + {"mel_n_fft", 512}, + {"mel_win_length", 400}, + {"mel_hop_length", 160}, + {"mel_chunk_length", 30}, + {"mel_length", 3000}, + {"tdt_encoder_hidden_size", 1024}, + {"tdt_pred_hidden_size", 640}, + {"tdt_pred_num_layers", 2}, + {"tdt_vocab_size", 8192}, + {"tdt_blank_id", 8192}, + {"tdt_encoder_layers", 24}, + {"max_source_positions", 375}, + {"subsampling_factor", 8}, + {"tdt_max_symbols_per_step", 10}, + {"tdt_att_context_left", -1}, + {"tdt_att_context_right", -1}, + }; + for (const auto& [name, expected] : dimensions) { + if (!raw.at(name).is_number_integer() || raw.at(name) != expected) + throw std::invalid_argument(std::string("unsupported Parakeet TDT dimension: ") + name); + } + if (raw.at("tdt_duration_values") != nlohmann::json({0, 1, 2, 3, 4}) || + raw.at("tdt_causal_downsampling") != false || raw.at("mel_normalize") != "per_feature" || + std::abs(raw.at("mel_preemph").get() - 0.97) > 1e-6) + throw std::invalid_argument("unsupported Parakeet TDT frontend or duration policy"); + + const auto mel = require_section(context.reader, "mel_filterbank"); + constexpr size_t values = 257 * 128; + if (mel.size() != 8 + values * sizeof(float)) + throw std::invalid_argument("invalid Parakeet TDT mel filterbank size"); + MelFilterbank filters; + std::memcpy(&filters.n_freq_bins, mel.data(), 4); + std::memcpy(&filters.n_mel_bins, mel.data() + 4, 4); + validate_tdt_mel_geometry(config, filters.n_freq_bins, filters.n_mel_bins); + filters.data.resize(values); + std::memcpy(filters.data.data(), mel.data() + 8, values * sizeof(float)); + for (float value : filters.data) + if (!std::isfinite(value)) + throw std::invalid_argument("nonfinite Parakeet mel coefficient"); + const auto tokenizer_data = require_section(context.reader, "tokenizer.json"); + const auto tokenizer_json = nlohmann::json::parse(tokenizer_data.begin(), tokenizer_data.end()); + if (tokenizer_json.at("model").at("type") != "BPE" || + tokenizer_json.at("decoder").at("type") != "Metaspace" || + tokenizer_json.at("decoder").at("replacement") != "\xe2\x96\x81" || + tokenizer_json.at("decoder").at("prepend_scheme") != "always") + throw std::invalid_argument( + "Parakeet TDT requires its native BPE tokenizer with Metaspace decoding"); + auto tokenizer = CreateBpeTokenizer(tokenizer_data.data(), tokenizer_data.size(), false); + if (!tokenizer) + throw std::invalid_argument("Parakeet TDT requires its native BPE tokenizer"); + auto encoder = load_engine(context, "encoder.plan"); + auto predictor = load_engine(context, "predictor.plan"); + auto joint = load_engine(context, "joint.plan"); + return new TdtPipeline(std::move(encoder), std::move(predictor), std::move(joint), + std::move(config), std::move(filters), std::move(tokenizer)); +} diff --git a/families/parakeet_tdt/runtime/resampler.cpp b/families/parakeet_tdt/runtime/resampler.cpp new file mode 100644 index 0000000000..fda7c33e9b --- /dev/null +++ b/families/parakeet_tdt/runtime/resampler.cpp @@ -0,0 +1,146 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "families/parakeet_tdt/runtime/resampler.h" + +#include +#include +#include +#include + +namespace trtmc::tdt { +namespace { + +std::int32_t output_length(std::int32_t num_samples, std::int32_t source_rate, + std::int32_t target_rate) { + return static_cast(static_cast(num_samples) * target_rate / + source_rate); +} + +double scaled_sinc(double distance, double cutoff) { + constexpr double pi = 3.14159265358979323846; + if (std::abs(distance) < 1.0e-12) + return cutoff; + return cutoff * std::sin(pi * distance * cutoff) / (pi * distance * cutoff); +} + +double hann_window(double distance, std::int32_t half_taps) { + constexpr double pi = 3.14159265358979323846; + const double position = + (distance + static_cast(half_taps)) / (2.0 * static_cast(half_taps)); + return 0.5 * (1.0 - std::cos(2.0 * pi * position)); +} + +float resample_at(const float* samples, std::int32_t num_samples, double source_position, + double cutoff, std::int32_t half_taps) { + const auto center = static_cast(std::floor(source_position)); + const auto low = std::max(0, center - half_taps + 1); + const auto high = std::min(num_samples - 1, center + half_taps); + double value = 0.0; + double weight_sum = 0.0; + for (std::int32_t index = low; index <= high; ++index) { + const auto distance = static_cast(index) - source_position; + const auto weight = scaled_sinc(distance, cutoff) * hann_window(distance, half_taps); + value += static_cast(samples[index]) * weight; + weight_sum += weight; + } + return weight_sum > 1.0e-12 ? static_cast(value / weight_sum) : 0.0F; +} + +struct PhaseTable { + std::int32_t rate_gcd{1}; + std::int32_t phase_count{1}; + std::int32_t half_taps{16}; + std::vector weights; +}; + +PhaseTable make_phase_table(std::int32_t source_rate, std::int32_t target_rate, double cutoff, + std::int32_t half_taps) { + PhaseTable table; + table.rate_gcd = std::gcd(source_rate, target_rate); + table.phase_count = target_rate / table.rate_gcd; + table.half_taps = half_taps; + const auto tap_count = 2 * half_taps; + table.weights.resize(static_cast(table.phase_count) * tap_count); + for (std::int32_t phase = 0; phase < table.phase_count; ++phase) { + const auto fraction = + static_cast(phase * table.rate_gcd) / static_cast(target_rate); + for (std::int32_t tap = 0; tap < tap_count; ++tap) { + const auto offset = tap - half_taps + 1; + const auto distance = static_cast(offset) - fraction; + table.weights[static_cast(phase) * tap_count + tap] = + scaled_sinc(distance, cutoff) * hann_window(distance, half_taps); + } + } + return table; +} + +std::vector resample_phases(const float* samples, std::int32_t num_samples, + std::int32_t source_rate, std::int32_t target_rate, + std::int32_t output_start, std::int32_t output_count, + double cutoff, std::int32_t half_taps) { + const auto table = make_phase_table(source_rate, target_rate, cutoff, half_taps); + const auto tap_count = 2 * half_taps; + std::vector output(static_cast(output_count)); + for (std::int32_t local = 0; local < output_count; ++local) { + const auto index = output_start + local; + const auto numerator = static_cast(index) * source_rate; + const auto center = static_cast(numerator / target_rate); + const auto phase = static_cast(numerator % target_rate) / table.rate_gcd; + const auto* weights = table.weights.data() + static_cast(phase) * tap_count; + const auto first = std::max(-half_taps + 1, -center); + const auto last = std::min(half_taps, num_samples - 1 - center); + double value = 0.0; + double weight_sum = 0.0; + for (std::int32_t offset = first; offset <= last; ++offset) { + const auto weight = weights[offset + half_taps - 1]; + value += static_cast(samples[center + offset]) * weight; + weight_sum += weight; + } + output[static_cast(local)] = + weight_sum > 1.0e-12 ? static_cast(value / weight_sum) : 0.0F; + } + return output; +} + +} // namespace + +std::vector resample_linear(const float* samples, std::int32_t num_samples, + std::int32_t source_rate, std::int32_t target_rate) { + if (num_samples <= 0) + return {}; + return resample_linear_range(samples, num_samples, source_rate, target_rate, 0, + output_length(num_samples, source_rate, target_rate)); +} + +std::vector resample_linear_range(const float* samples, std::int32_t num_samples, + std::int32_t source_rate, std::int32_t target_rate, + std::int32_t output_start, std::int32_t output_count) { + if (samples == nullptr || num_samples <= 0 || source_rate <= 0 || target_rate <= 0) + throw std::invalid_argument("resampler requires samples and positive dimensions"); + const auto full_length = output_length(num_samples, source_rate, target_rate); + const auto start = std::clamp(output_start, 0, full_length); + const auto count = std::clamp(output_count, 0, full_length - start); + if (count == 0) + return {}; + if (source_rate == target_rate) + return {samples + start, samples + start + count}; + constexpr std::int32_t half_taps = 16; + const auto cutoff = + std::min(1.0, static_cast(target_rate) / static_cast(source_rate)); + const auto phase_count = target_rate / std::gcd(source_rate, target_rate); + if (phase_count <= 2048) + return resample_phases(samples, num_samples, source_rate, target_rate, start, count, cutoff, + half_taps); + std::vector output(static_cast(count)); + for (std::int32_t local = 0; local < count; ++local) { + const auto position = static_cast(start + local) * source_rate / target_rate; + output[static_cast(local)] = + resample_at(samples, num_samples, position, cutoff, half_taps); + } + return output; +} + +} // namespace trtmc::tdt diff --git a/families/parakeet_tdt/runtime/resampler.h b/families/parakeet_tdt/runtime/resampler.h new file mode 100644 index 0000000000..73fea1dbc0 --- /dev/null +++ b/families/parakeet_tdt/runtime/resampler.h @@ -0,0 +1,19 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +namespace trtmc::tdt { + +std::vector resample_linear(const float* samples, std::int32_t num_samples, + std::int32_t source_rate, std::int32_t target_rate); +std::vector resample_linear_range(const float* samples, std::int32_t num_samples, + std::int32_t source_rate, std::int32_t target_rate, + std::int32_t output_start, std::int32_t output_count); + +} // namespace trtmc::tdt diff --git a/families/parakeet_tdt/runtime/tdt_config.h b/families/parakeet_tdt/runtime/tdt_config.h new file mode 100644 index 0000000000..183920ac5d --- /dev/null +++ b/families/parakeet_tdt/runtime/tdt_config.h @@ -0,0 +1,125 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include +#include + +namespace trtmc { + +struct TdtConfig { + int32_t sample_rate{16000}; + int32_t num_mel_bins{128}; + int32_t mel_n_fft{512}; + int32_t mel_win_length{400}; + int32_t mel_hop_length{160}; + int32_t mel_chunk_length{30}; + float mel_preemph{0.97F}; + int32_t mel_length{3000}; + int32_t encoder_hidden_size{1024}; + int32_t pred_hidden_size{640}; + int32_t pred_num_layers{2}; + int32_t vocab_size{8192}; + int32_t blank_id{8192}; + int32_t max_symbols_per_step{10}; + int32_t encoder_seq_len{375}; + int32_t encoder_layers{24}; + int32_t subsampling_factor{8}; + std::vector duration_values{0, 1, 2, 3, 4}; + + // Retained as inert fields so the family-private offline pipeline shares + // one constructor shape with the repository's speech pipeline machinery. + int32_t att_context_left{-1}; + int32_t att_context_right{-1}; + int32_t streaming_cache_left{0}; + int32_t streaming_time_cache{0}; + int32_t streaming_pre_encode_cache{0}; + int32_t streaming_drop_pre_encoded{0}; + bool causal_downsampling{false}; + bool has_prompt_kernel{false}; + int32_t num_prompts{0}; + std::unordered_map prompt_dictionary; + std::vector supported_right_contexts; +}; + +struct TdtGreedyDecision { + bool emit_token{false}; + int32_t frame_advance{0}; +}; + +// Streaming is deliberately not advertised for v3, but these private fields +// keep the inherited speech-session implementation buildable and fail closed +// unless a bundle contains matching streaming encoder plans. +struct TdtStreamingSchedule { + int32_t att_context_left{70}; + int32_t att_context_right{13}; + int32_t subsampling_factor{8}; + int32_t encoder_frame_ms{80}; + int32_t chunk_ms{1120}; + int32_t chunk_samples{17920}; + int32_t first_chunk_mel_frames{105}; + int32_t next_chunk_mel_frames{112}; + int32_t first_shift_mel_frames{105}; + int32_t next_shift_mel_frames{112}; + int32_t first_pre_encode_cache_mel_frames{0}; + int32_t next_pre_encode_cache_mel_frames{9}; + int32_t valid_encoder_frames{14}; + int32_t drop_extra_pre_encoded{2}; +}; + +inline TdtStreamingSchedule make_tdt_streaming_schedule(int32_t left, int32_t right, + int32_t sample_rate = 16000, + int32_t hop = 160, + int32_t subsampling = 8) { + if (left <= 0 || right < 0 || sample_rate <= 0 || hop <= 0 || subsampling <= 0) + throw std::invalid_argument("invalid private TDT streaming schedule"); + TdtStreamingSchedule out; + out.att_context_left = left; + out.att_context_right = right; + out.subsampling_factor = subsampling; + out.encoder_frame_ms = hop * subsampling * 1000 / sample_rate; + out.valid_encoder_frames = right + 1; + out.chunk_ms = out.valid_encoder_frames * out.encoder_frame_ms; + out.chunk_samples = sample_rate * out.chunk_ms / 1000; + out.first_chunk_mel_frames = 1 + subsampling * right; + out.next_chunk_mel_frames = subsampling * out.valid_encoder_frames; + out.first_shift_mel_frames = out.first_chunk_mel_frames; + out.next_shift_mel_frames = out.next_chunk_mel_frames; + out.next_pre_encode_cache_mel_frames = subsampling + 1; + out.drop_extra_pre_encoded = out.next_pre_encode_cache_mel_frames - (subsampling - 1); + return out; +} + +inline void validate_tdt_mel_geometry(const TdtConfig& config, int32_t filterbank_freq_bins, + int32_t filterbank_mel_bins) { + if (config.sample_rate <= 0 || config.mel_hop_length <= 0 || config.mel_n_fft <= 0 || + config.mel_chunk_length <= 0 || config.mel_length <= 0 || config.num_mel_bins <= 0) { + throw std::invalid_argument("invalid TDT mel dimensions in config"); + } + if (filterbank_freq_bins != config.mel_n_fft / 2 + 1 || + filterbank_mel_bins != config.num_mel_bins) { + throw std::invalid_argument("TDT mel filterbank dimensions do not match config"); + } +} + +inline TdtGreedyDecision make_tdt_greedy_decision(int32_t token_id, int32_t duration_index, + const std::vector& duration_values, + int32_t blank_id) { + if (duration_index < 0 || duration_index >= static_cast(duration_values.size())) + throw std::out_of_range("TDT duration index is outside the configured duration table"); + int32_t advance = duration_values[static_cast(duration_index)]; + if (advance < 0) + throw std::invalid_argument("TDT durations must be non-negative"); + const bool emit = token_id != blank_id; + if (advance == 0 && !emit) + advance = 1; + return {emit, advance}; +} + +} // namespace trtmc diff --git a/families/parakeet_tdt/runtime/tokenizer.h b/families/parakeet_tdt/runtime/tokenizer.h new file mode 100644 index 0000000000..072b12291b --- /dev/null +++ b/families/parakeet_tdt/runtime/tokenizer.h @@ -0,0 +1,30 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace trtmc::parakeet_tdt { + +class ITokenizer { + public: + virtual ~ITokenizer() = default; + virtual std::vector encode(const std::string& text) const = 0; + virtual std::string decode(const std::vector& ids) const = 0; + virtual std::int32_t id_for_token(std::string_view token) const = 0; + virtual std::string token_for_id(std::int32_t id) const = 0; +}; + +std::unique_ptr CreateBpeTokenizer(const char* tokenizer_json_data, + std::size_t tokenizer_json_size, + bool add_special_tokens = false); + +} // namespace trtmc::parakeet_tdt diff --git a/families/parakeet_tdt/support.py b/families/parakeet_tdt/support.py new file mode 100644 index 0000000000..b3f825e38b --- /dev/null +++ b/families/parakeet_tdt/support.py @@ -0,0 +1,12 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from tensorrt_model_connect.model_support import family_support + +describe = family_support( + model_types=("parakeet_tdt",), + architectures=("ParakeetForTDT",), + tasks=("speech_transcription",), + default_task="speech_transcription", + default_precision="fp32", +) diff --git a/families/parakeet_tdt/tests/README.md b/families/parakeet_tdt/tests/README.md new file mode 100644 index 0000000000..c985c94c04 --- /dev/null +++ b/families/parakeet_tdt/tests/README.md @@ -0,0 +1,54 @@ +# CPU migration tests + +`config.json` is NVIDIA's configuration from +[`nvidia/parakeet-tdt-0.6b-v3`](https://huggingface.co/nvidia/parakeet-tdt-0.6b-v3/blob/541d1f99c6b0c3cd0b11a95167540bb8edefd82b/config.json), +revision `541d1f99c6b0c3cd0b11a95167540bb8edefd82b`, retrieved September 18, 2026. +The publisher lists the checkpoint under CC BY 4.0. No model weights are included. + +The bundle-composition tests replace TensorRT compilation with byte payloads. +They exercise the real BuildRequest, family dispatch, configuration validation, +and BundleWriter. They do not establish engine validity, numerical equivalence, +transcript quality, GPU execution, or performance. + +`test_native.py` compiles five CPU executables with a C++17 compiler (`CXX`, +default `c++`): Task audio-input validation/downmix/resampling, the existing +mel/FFT/incremental-resampling tests, and the existing TDT duration/geometry +tests, semantic transcription pipeline orchestration with fake engines, and BPE +Metaspace decoding. The original mel and duration assertions from PR #1060 are +retained. Pipeline tests need CUDA headers (`CUDA_HOME` or `CUDA_PATH`) but use +no GPU or CUDA/TensorRT libraries. Tokenizer tests need `nlohmann/json.hpp` +(`NLOHMANN_JSON_INCLUDE_DIR`, default `/usr/include`). Missing headers produce +explicit skips, not passing evidence. + +The native pipeline now connects the audio helper to encoder, predictor, joint, +and tokenizer calls. Tests check request state reset, invalid options, malformed +or missing engine outputs, and recovery after failed requests. + +`test_factory.py` writes real bundles with the shared Python `BundleWriter`, +then compiles and executes the real C++ `BundleReader` and family factory. +Its 22 cases cover semantic task binding, engine section order/content, +invalid identity/config/tokenizer/frontend data, missing sections, failed engine +loads, and destruction of modules after partial construction failures. Engine +loading is faked; CUDA headers and nlohmann JSON are required as above. This +does not prove TensorRT deserialization, engine validity, model transcript parity, +or streaming support. + +An additional September 19, 2026 check compared the pinned official tokenizer +above against Hugging Face `tokenizers==0.22.2`: all 8,193 singleton vocabulary +IDs, 1,000 deterministic random token sequences, and two special/empty cases +matched (9,195 cases). It caught and fixed incorrect conversion of the literal +`Ġ` character to a space. The synthetic regression remains in the repository; +this broader check requires downloading the official tokenizer JSON, not weights. + +From the repository root: + +```sh +PYTHONPATH=core/builder:. python -m pytest families/parakeet_tdt/tests -q +``` + +The selected E2E builds `test_parakeet_tdt_sdk_cpp` in the configured +`TRTMC_NATIVE_BUILD_DIR` before invoking the public SDK consumer. Community GPU +builds the family DSO and registered CTests, so an unregistered consumer cannot +be assumed to exist. `test_gpu_ci.py` checks the selected pinned checkpoint and +uses a real minimal CMake build to verify the consumer helper. Checkpoint +identity tests live in `test_support.py` for public CPU discovery. diff --git a/families/parakeet_tdt/tests/config.json b/families/parakeet_tdt/tests/config.json new file mode 100644 index 0000000000..969c65410a --- /dev/null +++ b/families/parakeet_tdt/tests/config.json @@ -0,0 +1,49 @@ +{ + "architectures": [ + "ParakeetForTDT" + ], + "blank_token_id": 8192, + "decoder_hidden_size": 640, + "dtype": "float32", + "durations": [ + 0, + 1, + 2, + 3, + 4 + ], + "encoder_config": { + "activation_dropout": 0.1, + "attention_bias": false, + "attention_dropout": 0.1, + "conv_kernel_size": 9, + "convolution_bias": false, + "dropout": 0.1, + "dropout_positions": 0.0, + "hidden_act": "silu", + "hidden_size": 1024, + "initializer_range": 0.02, + "intermediate_size": 4096, + "layerdrop": 0.1, + "max_position_embeddings": 5000, + "model_type": "parakeet_encoder", + "num_attention_heads": 8, + "num_hidden_layers": 24, + "num_key_value_heads": 8, + "num_mel_bins": 128, + "scale_input": false, + "subsampling_conv_channels": 256, + "subsampling_conv_kernel_size": 3, + "subsampling_conv_stride": 2, + "subsampling_factor": 8 + }, + "hidden_act": "relu", + "initializer_range": 0.02, + "is_encoder_decoder": true, + "max_symbols_per_step": 10, + "model_type": "parakeet_tdt", + "num_decoder_layers": 2, + "pad_token_id": 2, + "transformers_version": "5.6.0.dev0", + "vocab_size": 8193 +} diff --git a/families/parakeet_tdt/tests/cpp/fake_module.h b/families/parakeet_tdt/tests/cpp/fake_module.h new file mode 100644 index 0000000000..8385b47662 --- /dev/null +++ b/families/parakeet_tdt/tests/cpp/fake_module.h @@ -0,0 +1,70 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once +#include "trtmc/runtime/trt_module.h" +using namespace trtmc; + +class FakeModule final : public ITrtModule { + public: + explicit FakeModule(int role) : role_(role) { ++live_instances; } + ~FakeModule() override { --live_instances; } + inline static int live_instances{0}; + int calls{0}; + bool reset_seen{false}; + bool malformed{false}; + bool missing{false}; + TensorMap forward(const TensorMap& in) override { + ++calls; + if (missing) + return {}; + if (role_ == 0) + return { + {"encoder_output", {values_, {1, 2}, malformed ? DType::kInt32 : DType::kFloat32}}}; + if (role_ == 1) { + if (*static_cast(in.at("token_id").data) == 2) + reset_seen = *static_cast(in.at("state_h_0").data) == 0; + return {{"pred_output", {values_, {1, malformed ? 1 : 2}, DType::kFloat32}}, + {"next_h_0", {values_, {1, 2}, DType::kFloat32}}, + {"next_c_0", {values_, {1, 2}, DType::kFloat32}}}; + } + return {{"token_logits", {malformed ? nullptr : tokens_, {1, 3}, DType::kFloat32}}, + {"duration_logits", {durations_, {1, 2}, DType::kFloat32}}}; + } + DeviceTensorMap forward_device(const DeviceTensorMap&) override { return {}; } + void forward_device_async(const DeviceTensorMap&) override {} + void forward_async(const TensorMap&) override {} + void sync() override {} + cudaStream_t stream() const override { return nullptr; } + void enable_cuda_graph() override {} + bool cuda_graph_active() const override { return false; } + bool cuda_graph_captured() const override { return false; } + int32_t profile_idx() const override { return 0; } + std::vector input_info() const override { return {}; } + std::vector output_info() const override { return {}; } + bool has_input(const std::string&) const override { return false; } + bool has_output(const std::string&) const override { return true; } + DType tensor_dtype(const std::string&) const override { return DType::kFloat32; } + std::vector tensor_shape(const std::string&) const override { return {}; } + std::vector input_profile_shape(const std::string&, int32_t, + ProfileShapeSelector) const override { + return {}; + } + int32_t optimization_profile_count() const override { return 1; } + void* device_ptr(const std::string&) const override { return nullptr; } + void bind_external(const std::string&, void*) override {} + void bind_external(const std::string&, void*, const std::vector&) override {} + int32_t input_rank(const std::string&) const override { return 0; } + bool input_is_dynamic(const std::string&) const override { return false; } + void reset_execution_context() override {} + void set_timing_label(std::string) override {} + bool ok() const override { return true; } + void keep_alive(std::shared_ptr) override {} + + private: + int role_; + float values_[2]{1, 1}; + float tokens_[3]{2, 0, -1}; + float durations_[2]{0, 1}; +}; diff --git a/families/parakeet_tdt/tests/cpp/test_audio_helpers.cpp b/families/parakeet_tdt/tests/cpp/test_audio_helpers.cpp new file mode 100644 index 0000000000..b017b5e066 --- /dev/null +++ b/families/parakeet_tdt/tests/cpp/test_audio_helpers.cpp @@ -0,0 +1,264 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Unit tests for TDT-owned audio feature extraction helpers. + +#include "families/parakeet_tdt/runtime/audio_helpers.h" +#include "families/parakeet_tdt/runtime/resampler.h" + +#include +#include +#include +#include +#include + +namespace { + +int g_failures = 0; + +void check(bool condition, const char* name) { + if (!condition) { + std::cerr << "FAIL: " << name << '\n'; + ++g_failures; + } +} + +void check_close(float actual, float expected, float tolerance, const char* name) { + if (std::fabs(actual - expected) > tolerance) { + std::cerr << "FAIL: " << name << " actual=" << actual << " expected=" << expected << '\n'; + ++g_failures; + } +} + +std::vector make_identity_filterbank(int32_t n_freq_bins) { + std::vector filters(static_cast(n_freq_bins) * n_freq_bins, 0.0F); + for (int32_t i = 0; i < n_freq_bins; ++i) { + filters[static_cast(i) * n_freq_bins + i] = 1.0F; + } + return filters; +} + +std::vector reference_rfft_power(const std::vector& input) { + const int32_t n = static_cast(input.size()); + const int32_t n_out = n / 2 + 1; + const double pi2 = 2.0 * 3.14159265358979323846; + std::vector power(static_cast(n_out), 0.0F); + for (int32_t k = 0; k < n_out; ++k) { + double real = 0.0; + double imag = 0.0; + for (int32_t t = 0; t < n; ++t) { + const double angle = + pi2 * static_cast(k) * static_cast(t) / static_cast(n); + real += static_cast(input[static_cast(t)]) * std::cos(angle); + imag -= static_cast(input[static_cast(t)]) * std::sin(angle); + } + power[static_cast(k)] = static_cast(real * real + imag * imag); + } + return power; +} + +void check_vectors_close(const std::vector& actual, const std::vector& expected, + float tolerance, const char* name) { + check(actual.size() == expected.size(), name); + for (std::size_t i = 0; i < std::min(actual.size(), expected.size()); ++i) { + const float scaled_tolerance = tolerance * std::max(1.0F, std::abs(expected[i])); + check_close(actual[i], expected[i], scaled_tolerance, name); + } +} + +void test_tdt_fft_matches_direct_dft() { + std::vector input(512); + for (std::size_t i = 0; i < input.size(); ++i) { + input[i] = static_cast(std::sin(static_cast(i) * 0.17) + + 0.25 * std::cos(static_cast(i) * 0.07)); + } + check_vectors_close(trtmc::tdt::detail::rfft_power(input), reference_rfft_power(input), 1e-5F, + "tdt 512-point FFT matches direct DFT"); +} + +trtmc::tdt::MelSpectrogramOptions make_incremental_options() { + trtmc::tdt::MelSpectrogramOptions options; + options.n_fft = 8; + options.win_length = 6; + options.hop_length = 2; + options.chunk_length_s = 1; + options.sample_rate = 32; + options.symmetric_window = true; + options.center_window_in_fft = true; + options.preemphasis = 0.97F; + options.log_scale = trtmc::tdt::MelLogScale::kNaturalLog; + options.normalize_per_feature = false; + return options; +} + +std::vector make_test_audio(int32_t count) { + std::vector audio(static_cast(count)); + for (int32_t i = 0; i < count; ++i) { + audio[static_cast(i)] = + static_cast(0.35 * std::sin(static_cast(i) * 0.47) + + 0.2 * std::cos(static_cast(i) * 0.19)); + } + return audio; +} + +std::vector materialize_incremental(const trtmc::tdt::IncrementalMelSpectrogram& state) { + std::vector output; + output.reserve(static_cast(state.n_mels()) * state.frame_count()); + for (int32_t mel = 0; mel < state.n_mels(); ++mel) { + for (int32_t frame = 0; frame < state.frame_count(); ++frame) { + output.push_back(state.value(mel, frame)); + } + } + return output; +} + +std::vector run_resampled_chunks(const std::vector& source_audio, + const std::vector& chunk_sizes, + trtmc::tdt::IncrementalMelStats& stats) { + const auto options = make_incremental_options(); + const auto filters = make_identity_filterbank(5); + trtmc::tdt::IncrementalMelSpectrogram state(filters.data(), 5, 5, options, 96); + int32_t accepted = 0; + for (const int32_t chunk_size : chunk_sizes) { + const int32_t take = + std::min(chunk_size, static_cast(source_audio.size()) - accepted); + state.accept_audio(source_audio.data() + accepted, take); + accepted += take; + const bool final = accepted == static_cast(source_audio.size()); + if (final) { + state.ensure_frames(16, true); + break; + } + + const int32_t stable_target = state.stable_target_sample_count(); + const int32_t stable_frames = + stable_target >= options.n_fft / 2 + ? 1 + (stable_target - options.n_fft / 2) / options.hop_length + : 0; + state.ensure_frames(stable_frames, false); + } + stats = state.stats(); + return materialize_incremental(state); +} + +void test_incremental_mel_processes_each_frame_once() { + const auto options = make_incremental_options(); + const auto filters = make_identity_filterbank(5); + const auto audio = make_test_audio(32); + const auto offline = trtmc::tdt::extract_configured_mel_spectrogram( + audio.data(), static_cast(audio.size()), filters.data(), 5, 5, options); + + trtmc::tdt::IncrementalMelSpectrogram state(filters.data(), 5, 5, options, 32); + state.accept_audio(audio.data(), 8); + state.ensure_frames(2, false); + state.accept_audio(audio.data() + 8, 8); + state.ensure_frames(6, false); + state.accept_audio(audio.data() + 16, 16); + state.ensure_frames(16, true); + + check_vectors_close(materialize_incremental(state), offline.data, 1e-5F, + "incremental mel matches one-shot features"); + const auto stats = state.stats(); + check(stats.accepted_source_samples == 32, "incremental mel accepts each source sample once"); + check(stats.generated_resampled_samples == 32, + "incremental mel materializes each identity-rate sample once"); + check(stats.computed_mel_frames == 16, "incremental mel computes each frame once"); +} + +void test_incremental_resample_matches_one_shot() { + const auto options = make_incremental_options(); + const auto filters = make_identity_filterbank(5); + const auto source_audio = make_test_audio(96); + const auto resampled = trtmc::tdt::resample_linear(source_audio.data(), 96, 96, 32); + const auto offline = trtmc::tdt::extract_configured_mel_spectrogram( + resampled.data(), static_cast(resampled.size()), filters.data(), 5, 5, options); + + trtmc::tdt::IncrementalMelSpectrogram state(filters.data(), 5, 5, options, 96); + state.accept_audio(source_audio.data(), 72); + state.ensure_frames(6, false); + state.accept_audio(source_audio.data() + 72, 24); + state.ensure_frames(16, true); + + check_vectors_close(materialize_incremental(state), offline.data, 1e-5F, + "incremental resampling matches one-shot features"); + const auto stats = state.stats(); + check(stats.accepted_source_samples == 96, + "incremental resampling accepts each source sample once"); + check(stats.generated_resampled_samples == 32, + "incremental resampling generates each target sample once"); + check(stats.computed_mel_frames == 16, "incremental resampling computes each mel frame once"); +} + +void test_incremental_resample_is_chunk_size_independent() { + const auto source_audio = make_test_audio(96); + trtmc::tdt::IncrementalMelStats coarse_stats; + trtmc::tdt::IncrementalMelStats irregular_stats; + const auto coarse = run_resampled_chunks(source_audio, {24, 24, 24, 24}, coarse_stats); + const auto irregular = run_resampled_chunks(source_audio, {17, 31, 11, 37}, irregular_stats); + + check_vectors_close(irregular, coarse, 0.0F, + "incremental features are independent of source chunk sizes"); + check(coarse_stats.accepted_source_samples == 96 && + irregular_stats.accepted_source_samples == 96, + "chunk schedules accept each source sample once"); + check(coarse_stats.generated_resampled_samples == 32 && + irregular_stats.generated_resampled_samples == 32, + "chunk schedules generate each resampled sample once"); + check(coarse_stats.computed_mel_frames == 16 && irregular_stats.computed_mel_frames == 16, + "chunk schedules compute each mel frame once"); +} + +void test_tdt_mel_matches_configured_owned_options() { + const std::vector samples{0.0F, 0.1F, -0.2F, 0.4F, -0.1F, 0.0F, 0.2F, -0.3F}; + const int32_t n_fft = 4; + const int32_t n_freq_bins = 3; + const int32_t n_mel_bins = 2; + const std::vector mel_filters{ + 0.5F, 0.1F, 0.2F, 0.7F, 0.3F, 0.4F, + }; + + const auto tdt = trtmc::tdt::extract_tdt_mel_spectrogram( + samples.data(), static_cast(samples.size()), mel_filters.data(), n_freq_bins, + n_mel_bins, n_fft, /*win_length=*/3, /*hop_length=*/2, /*chunk_length_s=*/1, + /*sample_rate=*/8, /*preemphasis=*/0.97F); + + trtmc::tdt::MelSpectrogramOptions options; + options.n_fft = n_fft; + options.win_length = 3; + options.hop_length = 2; + options.chunk_length_s = 1; + options.sample_rate = 8; + options.symmetric_window = true; + options.center_window_in_fft = true; + options.preemphasis = 0.97F; + options.log_scale = trtmc::tdt::MelLogScale::kNaturalLog; + options.normalize_per_feature = true; + const auto configured = trtmc::tdt::extract_configured_mel_spectrogram( + samples.data(), static_cast(samples.size()), mel_filters.data(), n_freq_bins, + n_mel_bins, options); + + check(tdt.n_mels == configured.n_mels, "tdt mel keeps configured mel count"); + check(tdt.n_frames == configured.n_frames, "tdt mel keeps configured frame count"); + check(tdt.data.size() == configured.data.size(), "tdt mel keeps configured data size"); + for (std::size_t i = 0; i < tdt.data.size(); ++i) { + check_close(tdt.data[i], configured.data[i], 1e-6F, "tdt mel value matches configured"); + } +} + +} // namespace + +int main() { + test_tdt_fft_matches_direct_dft(); + test_tdt_mel_matches_configured_owned_options(); + test_incremental_mel_processes_each_frame_once(); + test_incremental_resample_matches_one_shot(); + test_incremental_resample_is_chunk_size_independent(); + + if (g_failures != 0) { + std::cerr << g_failures << " tdt audio helper test(s) failed\n"; + return 1; + } + return 0; +} diff --git a/families/parakeet_tdt/tests/cpp/test_audio_input.cpp b/families/parakeet_tdt/tests/cpp/test_audio_input.cpp new file mode 100644 index 0000000000..ba920a8c1e --- /dev/null +++ b/families/parakeet_tdt/tests/cpp/test_audio_input.cpp @@ -0,0 +1,79 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "families/parakeet_tdt/runtime/audio_input.h" + +#include +#include +#include +#include +#include + +int failures = 0; +void check(bool value, const char* label) { + if (!value) { + std::cerr << label << '\n'; + ++failures; + } +} +template +void rejects(F run, const char* label) { + try { + run(); + check(false, label); + } catch (const std::invalid_argument&) { + } +} + +int main() { + using trtmc::internal::SpeechTranscriptionRequest; + using trtmc::tdt::prepare_audio; + float stereo[] = {1, 3, -2, 2, 4, 2}; + SpeechTranscriptionRequest request; + request.audio = {{stereo, 6}, 16000, 2}; + auto mono = prepare_audio(request); + check(mono == std::vector({2, 0, 3}), "downmix interleaved frames"); + check(stereo[0] == 1 && stereo[1] == 3, "preserve borrowed input"); + request.audio.sample_rate.reset(); + check(prepare_audio(request) == mono, "omitted rate uses model 16 kHz default"); + request.audio.channels = 0; + rejects([&] { prepare_audio(request); }, "reject missing channels"); + request.audio.channels = 4; + rejects([&] { prepare_audio(request); }, "reject partial frame"); + request.audio.channels = 2; + request.audio.sample_rate = 0; + rejects([&] { prepare_audio(request); }, "reject zero rate"); + request.audio.sample_rate = std::numeric_limits::max(); + rejects([&] { prepare_audio(request); }, "reject out-of-range rate"); + request.audio.sample_rate = 16000; + request.source_language = "en"; + rejects([&] { prepare_audio(request); }, "reject unsupported language control"); + request.source_language.reset(); + request.audio.samples = {nullptr, 6}; + rejects([&] { prepare_audio(request); }, "reject null borrowed buffer"); + request.audio.samples = {}; + rejects([&] { prepare_audio(request); }, "reject empty input"); + request.audio.samples = {stereo, 6}; + stereo[0] = std::numeric_limits::quiet_NaN(); + rejects([&] { prepare_audio(request); }, "reject NaN"); + stereo[0] = std::numeric_limits::infinity(); + rejects([&] { prepare_audio(request); }, "reject infinity"); + stereo[0] = stereo[1] = std::numeric_limits::max(); + check(std::isfinite(prepare_audio(request)[0]), "downmix avoids intermediate float overflow"); + std::vector pcm(48000, 0.25F); + request.audio = {{pcm.data(), pcm.size()}, 48000, 1}; + auto resampled = prepare_audio(request); + check(resampled.size() == 16000, "resample 48 kHz to 16 kHz"); + check(std::abs(resampled[8000] - 0.25F) < 1e-6F, "resampling preserves constant signal"); + pcm.resize(16000 * 30); + request.audio = {{pcm.data(), pcm.size()}, 16000, 1}; + check(prepare_audio(request).size() == pcm.size(), "accept exact 30-second window"); + pcm.push_back(0); + request.audio.samples = {pcm.data(), pcm.size()}; + rejects([&] { prepare_audio(request); }, "reject overlong input instead of truncating"); + request.audio.samples = {stereo, static_cast(std::numeric_limits::max()) + 1}; + rejects([&] { prepare_audio(request); }, "reject sample count before narrowing or dereference"); + return failures == 0 ? 0 : 1; +} diff --git a/families/parakeet_tdt/tests/cpp/test_decode_policy.cpp b/families/parakeet_tdt/tests/cpp/test_decode_policy.cpp new file mode 100644 index 0000000000..5897c9e1b5 --- /dev/null +++ b/families/parakeet_tdt/tests/cpp/test_decode_policy.cpp @@ -0,0 +1,111 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "families/parakeet_tdt/runtime/tdt_config.h" + +#include +#include + +namespace { + +int failures = 0; + +void check(bool condition, const char* message) { + if (!condition) { + std::cerr << "FAIL: " << message << '\n'; + ++failures; + } +} + +void test_duration_value_controls_frame_advance() { + const auto decision = trtmc::make_tdt_greedy_decision(7, 2, {0, 1, 3}, 9); + check(decision.emit_token, "nonblank should emit"); + check(decision.frame_advance == 3, "duration table value should advance frames"); +} + +void test_zero_duration_blank_forces_progress_without_emit() { + const auto decision = trtmc::make_tdt_greedy_decision(9, 0, {0, 1, 3}, 9); + check(!decision.emit_token, "blank should not emit"); + check(decision.frame_advance == 1, "zero-duration blank must force progress"); +} + +void test_zero_duration_nonblank_stays_on_frame() { + const auto decision = trtmc::make_tdt_greedy_decision(7, 0, {0, 1, 3}, 9); + check(decision.emit_token, "zero-duration nonblank should emit"); + check(decision.frame_advance == 0, "zero-duration nonblank must stay on the frame"); +} + +void test_invalid_duration_index_is_rejected() { + bool threw = false; + try { + (void)trtmc::make_tdt_greedy_decision(7, 3, {0, 1, 2}, 9); + } catch (const std::out_of_range&) { + threw = true; + } + check(threw, "duration indices outside the table must throw"); +} + +void test_negative_duration_is_rejected() { + bool threw = false; + try { + (void)trtmc::make_tdt_greedy_decision(7, 1, {0, -1, 2}, 9); + } catch (const std::invalid_argument&) { + threw = true; + } + check(threw, "negative duration values must throw"); +} + +void test_streaming_schedule_derives_timing_and_cache_geometry() { + const auto schedule = trtmc::make_tdt_streaming_schedule(20, 3, 8000, 80, 4); + check(schedule.encoder_frame_ms == 40, "encoder frame timing derives from hop and stride"); + check(schedule.chunk_samples == 1280, "chunk samples derive from schedule timing"); + check(schedule.next_pre_encode_cache_mel_frames == 5, + "pre-encode cache derives from subsampling"); + check(schedule.drop_extra_pre_encoded == 2, "pre-encode drop derives from cache overlap"); +} + +void test_mel_geometry_accepts_matching_dimensions() { + trtmc::TdtConfig config; + trtmc::validate_tdt_mel_geometry(config, 257, 128); +} + +void test_mel_geometry_rejects_invalid_or_mismatched_dimensions() { + trtmc::TdtConfig config; + config.mel_n_fft = 0; + bool invalid_threw = false; + try { + trtmc::validate_tdt_mel_geometry(config, 257, 128); + } catch (const std::invalid_argument&) { + invalid_threw = true; + } + check(invalid_threw, "non-positive mel dimensions must throw"); + + config.mel_n_fft = 512; + bool mismatch_threw = false; + try { + trtmc::validate_tdt_mel_geometry(config, 256, 128); + } catch (const std::invalid_argument&) { + mismatch_threw = true; + } + check(mismatch_threw, "filterbank dimensions must match the runtime config"); +} + +} // namespace + +int main() { + test_duration_value_controls_frame_advance(); + test_zero_duration_blank_forces_progress_without_emit(); + test_zero_duration_nonblank_stays_on_frame(); + test_invalid_duration_index_is_rejected(); + test_negative_duration_is_rejected(); + test_streaming_schedule_derives_timing_and_cache_geometry(); + test_mel_geometry_accepts_matching_dimensions(); + test_mel_geometry_rejects_invalid_or_mismatched_dimensions(); + if (failures) { + std::cerr << failures << " TDT decode policy test(s) failed\n"; + return 1; + } + return 0; +} diff --git a/families/parakeet_tdt/tests/cpp/test_factory.cpp b/families/parakeet_tdt/tests/cpp/test_factory.cpp new file mode 100644 index 0000000000..fdee4c999a --- /dev/null +++ b/families/parakeet_tdt/tests/cpp/test_factory.cpp @@ -0,0 +1,70 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#include "families/parakeet_tdt/tests/cpp/fake_module.h" +#include "trtmc/internal/model.h" +#include "trtmc/runtime/family_factory.h" +#include "trtmc/runtime/trt_backend.h" + +#include +#include + +class FakeBackend final : public IBackend { + public: + int calls{0}; + int fail_at{0}; + std::unique_ptr create_module(const void* data, size_t size, + const ModuleCreateOptions&) override { + const std::string payload(static_cast(data), size); + const char* expected[] = {"encoder", "predictor", "joint"}; + if (calls >= 3 || payload != expected[calls]) + throw std::runtime_error("incorrect engine section order or content"); + ++calls; + return calls == fail_at ? nullptr : std::make_unique(calls - 1); + } + std::unique_ptr + create_module_prebound(const void*, size_t, const ModuleCreateOptions&, + const std::vector&) override { + throw std::runtime_error("unexpected prebound call"); + } + BackendDualProfileModules create_dual_profile_modules(const void*, size_t, + const ModuleCreateOptions&) override { + throw std::runtime_error("unexpected dual profile call"); + } + const char* name() const override { return "trt"; } +}; + +int main(int argc, char** argv) { + if (argc != 5) + return 1; + FakeBackend backend; + const std::string mode = argv[4]; + backend.fail_at = mode == "fail-engine" ? 1 + : mode == "fail-predictor" ? 2 + : mode == "fail-joint" ? 3 + : 0; + const auto kv = std::string(argv[4]) == "kv" ? 1U : 0U; + const std::string expected_error = argv[2]; + try { + BundleReader reader(argv[1]); + std::unique_ptr task(trtmc_create_family({reader, backend, kv})); + if (!expected_error.empty()) + throw std::logic_error("factory accepted invalid bundle"); + auto* model = dynamic_cast(task.get()); + if (!model || model->task_bindings().size() != 1 || + model->task_bindings()[0].key.id != "speech_transcription") + return 2; + } catch (const std::exception& error) { + if (expected_error.empty() || + std::string(error.what()).find(expected_error) == std::string::npos) { + std::cerr << error.what() << '\n'; + return 3; + } + } + if (backend.calls != std::stoi(argv[3])) + return 4; + if (FakeModule::live_instances != 0) + return 5; + std::cout << "factory bundle contract passed\n"; +} diff --git a/families/parakeet_tdt/tests/cpp/test_pipeline.cpp b/families/parakeet_tdt/tests/cpp/test_pipeline.cpp new file mode 100644 index 0000000000..e704d3de32 --- /dev/null +++ b/families/parakeet_tdt/tests/cpp/test_pipeline.cpp @@ -0,0 +1,90 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#include "families/parakeet_tdt/runtime/pipeline.h" +#include "families/parakeet_tdt/tests/cpp/fake_module.h" + +#include +#include + +using namespace trtmc; +using namespace trtmc::parakeet_tdt; + +class FakeTokenizer final : public ITokenizer { + public: + std::vector encode(const std::string&) const override { return {}; } + std::string decode(const std::vector& ids) const override { + return ids == std::vector{0} ? "hello" : "unexpected"; + } + int32_t id_for_token(std::string_view) const override { return -1; } + std::string token_for_id(int32_t) const override { return {}; } +}; + +int main() { + auto encoder = std::make_unique(0); + auto predictor = std::make_unique(1); + auto joint = std::make_unique(2); + auto* enc = encoder.get(); + auto* pred = predictor.get(); + auto* joint_ptr = joint.get(); + TdtConfig cfg; + cfg.num_mel_bins = 2; + cfg.mel_n_fft = 4; + cfg.mel_win_length = 4; + cfg.mel_hop_length = 2; + cfg.mel_length = 8; + cfg.encoder_seq_len = 1; + cfg.encoder_hidden_size = 2; + cfg.pred_hidden_size = 2; + cfg.pred_num_layers = 1; + cfg.blank_id = 2; + cfg.vocab_size = 2; + cfg.duration_values = {0, 1}; + MelFilterbank filters{{1, 1, 1, 1, 1, 1}, 3, 2}; + TdtPipeline pipeline(std::move(encoder), std::move(predictor), std::move(joint), cfg, filters, + std::make_shared()); + auto bindings = pipeline.task_bindings(); + if (bindings.size() != 1 || bindings[0].key.id != "speech_transcription") + return 1; + auto* task = static_cast(bindings[0].implementation); + float pcm[16]{}; + internal::SpeechTranscriptionRequest req{{{pcm, 16}, 16000, 1}, {}}; + for (int i = 0; i < 2; ++i) { + pred->reset_seen = false; + auto result = task->run(req, {}); + if (result.text != "hello" || result.token_ids != std::vector{0} || + !pred->reset_seen) + return 2; + } + internal::ConfigEntry bad{"max_new_tokens", int64_t{0}}; + const int before = enc->calls; + try { + task->run(req, {&bad, 1}); + return 3; + } catch (const std::invalid_argument&) { + } + if (enc->calls != before) + return 4; + for (auto* module : {enc, pred, joint_ptr}) { + module->malformed = true; + try { + task->run(req, {}); + return 5; + } catch (const std::runtime_error&) { + } + module->malformed = false; + module->missing = true; + try { + task->run(req, {}); + return 6; + } catch (const std::runtime_error&) { + } + module->missing = false; + } + // A failed request must not poison the predictor state of the next request. + pred->reset_seen = false; + if (task->run(req, {}).text != "hello" || !pred->reset_seen) + return 7; + std::cout << "semantic pipeline orchestration passed (fake engines)\n"; +} diff --git a/families/parakeet_tdt/tests/cpp/test_tokenizer.cpp b/families/parakeet_tdt/tests/cpp/test_tokenizer.cpp new file mode 100644 index 0000000000..868dda1400 --- /dev/null +++ b/families/parakeet_tdt/tests/cpp/test_tokenizer.cpp @@ -0,0 +1,39 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#include "families/parakeet_tdt/runtime/tokenizer.h" + +#include +#include +#include +#include + +int main(int argc, char** argv) { + using namespace trtmc::parakeet_tdt; + std::string source = + R"({"model":{"type":"BPE","vocab":{"":0,"":1,"▁hello":2,"▁world":3,"!":4,"▁你好":5,"Ġ":6},"merges":[]},"added_tokens":[{"id":0,"content":"","special":true},{"id":1,"content":"","special":true}],"decoder":{"type":"Metaspace","replacement":"▁","prepend_scheme":"always","split":true}})"; + nlohmann::json cases = {{{"ids", {6}}, {"text", "Ġ"}}, + {{"ids", {2, 3, 4}}, {"text", "hello world!"}}, + {{"ids", {0, 1, 5, 4}}, {"text", "你好!"}}, + {{"ids", nlohmann::json::array()}, {"text", ""}}}; + if (argc == 3) { + std::ifstream input(argv[1]); + std::ostringstream buffer; + buffer << input.rdbuf(); + source = buffer.str(); + std::ifstream golden(argv[2]); + golden >> cases; + } + auto tokenizer = CreateBpeTokenizer(source.data(), source.size(), false); + if (!tokenizer) + return 1; + for (const auto& item : cases) { + const auto ids = item.at("ids").get>(); + if (tokenizer->decode(ids) != item.at("text").get()) { + std::cerr << "decode mismatch for " << item.at("ids") << '\n'; + return 2; + } + } + std::cout << cases.size() << " tokenizer decode cases passed\n"; +} diff --git a/families/parakeet_tdt/tests/data/Recording.wav b/families/parakeet_tdt/tests/data/Recording.wav new file mode 100644 index 0000000000..db2716e688 Binary files /dev/null and b/families/parakeet_tdt/tests/data/Recording.wav differ diff --git a/families/parakeet_tdt/tests/data/asr_probes/README.md b/families/parakeet_tdt/tests/data/asr_probes/README.md new file mode 100644 index 0000000000..b41549c1cb --- /dev/null +++ b/families/parakeet_tdt/tests/data/asr_probes/README.md @@ -0,0 +1,21 @@ +# ASR Probe Inputs + +These WAV files are generated from the existing ASR fixture at +`../Recording.wav` from the owning model family data folder. They exercise TensorRT Model Connect audio +input handling paths and TRT-vs-reference transcript parity. They are not +intended to benchmark ASR model quality. + +| ID | File | SR | Ch | Duration | Readiness | Model Connect path | +|---|---|---:|---:|---:|---|---| +| probe_01_clean_48k_stereo_baseline | `probe_01_clean_48k_stereo_baseline.wav` | 48000 | 2 | 4.16s | ready_now | WAV decode + stereo input + model-specific resample. | +| probe_02_clean_16k_mono_no_resample | `probe_02_clean_16k_mono_no_resample.wav` | 16000 | 1 | 4.16s | ready_now | WAV decode without model-side downmix; should avoid 48k->16k resample in ASR preprocessing. | +| probe_03_clean_48k_mono_resample | `probe_03_clean_48k_mono_resample.wav` | 48000 | 1 | 4.16s | ready_now | WAV decode + 48k->16k resample, without stereo downmix. | +| probe_04_clean_48k_stereo_gain_skew | `probe_04_clean_48k_stereo_gain_skew.wav` | 48000 | 2 | 4.16s | ready_now | Stereo downmix should be robust when L/R channels are not identical. | +| probe_05_low_volume_48k_stereo | `probe_05_low_volume_48k_stereo.wav` | 48000 | 2 | 4.16s | ready_now | PCM int16 decode + float normalization + ASR preprocessing. | +| probe_06_leading_trailing_silence_48k_stereo | `probe_06_leading_trailing_silence_48k_stereo.wav` | 48000 | 2 | 6.16s | ready_now | Longer waveform decode + silence handling + transcript stability. | +| probe_08_noisy_48k_stereo_snr20 | `probe_08_noisy_48k_stereo_snr20.wav` | 48000 | 2 | 4.16s | nightly_optional | PCM decode + preprocessing under non-clean waveform. | + +Use `manifest.json` for detailed purpose and notes. Re-run +`generate_asr_probe_inputs.py` from this directory to regenerate the WAVs. + + diff --git a/families/parakeet_tdt/tests/data/asr_probes/generate_asr_probe_inputs.py b/families/parakeet_tdt/tests/data/asr_probes/generate_asr_probe_inputs.py new file mode 100644 index 0000000000..ab432fb08c --- /dev/null +++ b/families/parakeet_tdt/tests/data/asr_probes/generate_asr_probe_inputs.py @@ -0,0 +1,234 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Generate ASR probe WAV inputs for TensorRT Model Connect E2E coverage.""" + +from __future__ import annotations + +import json +import shutil +import wave +from pathlib import Path + +import numpy as np + + +ROOT = Path(__file__).resolve().parent +SOURCE = ROOT.parent / "Recording.wav" + + +def read_wav(path: Path) -> tuple[int, np.ndarray]: + with wave.open(str(path), "rb") as wav: + channels = wav.getnchannels() + sample_width = wav.getsampwidth() + sample_rate = wav.getframerate() + frames = wav.readframes(wav.getnframes()) + if sample_width != 2: + raise ValueError(f"Only int16 WAV is supported, got sample_width={sample_width}") + audio_i16 = np.frombuffer(frames, dtype=" None: + audio = np.asarray(audio, dtype=np.float32) + if audio.ndim == 1: + audio = audio[:, None] + audio_i16 = np.clip(audio, -1.0, 1.0) + audio_i16 = np.round(audio_i16 * 32767.0).astype(" np.ndarray: + if audio.ndim == 1: + return audio + return audio.mean(axis=1) + + +def stereo_from_mono(audio: np.ndarray, *, right_gain: float = 1.0) -> np.ndarray: + audio = mono(audio) + return np.stack([audio, audio * right_gain], axis=1) + + +def resample_linear(audio: np.ndarray, src_rate: int, dst_rate: int) -> np.ndarray: + if src_rate == dst_rate: + return audio.copy() + if audio.ndim == 1: + audio_2d = audio[:, None] + else: + audio_2d = audio + src_n = audio_2d.shape[0] + dst_n = int(round(src_n * dst_rate / src_rate)) + src_x = np.arange(src_n, dtype=np.float64) + dst_x = np.linspace(0, src_n - 1, dst_n, dtype=np.float64) + channels = [ + np.interp(dst_x, src_x, audio_2d[:, c]).astype(np.float32) + for c in range(audio_2d.shape[1]) + ] + out = np.stack(channels, axis=1) + return out[:, 0] if audio.ndim == 1 else out + + +def add_white_noise(audio: np.ndarray, snr_db: float, seed: int = 20260611) -> np.ndarray: + rng = np.random.default_rng(seed) + signal_power = float(np.mean(np.square(audio))) or 1e-12 + noise_power = signal_power / (10.0 ** (snr_db / 10.0)) + noise = rng.normal(0.0, np.sqrt(noise_power), size=audio.shape).astype(np.float32) + return np.clip(audio + noise, -1.0, 1.0) + + +def main(output_dir: Path | None = None) -> None: + output_dir = ROOT if output_dir is None else output_dir + output_dir.mkdir(parents=True, exist_ok=True) + sample_rate, original = read_wav(SOURCE) + mono_48k = mono(original) + mono_16k = resample_linear(mono_48k, sample_rate, 16000) + silence_48k = np.zeros((sample_rate,), dtype=np.float32) + cases: list[dict[str, object]] = [] + + def add_case( + name: str, + sample_rate_hz: int, + audio: np.ndarray, + *, + purpose: str, + model_connect_path: str, + readiness: str, + notes: str, + ) -> None: + file_name = f"{name}.wav" + write_wav(output_dir / file_name, sample_rate_hz, audio) + channels = 1 if audio.ndim == 1 else audio.shape[1] + duration_s = round(float(audio.shape[0]) / sample_rate_hz, 3) + cases.append( + { + "id": name, + "file": file_name, + "sample_rate_hz": sample_rate_hz, + "channels": channels, + "duration_s": duration_s, + "purpose": purpose, + "model_connect_path": model_connect_path, + "readiness": readiness, + "contract_oracle": ( + "Compare TRT transcript with reference backend transcript; " + "do not judge model accuracy." + ), + "notes": notes, + } + ) + + shutil.copyfile(SOURCE, output_dir / "probe_01_clean_48k_stereo_baseline.wav") + cases.append( + { + "id": "probe_01_clean_48k_stereo_baseline", + "file": "probe_01_clean_48k_stereo_baseline.wav", + "sample_rate_hz": sample_rate, + "channels": original.shape[1], + "duration_s": round(float(original.shape[0]) / sample_rate, 3), + "purpose": "Current smoke input copied as baseline.", + "model_connect_path": "WAV decode + stereo input + model-specific resample.", + "readiness": "ready_now", + "contract_oracle": ( + "Compare TRT transcript with reference backend transcript; " + "do not judge model accuracy." + ), + "notes": "Expected transcript is the existing short weather sentence.", + } + ) + + add_case( + "probe_02_clean_16k_mono_no_resample", + 16000, + mono_16k, + purpose="Cover native 16k mono path.", + model_connect_path=( + "WAV decode without model-side downmix; should avoid 48k->16k " + "resample in ASR preprocessing." + ), + readiness="ready_now", + notes="Useful to isolate resampling differences from core decode/runtime differences.", + ) + add_case( + "probe_03_clean_48k_mono_resample", + 48000, + mono_48k, + purpose="Cover mono 48k resample path.", + model_connect_path="WAV decode + 48k->16k resample, without stereo downmix.", + readiness="ready_now", + notes="Separates resample handling from stereo downmix handling.", + ) + add_case( + "probe_04_clean_48k_stereo_gain_skew", + 48000, + stereo_from_mono(mono_48k, right_gain=0.55), + purpose="Cover stereo-to-mono with asymmetric channel levels.", + model_connect_path="Stereo downmix should be robust when L/R channels are not identical.", + readiness="ready_now", + notes="Still the same speech content; intended to catch channel handling mistakes.", + ) + add_case( + "probe_05_low_volume_48k_stereo", + 48000, + np.clip(stereo_from_mono(mono_48k) * 0.12, -1.0, 1.0), + purpose="Cover low-amplitude PCM handling.", + model_connect_path="PCM int16 decode + float normalization + ASR preprocessing.", + readiness="ready_now", + notes="Reference and TRT should see the same low-volume input.", + ) + add_case( + "probe_06_leading_trailing_silence_48k_stereo", + 48000, + np.concatenate( + [ + stereo_from_mono(silence_48k), + stereo_from_mono(mono_48k), + stereo_from_mono(silence_48k), + ] + ), + purpose="Cover non-speech padding around valid speech.", + model_connect_path="Longer waveform decode + silence handling + transcript stability.", + readiness="ready_now", + notes="Detects unexpected trimming, transcript drift, or EOS issues.", + ) + add_case( + "probe_08_noisy_48k_stereo_snr20", + 48000, + add_white_noise(stereo_from_mono(mono_48k), snr_db=20.0), + purpose="Cover deterministic noisy input handling.", + model_connect_path="PCM decode + preprocessing under non-clean waveform.", + readiness="nightly_optional", + notes="Use only as TRT-vs-reference parity; do not evaluate model noise robustness.", + ) + + manifest = { + "source_audio": "families/parakeet_tdt/tests/data/Recording.wav", + "source_properties": { + "sample_rate_hz": sample_rate, + "channels": int(original.shape[1]), + "duration_s": round(float(original.shape[0]) / sample_rate, 3), + }, + "principle": ( + "These probes test Model Connect input handling and TRT-vs-reference " + "contract parity, not model accuracy." + ), + "recommended_metrics": [ + "CER for ASR-specific character-level transcript parity", + "WER substitution/insertion/deletion breakdown for ASR failure triage", + "no-speech/blank-audio state reporting for future ASR contract work", + ], + "cases": cases, + } + (output_dir / "manifest.json").write_text( + json.dumps(manifest, indent=2) + "\n", + encoding="utf-8", + ) + + +if __name__ == "__main__": + main() diff --git a/families/parakeet_tdt/tests/data/asr_probes/manifest.json b/families/parakeet_tdt/tests/data/asr_probes/manifest.json new file mode 100644 index 0000000000..8116e9fa5f --- /dev/null +++ b/families/parakeet_tdt/tests/data/asr_probes/manifest.json @@ -0,0 +1,100 @@ +{ + "source_audio": "families/parakeet_tdt/tests/data/Recording.wav", + "source_properties": { + "sample_rate_hz": 48000, + "channels": 2, + "duration_s": 4.16 + }, + "principle": "These probes test Model Connect input handling and TRT-vs-reference contract parity, not model accuracy.", + "recommended_metrics": [ + "CER for ASR-specific character-level transcript parity", + "WER substitution/insertion/deletion breakdown for ASR failure triage", + "no-speech/blank-audio state reporting for future ASR contract work" + ], + "cases": [ + { + "id": "probe_01_clean_48k_stereo_baseline", + "file": "probe_01_clean_48k_stereo_baseline.wav", + "sample_rate_hz": 48000, + "channels": 2, + "duration_s": 4.16, + "purpose": "Current smoke input copied as baseline.", + "model_connect_path": "WAV decode + stereo input + model-specific resample.", + "readiness": "ready_now", + "contract_oracle": "Compare TRT transcript with reference backend transcript; do not judge model accuracy.", + "notes": "Expected transcript is the existing short weather sentence." + }, + { + "id": "probe_02_clean_16k_mono_no_resample", + "file": "probe_02_clean_16k_mono_no_resample.wav", + "sample_rate_hz": 16000, + "channels": 1, + "duration_s": 4.16, + "purpose": "Cover native 16k mono path.", + "model_connect_path": "WAV decode without model-side downmix; should avoid 48k->16k resample in ASR preprocessing.", + "readiness": "ready_now", + "contract_oracle": "Compare TRT transcript with reference backend transcript; do not judge model accuracy.", + "notes": "Useful to isolate resampling differences from core decode/runtime differences." + }, + { + "id": "probe_03_clean_48k_mono_resample", + "file": "probe_03_clean_48k_mono_resample.wav", + "sample_rate_hz": 48000, + "channels": 1, + "duration_s": 4.16, + "purpose": "Cover mono 48k resample path.", + "model_connect_path": "WAV decode + 48k->16k resample, without stereo downmix.", + "readiness": "ready_now", + "contract_oracle": "Compare TRT transcript with reference backend transcript; do not judge model accuracy.", + "notes": "Separates resample handling from stereo downmix handling." + }, + { + "id": "probe_04_clean_48k_stereo_gain_skew", + "file": "probe_04_clean_48k_stereo_gain_skew.wav", + "sample_rate_hz": 48000, + "channels": 2, + "duration_s": 4.16, + "purpose": "Cover stereo-to-mono with asymmetric channel levels.", + "model_connect_path": "Stereo downmix should be robust when L/R channels are not identical.", + "readiness": "ready_now", + "contract_oracle": "Compare TRT transcript with reference backend transcript; do not judge model accuracy.", + "notes": "Still the same speech content; intended to catch channel handling mistakes." + }, + { + "id": "probe_05_low_volume_48k_stereo", + "file": "probe_05_low_volume_48k_stereo.wav", + "sample_rate_hz": 48000, + "channels": 2, + "duration_s": 4.16, + "purpose": "Cover low-amplitude PCM handling.", + "model_connect_path": "PCM int16 decode + float normalization + ASR preprocessing.", + "readiness": "ready_now", + "contract_oracle": "Compare TRT transcript with reference backend transcript; do not judge model accuracy.", + "notes": "Reference and TRT should see the same low-volume input." + }, + { + "id": "probe_06_leading_trailing_silence_48k_stereo", + "file": "probe_06_leading_trailing_silence_48k_stereo.wav", + "sample_rate_hz": 48000, + "channels": 2, + "duration_s": 6.16, + "purpose": "Cover non-speech padding around valid speech.", + "model_connect_path": "Longer waveform decode + silence handling + transcript stability.", + "readiness": "ready_now", + "contract_oracle": "Compare TRT transcript with reference backend transcript; do not judge model accuracy.", + "notes": "Detects unexpected trimming, transcript drift, or EOS issues." + }, + { + "id": "probe_08_noisy_48k_stereo_snr20", + "file": "probe_08_noisy_48k_stereo_snr20.wav", + "sample_rate_hz": 48000, + "channels": 2, + "duration_s": 4.16, + "purpose": "Cover deterministic noisy input handling.", + "model_connect_path": "PCM decode + preprocessing under non-clean waveform.", + "readiness": "nightly_optional", + "contract_oracle": "Compare TRT transcript with reference backend transcript; do not judge model accuracy.", + "notes": "Use only as TRT-vs-reference parity; do not evaluate model noise robustness." + } + ] +} diff --git a/families/parakeet_tdt/tests/manifests/parakeet-tdt-0.6b-v3.json b/families/parakeet_tdt/tests/manifests/parakeet-tdt-0.6b-v3.json new file mode 100644 index 0000000000..c16922082d --- /dev/null +++ b/families/parakeet_tdt/tests/manifests/parakeet-tdt-0.6b-v3.json @@ -0,0 +1,53 @@ +{ + "name": "parakeet-tdt-0.6b-v3", + "hf_id": "nvidia/parakeet-tdt-0.6b-v3", + "hf_revision": "541d1f99c6b0c3cd0b11a95167540bb8edefd82b", + "bundle": "parakeet-tdt-0.6b-v3.bundle", + "family": "parakeet_tdt", + "precision": "fp16", + "testcases": [ + { + "name": "parakeet-tdt-0.6b-v3", + "max_new_tokens": 50, + "test_input_audio": "data/Recording.wav", + "premerge": true + }, + { + "name": "parakeet-tdt-0.6b-v3-asr-probe01", + "max_new_tokens": 50, + "test_input_audio": "data/asr_probes/probe_01_clean_48k_stereo_baseline.wav" + }, + { + "name": "parakeet-tdt-0.6b-v3-asr-probe02", + "max_new_tokens": 50, + "test_input_audio": "data/asr_probes/probe_02_clean_16k_mono_no_resample.wav" + }, + { + "name": "parakeet-tdt-0.6b-v3-asr-probe03", + "max_new_tokens": 50, + "test_input_audio": "data/asr_probes/probe_03_clean_48k_mono_resample.wav" + }, + { + "name": "parakeet-tdt-0.6b-v3-asr-probe04", + "max_new_tokens": 50, + "test_input_audio": "data/asr_probes/probe_04_clean_48k_stereo_gain_skew.wav" + }, + { + "name": "parakeet-tdt-0.6b-v3-asr-probe05", + "max_new_tokens": 50, + "test_input_audio": "data/asr_probes/probe_05_low_volume_48k_stereo.wav" + }, + { + "name": "parakeet-tdt-0.6b-v3-asr-probe06", + "max_new_tokens": 50, + "test_input_audio": "data/asr_probes/probe_06_leading_trailing_silence_48k_stereo.wav" + }, + { + "name": "parakeet-tdt-0.6b-v3-asr-probe08", + "max_new_tokens": 50, + "test_input_audio": "data/asr_probes/probe_08_noisy_48k_stereo_snr20.wav" + } + ], + "task": "speech_transcription", + "tensor_parallel_size": 1 +} diff --git a/families/parakeet_tdt/tests/sdk_consumer.cpp b/families/parakeet_tdt/tests/sdk_consumer.cpp new file mode 100644 index 0000000000..f681732162 --- /dev/null +++ b/families/parakeet_tdt/tests/sdk_consumer.cpp @@ -0,0 +1,46 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#include +#include +#include +#include +#include +#include + +int main(int argc, char** argv) { + if (argc != 7) { + std::cerr << "Usage: consumer BUNDLE RUNTIME_ROOT PCM_F32 RATE CHANNELS MAX_TOKENS\n"; + return 2; + } + try { + const auto rate = std::stoul(argv[4]); + const auto channels = std::stoul(argv[5]); + const auto limit = std::stoll(argv[6]); + if (rate > UINT32_MAX || !rate || channels > UINT32_MAX || !channels || limit <= 0) + throw std::invalid_argument("invalid audio metadata or token limit"); + std::ifstream input(argv[3], std::ios::binary | std::ios::ate); + const auto bytes = input.tellg(); + if (!input || bytes <= 0 || bytes % sizeof(float) != 0) + throw std::invalid_argument("input must be nonempty float32 PCM"); + std::vector pcm(static_cast(bytes) / sizeof(float)); + input.seekg(0); + input.read(reinterpret_cast(pcm.data()), bytes); + if (!input) + throw std::runtime_error("incomplete PCM input"); + trtmc::LoadOptions options; + options.runtime_root = argv[2]; + auto model = trtmc::Model::load(argv[1], options); + const auto task = model.task(); + trtmc::SpeechTranscriptionRequest request{{{pcm.data(), pcm.size()}, + static_cast(rate), + static_cast(channels)}, + {}}; + auto result = task.run(request, {{"max_new_tokens", static_cast(limit)}}); + std::cout << nlohmann::json({{"text", std::string(result.text())}}).dump() << '\n'; + } catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/families/parakeet_tdt/tests/test_build.py b/families/parakeet_tdt/tests/test_build.py new file mode 100644 index 0000000000..1b0596daeb --- /dev/null +++ b/families/parakeet_tdt/tests/test_build.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU checks for dispatch and bundle composition, not TensorRT graph proof.""" + +import importlib +import json +import struct +import sys +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from tensorrt_model_connect.build import BuildRequest +from tensorrt_model_connect.bundle_writer import BundleWriter + + +def request(tmp_path, **changes): + base = BuildRequest( + model_dir=tmp_path / "checkpoint", output_path=tmp_path / "model.bundle", + family="parakeet_tdt", task="speech_transcription", precision="fp16", + ) + return replace(base, **changes) + + +@pytest.mark.parametrize(("option", "value"), [ + ("family", "whisper"), ("task", "transcription"), ("backend", "trt_rtx"), + ("precision", "bf16"), ("tensor_parallel_size", 2), + ("context_parallel_size", 2), ("max_batch_size", 2), + ("max_sequence_length", 128), ("image_height", 224), ("image_width", 224), + ("video_num_frames", 8), ("quantization", "int8"), + ("fp32_layers", (0,)), + ("graph_transform", lambda *args: None), +]) +def test_unsupported_options_fail_before_checkpoint_or_engine_loading(tmp_path, option, value): + build = importlib.import_module("families.parakeet_tdt.model").build + with pytest.raises(ValueError, match=option): + build(request(tmp_path, **{option: value}), BundleWriter(tmp_path / "bad.bundle")) + assert not (tmp_path / "bad.bundle").exists() + + +def test_missing_native_checkpoint_does_not_fall_back_to_nemo(tmp_path): + build = importlib.import_module("families.parakeet_tdt.model").build + req = request(tmp_path) + req.model_dir.mkdir() + (req.model_dir / "legacy.nemo").write_bytes(b"not a supported input") + with pytest.raises(FileNotFoundError, match="config.json"): + build(req, BundleWriter(req.output_path)) + + +def test_dynamic_kv_cache_is_explicitly_unsupported(tmp_path): + build = importlib.import_module("families.parakeet_tdt.model").build + req = request(tmp_path, dynamic_kv_cache=True) + with pytest.raises(NotImplementedError, match="does not support dynamic_kv_cache"): + build(req, BundleWriter(req.output_path)) + + +@pytest.mark.parametrize("precision", ["fp16", "fp32"]) +def test_build_packages_owned_sections_without_mutating_checkpoint(tmp_path, monkeypatch, precision): + build = importlib.import_module("families.parakeet_tdt.model").build + req = request(tmp_path, precision=precision) + req.model_dir.mkdir() + fixture = Path(__file__).with_name("config.json") + (req.model_dir / "config.json").write_bytes(fixture.read_bytes()) + (req.model_dir / "model.safetensors").write_bytes(b"weights validated in separate tests") + (req.model_dir / "tokenizer.json").write_text('{"model":{"type":"Unigram"}}') + before = {p.name: p.read_bytes() for p in req.model_dir.iterdir()} + seen = [] + + def compile_engines(model_dir, *, precision, verbose): + seen.append((model_dir, precision, verbose)) + return {"encoder.plan": b"encoder", "predictor.plan": b"predictor", + "joint.plan": b"joint", "mel_filterbank": b"mel"}, {"tdt_blank_id": 8192} + + monkeypatch.setitem(sys.modules, "families.parakeet_tdt.engines", + SimpleNamespace(compile_engines=compile_engines)) + writer = BundleWriter(req.output_path) + build(req, writer) + writer.finish() + data = req.output_path.read_bytes() + size = struct.unpack_from(" dict[str, tuple[Path, dict, dict]]: + result = {} + for path in sorted(MANIFEST_ROOT.glob("*.json")): + manifest = json.loads(path.read_text(encoding="utf-8")) + assert manifest["family"] == FAMILY + assert manifest["task"] in TASKS + for case in manifest["testcases"]: + name = str(case["name"]) + assert name not in result + result[name] = (path, manifest, case) + return result + + +CASES = _case_index() + + +def _selected_cases(config) -> tuple[list[str], bool]: + model_filters = set() + for raw in config.getoption("--e2e-model") or []: + model_filters.update((item.strip() for item in str(raw).split(",") if item.strip())) + models_file = config.getoption("--e2e-models-file") + if models_file: + model_filters.update( + ( + line.strip() + for line in Path(models_file).read_text(encoding="utf-8").splitlines() + if line.strip() and (not line.lstrip().startswith("#")) + ) + ) + testcase_filters = set() + for raw in config.getoption("--e2e-testcase") or []: + testcase_filters.update((item.strip() for item in str(raw).split(",") if item.strip())) + if not model_filters and (not testcase_filters): + return (sorted(CASES), False) + selected = [] + for name, (_, manifest, _) in CASES.items(): + model_match = ( + not model_filters + or FAMILY in model_filters + or name in model_filters + or (manifest["name"] in model_filters) + ) + testcase_match = not testcase_filters or name in testcase_filters + if model_match and testcase_match: + selected.append(name) + return (sorted(selected), True) + + +def pytest_generate_tests(metafunc) -> None: + if "case_name" in metafunc.fixturenames: + names, enabled = _selected_cases(metafunc.config) + parameters = names + if not enabled: + parameters = [ + pytest.param( + name, + marks=pytest.mark.skip( + reason="direct E2E requires one of the three explicit E2E selectors" + ), + ) + for name in names + ] + metafunc.parametrize("case_name", parameters, ids=names) + + +def _required_path(value: str | None, label: str) -> Path: + assert value, f"selected {FAMILY} E2E requires {label}" + path = Path(value) + assert path.exists(), f"selected {FAMILY} E2E {label} does not exist: {path}" + return path + + +def _model_dir(manifest: dict) -> Path: + explicit = os.environ.get(f"TRTMC_{FAMILY.upper()}_MODEL_DIR") + if explicit: + return _required_path(explicit, f"TRTMC_{FAMILY.upper()}_MODEL_DIR") + from huggingface_hub import snapshot_download + + try: + snapshot = snapshot_download( + repo_id=manifest["hf_id"], revision=manifest["hf_revision"], local_files_only=True + ) + except Exception as error: + raise AssertionError( + f"selected {FAMILY} E2E requires the exact cached checkpoint {manifest['hf_id']}" + ) from error + return Path(snapshot) + + +def _audio(case, generated_root): + path = TEST_ROOT / case["test_input_audio"] + if not path.is_file(): + from families.parakeet_tdt.tests.data.asr_probes.generate_asr_probe_inputs import main + path = generated_root / path.name + if not path.is_file(): + main(generated_root) + with wave.open(str(path), "rb") as source: + assert source.getsampwidth() == 2, "Parakeet fixtures must use PCM16" + rate, channels = source.getframerate(), source.getnchannels() + audio = np.frombuffer(source.readframes(source.getnframes()), dtype=" Path: + native = _required_path(os.environ.get("TRTMC_NATIVE_BUILD_DIR"), "TRTMC_NATIVE_BUILD_DIR") + subprocess.run( + ["cmake", "--build", str(native), "--parallel", "8", "--target", "test_parakeet_tdt_sdk_cpp"], + check=True, + timeout=600, + ) + consumer = native / "test_parakeet_tdt_sdk_cpp" + assert consumer.is_file(), "native build did not produce the family-owned SDK consumer" + return consumer + + +def test_official_checkpoint_e2e(case_name, tmp_path): + _, manifest, case = CASES[case_name] + runtime = _required_path(os.environ.get("TRTMC_RUNTIME_ROOT"), "TRTMC_RUNTIME_ROOT") + consumer = _sdk_consumer_binary() + assert (runtime / "libtrtmc_model_parakeet_tdt.so").is_file() + assert (runtime / "libtrtmc_backend_trt.so").is_file() + import torch + assert torch.cuda.is_available(), "selected Parakeet E2E requires CUDA" + model_dir = _model_dir(manifest) + record_evidence("inputs", {"manifest": manifest, "case": case}) + record_evidence("checkpoint", {"model_dir": str(model_dir), "hf_revision": manifest["hf_revision"]}) + bundle = tmp_path / manifest["bundle"] + with evidence_stage("build"): + build(BuildRequest(model_dir=model_dir, output_path=bundle, family=FAMILY, + task=manifest["task"], precision=manifest["precision"], + tensor_parallel_size=int(manifest["tensor_parallel_size"]))) + audio, rate, channels = _audio(case, tmp_path / "audio-probes") + pcm = tmp_path / "input.f32" + audio.astype(" 0 and rate > 0 and channels > 0 + assert list(source_root.iterdir()) == [] + assert len(list(output.glob("*.wav"))) == 7 diff --git a/families/parakeet_tdt/tests/test_factory.py b/families/parakeet_tdt/tests/test_factory.py new file mode 100644 index 0000000000..e48af05cb3 --- /dev/null +++ b/families/parakeet_tdt/tests/test_factory.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Real Python bundle writer -> C++ reader/factory, with fake engine loading.""" + +import json +import os +from pathlib import Path +import shutil +import struct +import subprocess + +import pytest + +from tensorrt_model_connect.bundle_writer import BundleWriter + + +@pytest.fixture(scope="module") +def factory_runner(tmp_path_factory): + compiler = shutil.which(os.environ.get("CXX", "c++")) + cuda = Path(os.environ.get("CUDA_HOME", os.environ.get("CUDA_PATH", "/usr/local/cuda"))) + json_include = Path(os.environ.get("NLOHMANN_JSON_INCLUDE_DIR", "/usr/include")) + if not compiler or not (cuda / "include/cuda_runtime_api.h").is_file(): + pytest.skip("factory contract requires C++17 and CUDA headers; no GPU is used") + if not (json_include / "nlohmann/json.hpp").is_file(): + pytest.skip("factory contract requires nlohmann/json.hpp (set NLOHMANN_JSON_INCLUDE_DIR)") + family = Path(__file__).resolve().parents[1] + root = family.parents[1] + output = tmp_path_factory.mktemp("factory") / "factory" + subprocess.run([ + compiler, "-std=c++17", "-Wall", "-Wextra", "-Werror", + "-I", str(root), "-I", str(root / "core"), + "-I", str(root / "core/runtime/include"), + "-I", str(cuda / "include"), "-I", str(json_include), + str(family / "tests/cpp/test_factory.cpp"), + *(str(family / "runtime" / name) for name in ( + "plugin.cpp", "pipeline.cpp", "bpe_tokenizer.cpp", "audio_helpers.cpp", "resampler.cpp", + )), + str(root / "core/runtime/bundle/bundle_format.cpp"), "-o", str(output), + ], check=True, capture_output=True, text=True, timeout=120) + return output + + +@pytest.mark.parametrize(("change", "error", "calls"), [ + ("valid", "", 3), + ("wrong-family", "identity", 0), + ("wrong-task", "identity", 0), + ("wrong-backend", "identity", 0), + ("kv", "kv-cache-size", 0), + ("dimension", "dimension", 0), + ("fractional-dimension", "dimension", 0), + ("policy", "duration policy", 0), + ("mel-length", "filterbank size", 0), + ("mel-nonfinite", "nonfinite", 0), + ("tokenizer-type", "native BPE tokenizer", 0), + ("tokenizer-decoder", "native BPE tokenizer", 0), + ("tokenizer-replacement", "native BPE tokenizer", 0), + ("tokenizer-prepend", "native BPE tokenizer", 0), + ("missing-runtime", "missing section", 0), + ("missing-tokenizer", "missing section", 0), + ("missing-encoder", "missing section", 0), + ("missing-predictor", "missing section", 1), + ("missing-joint", "missing section", 2), + ("fail-engine", "failed to load encoder.plan", 1), + ("fail-predictor", "failed to load predictor.plan", 2), + ("fail-joint", "failed to load joint.plan", 3), +]) +def test_native_factory_bundle_contract(tmp_path, factory_runner, change, error, calls): + runtime = { + "tensor_parallel_size": 1, "mel_sampling_rate": 16000, "num_mel_bins": 128, + "mel_n_fft": 512, "mel_win_length": 400, "mel_hop_length": 160, + "mel_chunk_length": 30, "mel_length": 3000, "tdt_encoder_hidden_size": 1024, + "tdt_pred_hidden_size": 640, "tdt_pred_num_layers": 2, "tdt_vocab_size": 8192, + "tdt_blank_id": 8192, "tdt_encoder_layers": 24, "max_source_positions": 375, + "subsampling_factor": 8, "tdt_max_symbols_per_step": 10, + "tdt_att_context_left": -1, "tdt_att_context_right": -1, + "tdt_duration_values": [0, 1, 2, 3, 4], "tdt_causal_downsampling": False, + "mel_normalize": "per_feature", "mel_preemph": 0.97, + } + tokenizer = { + "model": {"type": "BPE", "vocab": {"▁hello": 0}, "merges": []}, + "decoder": {"type": "Metaspace", "replacement": "▁", "prepend_scheme": "always"}, + } + if change == "dimension": + runtime["tdt_blank_id"] = 1024 + if change == "fractional-dimension": + runtime["tdt_blank_id"] = 8192.0 + if change == "policy": + runtime["tdt_duration_values"] = [0, 2] + if change == "tokenizer-type": + tokenizer["model"]["type"] = "Unigram" + if change == "tokenizer-decoder": + tokenizer["decoder"]["type"] = "ByteLevel" + if change == "tokenizer-replacement": + tokenizer["decoder"]["replacement"] = "_" + if change == "tokenizer-prepend": + tokenizer["decoder"]["prepend_scheme"] = "never" + mel = struct.pack("- + The family-local migration has no release-performance workload or receipt + yet. Target-GPU engine conversion and exact-transcript E2E parity remain + unverified; no performance qualification is claimed. - model: qwen3-embedding-0.6b reason: >- The migrated embedding task has no release-performance workload or receipt