From df41fc00fc6819416cd8ec4d6bcc850706da52cd Mon Sep 17 00:00:00 2001 From: Karen Mosoyan Date: Tue, 16 Jun 2026 17:27:18 -0700 Subject: [PATCH 1/6] Wire Parakeet-TDT cloud-handoff probe (mirrors Gemma4) Adds probe-driven cloud handoff for Parakeet ASR, mirroring the existing Gemma4 text path. The ~65k-param probe reads conformer layer-17 activations and predicts p_wrong (WER>0.30); confidence = 1 - p_wrong drives handoff. Transpile (tdt_runtime.py): ParakeetTDTEncoderProbe wraps the encoder so the audio_encoder CPU graph also emits probe_hidden (layer-17 norm_out). The NPU/ANE encoder is unchanged, so probe handoff is CPU-graph-only. Export (handoff_probe.py): export_parakeet_handoff_probe reads real dims from the checkpoint tensors (1024/32/128/64) rather than the hardcoded Gemma 1536, via a shared _write_probe_bundle; wired into the convert CLI. Engine (model.cpp): the load guard checks audio_encoder for probe_hidden when the model is PARAKEET_TDT; maybe_capture_handoff_probe_hidden now appends all frame rows (text decode emits a single row, so prior behaviour is preserved); transcribe_parakeet_tdt resets and captures the rollout. Decision (transcribe.cpp): probe-driven cloud handoff with threshold, cloud_transcribe_request, cloud_handoff_reason, and graceful local fallback. tests/transcribe.cpp: env knobs (CACTUS_TEST_AUTO_HANDOFF / _CONFIDENCE_THRESHOLD / _CLOUD_TIMEOUT_MS) to exercise the handoff path; default behaviour unchanged. Validated end-to-end: engine p_wrong=0.074 (confidence 0.926) on test.wav vs the reference kit's 0.060; "above threshold" keeps local, forced low confidence triggers handoff with graceful fallback. --- cactus-engine/src/model.cpp | 40 ++++-- cactus-engine/src/transcribe.cpp | 77 +++++++++++- cactus-engine/tests/transcribe.cpp | 13 +- python/cactus/cli/model.py | 9 +- python/cactus/convert/handoff_probe.py | 163 +++++++++++++++++-------- python/cactus/transpile/tdt_runtime.py | 50 +++++++- 6 files changed, 288 insertions(+), 64 deletions(-) diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index a955b871f..45a72990d 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -352,12 +352,16 @@ void Model::maybe_capture_handoff_probe_hidden(const Component& comp) { if (desc.total_size < handoff_probe_feat_dim_) return; size_t rows = desc.total_size / handoff_probe_feat_dim_; if (rows == 0) return; - size_t row = rows - 1; const auto* data = static_cast(comp.graph->get_output(node)); if (!data) return; - size_t base = row * static_cast(handoff_probe_feat_dim_); - for (size_t i = 0; i < handoff_probe_feat_dim_; ++i) { - handoff_probe_hidden_.push_back(read_scalar_value(desc.precision, data, base + i)); + // Append every row of probe_hidden. Text decode steps emit a single row + // (one token), so this preserves prior behaviour; the Parakeet audio encoder + // emits all T time frames in one pass, which the probe then pools over. + for (size_t row = 0; row < rows; ++row) { + size_t base = row * static_cast(handoff_probe_feat_dim_); + for (size_t i = 0; i < handoff_probe_feat_dim_; ++i) { + handoff_probe_hidden_.push_back(read_scalar_value(desc.precision, data, base + i)); + } } } @@ -636,10 +640,24 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, CACTUS_LOG_WARN("model", "NPU source encoder load failed for " << full_path << "; falling back to CPU"); } } - if (load_handoff_probe() && decoder_ && output_index(*decoder_, "probe_hidden") < 0) { - CACTUS_LOG_WARN("cloud_handoff", "Handoff probe is packaged, but decoder_step does not expose probe_hidden; " - "reconvert Gemma4 to enable probe-based handoff"); - handoff_probe_loaded_ = false; + if (load_handoff_probe()) { + // The probe reads a "probe_hidden" graph output. For Parakeet ASR it is + // exposed by the audio_encoder (layer-17 activations); for text models + // by the decoder_step. Disable the probe if the bundle predates it. + bool exposes_probe_hidden = false; + const char* reconvert_hint = + "decoder_step does not expose probe_hidden; reconvert Gemma4 to enable probe-based handoff"; + if (config_.model_type == Config::ModelType::PARAKEET_TDT) { + exposes_probe_hidden = audio_encoder_ && output_index(*audio_encoder_, "probe_hidden") >= 0; + reconvert_hint = + "audio_encoder does not expose probe_hidden; reconvert Parakeet to enable probe-based handoff"; + } else { + exposes_probe_hidden = decoder_ && output_index(*decoder_, "probe_hidden") >= 0; + } + if (!exposes_probe_hidden) { + CACTUS_LOG_WARN("cloud_handoff", "Handoff probe is packaged, but " << reconvert_hint); + handoff_probe_loaded_ = false; + } } initialized_ = true; @@ -2883,6 +2901,8 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a std::vector emitted; double raw_decode_ms = 0.0; + reset_handoff_probe_rollout(); + Component* audio_enc = components_.count("audio_encoder") ? &components_.at("audio_encoder") : nullptr; Component* dec = components_.count("decoder") ? &components_.at("decoder") : nullptr; if (!audio_enc || !dec) { @@ -2967,6 +2987,10 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a } if (!used_npu) { audio_enc->graph->execute(); + // Capture layer-17 activations for the cloud-handoff probe. Only the CPU + // graph exposes probe_hidden; the NPU encoder emits only "encoded", so + // probe-based handoff is unavailable on the NPU path. + maybe_capture_handoff_probe_hidden(*audio_enc); } int hidden_idx = output_index(*audio_enc, "encoder_hidden_states"); diff --git a/cactus-engine/src/transcribe.cpp b/cactus-engine/src/transcribe.cpp index 20f9c2606..7d8ad5643 100644 --- a/cactus-engine/src/transcribe.cpp +++ b/cactus-engine/src/transcribe.cpp @@ -1,10 +1,12 @@ #include "../cactus_engine.h" #include "utils.h" +#include "cloud.h" #include "cactus_kernels.h" #include "wav.h" #include #include #include +#include #include #include @@ -237,9 +239,11 @@ int cactus_transcribe( const size_t mel_bins = std::max(1, static_cast(handle->model->get_config().num_mel_bins)); std::vector audio_features; + std::vector handoff_audio_samples; // 16k mono, pre-emphasis-free, for cloud handoff if (is_parakeet) { auto cfg = cactus::audio::get_parakeet_spectrogram_config(); const size_t waveform_samples = audio_samples.size(); + handoff_audio_samples = audio_samples; cactus::audio::apply_preemphasis(audio_samples, 0.97f); audio_features = cactus::audio::compute_spectrogram_graph( audio_samples, cfg, mel_bins, 0.0f, 8000.0f, @@ -423,6 +427,76 @@ int cactus_transcribe( : 0.0f; float confidence = 1.0f - mean_entropy; + // Cloud-handoff decision for Parakeet, driven by the handoff probe's + // p_wrong (confidence = 1 - p_wrong). Mirrors the Gemma4 path in + // complete.cpp. TDT decoding produces no token entropy, so without the + // probe confidence stays 1.0 and no handoff occurs. + bool cloud_handoff_used = false; + std::string handoff_reason; + float reported_threshold = -1.0f; + if (is_parakeet) { + const bool cloud_disabled = env_flag_enabled("CACTUS_DISABLE_CLOUD_HANDOFF"); + const bool cloud_eligible = + !cloud_disabled && !handle->cloud_handoff_disabled && options.auto_handoff; + float threshold = options.confidence_threshold; + if (threshold < 0.0f) { + const float model_default = handle->model->get_config().default_cloud_handoff_threshold; + threshold = handle->model->has_handoff_probe() + ? 0.50f + : (model_default > 0.0f ? model_default : 0.7f); + } + reported_threshold = threshold; + + if (!options.auto_handoff) { + handoff_reason = "handoff off"; + } else if (cloud_disabled) { + handoff_reason = "disabled (env)"; + } else if (handle->cloud_handoff_disabled) { + handoff_reason = "disabled (auth failed earlier)"; + } + + if (cloud_eligible && handle->model->has_handoff_probe_rollout()) { + float p_wrong = handle->model->handoff_probe_wrong_probability(); + if (std::isfinite(p_wrong)) { + confidence = std::max(0.0f, std::min(1.0f, 1.0f - p_wrong)); + CACTUS_LOG_DEBUG("cloud_handoff", "Parakeet handoff probe p_wrong=" + << p_wrong << " confidence=" << confidence); + if (confidence < threshold && !handoff_audio_samples.empty()) { + CACTUS_LOG_WARN("cloud_handoff", "Cloud transcription handoff triggered: p_wrong=" + << p_wrong << " confidence=" << confidence << " threshold=" << threshold); + std::vector pcm16(handoff_audio_samples.size()); + for (size_t i = 0; i < handoff_audio_samples.size(); ++i) { + float clamped = std::max(-1.0f, std::min(1.0f, handoff_audio_samples[i])); + pcm16[i] = static_cast(std::lround(clamped * 32767.0f)); + } + std::vector wav = cactus::ffi::cloud_build_wav( + reinterpret_cast(pcm16.data()), pcm16.size() * sizeof(int16_t)); + std::string audio_b64 = cactus::ffi::cloud_base64_encode(wav.data(), wav.size()); + std::string cloud_key = cactus::ffi::resolve_cloud_api_key(nullptr); + CloudResponse cloud = cactus::ffi::cloud_transcribe_request( + audio_b64, final_text, + static_cast(options.cloud_timeout_ms / 1000), + cloud_key.empty() ? nullptr : cloud_key.c_str()); + if (cloud.used_cloud && !cloud.transcript.empty()) { + final_text = cloud.transcript; + cloud_handoff_used = true; + handoff_reason = "low confidence (probe)"; + } else { + if (cloud.error.rfind("http_401", 0) == 0 || cloud.error.rfind("http_403", 0) == 0) { + handle->cloud_handoff_disabled = true; + } + handoff_reason = "handoff failed: " + cloud.error; + CACTUS_LOG_WARN("cloud_handoff", "Cloud transcription failed, keeping local output: " + << cloud.error); + } + } + } + } + if (handoff_reason.empty()) { + handoff_reason = (confidence >= threshold) ? "above threshold" : "kept local"; + } + } + auto end_time = std::chrono::high_resolution_clock::now(); double total_time_ms = std::chrono::duration_cast(end_time - start_time).count() / 1000.0; @@ -436,7 +510,8 @@ int cactus_transcribe( std::string json = construct_response_json( final_text, {}, time_to_first_token, total_time_ms, prefill_tps, decode_tps, prompt_tokens, completion_tokens, - confidence, false, "", segments); + confidence, cloud_handoff_used, "", segments, "", + reported_threshold, handoff_reason); if (json.size() >= buffer_size) { handle_error_response("Response buffer too small", response_buffer, buffer_size); return -1; diff --git a/cactus-engine/tests/transcribe.cpp b/cactus-engine/tests/transcribe.cpp index 84d22b0e9..084207df3 100644 --- a/cactus-engine/tests/transcribe.cpp +++ b/cactus-engine/tests/transcribe.cpp @@ -104,7 +104,18 @@ static int transcribe_file(cactus_model_t model, const std::string& audio_path, const std::string& model_path, const std::string& language) { std::string prompt = whisper_prompt(model_path, language); std::vector response_buffer(RESPONSE_BUFFER_SIZE, 0); - std::string options = "{\"max_tokens\":500,\"auto_handoff\":false,\"language\":\"" + language + "\"}"; + // Cloud-handoff knobs are env-overridable so this binary can exercise the + // Parakeet handoff-probe path (CACTUS_TEST_AUTO_HANDOFF=1, optional + // CACTUS_TEST_CONFIDENCE_THRESHOLD / CACTUS_TEST_CLOUD_TIMEOUT_MS). + const bool auto_handoff = std::getenv("CACTUS_TEST_AUTO_HANDOFF") != nullptr; + const char* thr = std::getenv("CACTUS_TEST_CONFIDENCE_THRESHOLD"); + const char* cto = std::getenv("CACTUS_TEST_CLOUD_TIMEOUT_MS"); + std::ostringstream opts; + opts << "{\"max_tokens\":500,\"language\":\"" << language << "\"" + << ",\"auto_handoff\":" << (auto_handoff ? "true" : "false") + << ",\"confidence_threshold\":" << (thr ? thr : "-1") + << ",\"cloud_timeout_ms\":" << (cto ? cto : "15000") << "}"; + std::string options = opts.str(); auto start = std::chrono::steady_clock::now(); int rc = cactus_transcribe( diff --git a/python/cactus/cli/model.py b/python/cactus/cli/model.py index afc2baac5..1f01c49ac 100644 --- a/python/cactus/cli/model.py +++ b/python/cactus/cli/model.py @@ -336,12 +336,17 @@ def ensure_bundle(model_id, *, bits=4, token=None, raise RuntimeError(f"Transpilation failed for {model_id}") try: - from cactus.convert.handoff_probe import export_gemma4_handoff_probe + from cactus.convert.handoff_probe import ( + export_gemma4_handoff_probe, + export_parakeet_handoff_probe, + ) if export_gemma4_handoff_probe(output_dir, model_id=model_id): print_color(GREEN, f"Gemma4 cloud handoff probe packaged into {output_dir}") + if export_parakeet_handoff_probe(output_dir, model_id=model_id): + print_color(GREEN, f"Parakeet cloud handoff probe packaged into {output_dir}") except Exception as e: - print_color(YELLOW, f"Warning: failed to package Gemma4 cloud handoff probe: {e}") + print_color(YELLOW, f"Warning: failed to package cloud handoff probe: {e}") print_color(GREEN, f"Model converted and transpiled to {output_dir}") return output_dir diff --git a/python/cactus/convert/handoff_probe.py b/python/cactus/convert/handoff_probe.py index 8fc45058e..3d82472a5 100644 --- a/python/cactus/convert/handoff_probe.py +++ b/python/cactus/convert/handoff_probe.py @@ -3,6 +3,7 @@ import io import hashlib import json +import os import struct import zipfile from pathlib import Path @@ -25,61 +26,57 @@ ) -def _candidate_probe_files(output_dir: Path) -> list[Path]: +def _candidate_probe_files(output_dir: Path, names: tuple[str, ...]) -> list[Path]: cwd = Path.cwd() - return [ - output_dir / "probe.pt", - output_dir / "global_attn_probe_v10p6.pt", - cwd / "probe.pt", - cwd / "v10p6_probe_release" / "global_attn_probe_v10p6.pt", - Path.home() / "Downloads" / "probe.pt", - Path.home() / "Downloads" / "v10p6_probe_release" / "global_attn_probe_v10p6.pt", - ] - - -def _candidate_probe_zips(output_dir: Path) -> list[Path]: + roots = [output_dir, cwd, Path.home() / "Downloads"] + out: list[Path] = [] + for root in roots: + for name in names: + out.append(root / name) + # also look one release-dir deep (e.g. v10p6_probe_release/, the kit dir) + out.append(root / "v10p6_probe_release" / "global_attn_probe_v10p6.pt") + out.append(root / "parakeet_handoff_kit" / "parakeet_handoff_probe.pt") + out.append(root / "parakeet_handoff_kit" / "parakeet_handoff_kit" / "parakeet_handoff_probe.pt") + return out + + +def _candidate_probe_zips(output_dir: Path, zip_names: tuple[str, ...]) -> list[Path]: cwd = Path.cwd() - return [ - output_dir / "v10p6_probe_release.zip", - cwd / "v10p6_probe_release.zip", - Path.home() / "Downloads" / "v10p6_probe_release.zip", - ] + return [root / name for root in (output_dir, cwd, Path.home() / "Downloads") for name in zip_names] -def _load_checkpoint_from_zip(zip_path: Path) -> Any: +def _load_checkpoint_from_zip(zip_path: Path, inner_names: tuple[str, ...]) -> Any: import torch with zipfile.ZipFile(zip_path) as zf: - names = set(zf.namelist()) - for name in ( - "v10p6_probe_release/global_attn_probe_v10p6.pt", - "global_attn_probe_v10p6.pt", - "probe.pt", - ): - if name in names: + available = set(zf.namelist()) + for name in inner_names: + if name in available: with zf.open(name) as f: - return torch.load(io.BytesIO(f.read()), map_location="cpu") + return torch.load(io.BytesIO(f.read()), map_location="cpu", weights_only=False) raise FileNotFoundError(f"no probe checkpoint found in {zip_path}") def _state_dict_from_checkpoint(checkpoint: Any) -> dict[str, Any]: if isinstance(checkpoint, dict): - for key in ("state_dict", "model_state"): + for key in ("model_state", "state_dict"): value = checkpoint.get(key) if isinstance(value, dict): return value return checkpoint -def _load_checkpoint(output_dir: Path) -> tuple[Any, str] | tuple[None, None]: +def _load_checkpoint( + output_dir: Path, file_names: tuple[str, ...], zip_names: tuple[str, ...], inner_names: tuple[str, ...] +) -> tuple[Any, str] | tuple[None, None]: import torch - for path in _candidate_probe_files(output_dir): + for path in _candidate_probe_files(output_dir, file_names): if path.exists(): - return torch.load(path, map_location="cpu"), str(path) - for path in _candidate_probe_zips(output_dir): + return torch.load(path, map_location="cpu", weights_only=False), str(path) + for path in _candidate_probe_zips(output_dir, zip_names): if path.exists(): - return _load_checkpoint_from_zip(path), str(path) + return _load_checkpoint_from_zip(path, inner_names), str(path) return None, None @@ -93,39 +90,105 @@ def _tensor_hash(state: dict[str, Any]) -> str: return digest.hexdigest() -def export_gemma4_handoff_probe(output_dir: str | Path, *, model_id: str | None = None) -> bool: - """Package the Gemma4 v10p6 cloud-handoff probe into a C++-readable bundle file.""" - model_key = (model_id or "").lower() - if model_key and "gemma-4" not in model_key and "gemma4" not in model_key: - return False - - out_dir = Path(output_dir) - checkpoint, source = _load_checkpoint(out_dir) - if checkpoint is None: - return False - - state = _state_dict_from_checkpoint(checkpoint) +def _probe_dims(state: dict[str, Any]) -> tuple[int, int, int, int]: + """Read (feat_dim, t_h, h1, h2) from tensor shapes and validate the + architecture the C++ loader (Model::load_handoff_probe) assumes.""" missing = [key for key in _ORDERED_KEYS if key not in state] if missing: raise RuntimeError(f"handoff probe checkpoint missing tensors: {', '.join(missing)}") + feat_dim = int(state["norm.weight"].shape[0]) + t_h = int(state["proj.weight"].shape[0]) + h1 = int(state["head.0.weight"].shape[0]) + h2 = int(state["head.2.weight"].shape[0]) + expect = { + "norm.weight": (feat_dim,), "norm.bias": (feat_dim,), + "proj.weight": (t_h, feat_dim), "proj.bias": (t_h,), + "attn_query": (t_h,), + "head.0.weight": (h1, t_h), "head.0.bias": (h1,), + "head.2.weight": (h2, h1), "head.2.bias": (h2,), + "head.4.weight": (1, h2), "head.4.bias": (1,), + } + for key, shape in expect.items(): + got = tuple(int(d) for d in state[key].shape) + if got != shape: + raise RuntimeError(f"handoff probe tensor {key}: expected {shape}, got {got}") + return feat_dim, t_h, h1, h2 + +def _write_probe_bundle(state: dict[str, Any], out_dir: Path, source: str, *, layer: int, fmt: str) -> dict: + feat_dim, t_h, h1, h2 = _probe_dims(state) out_dir.mkdir(parents=True, exist_ok=True) probe_path = out_dir / "handoff_probe.bin" with probe_path.open("wb") as f: f.write(_PROBE_MAGIC) - f.write(struct.pack(" bool: + """Package the Gemma4 v10p6 cloud-handoff probe into a C++-readable bundle file.""" + model_key = (model_id or "").lower() + if model_key and "gemma-4" not in model_key and "gemma4" not in model_key: + return False + + out_dir = Path(output_dir) + checkpoint, source = _load_checkpoint( + out_dir, + file_names=("probe.pt", "global_attn_probe_v10p6.pt"), + zip_names=("v10p6_probe_release.zip",), + inner_names=( + "v10p6_probe_release/global_attn_probe_v10p6.pt", + "global_attn_probe_v10p6.pt", + "probe.pt", + ), + ) + if checkpoint is None: + return False + state = _state_dict_from_checkpoint(checkpoint) + _write_probe_bundle(state, out_dir, source, layer=28, fmt="cactus_handoff_probe_v10p6") + return True + + +def export_parakeet_handoff_probe(output_dir: str | Path, *, model_id: str | None = None) -> bool: + """Package the Parakeet-TDT cloud-handoff probe into a C++-readable bundle. + + Reads the probe dims from the checkpoint tensors (feat_dim=1024, t_h=32, + h1=128, h2=64 for parakeet-tdt-0.6b-v2), so it is not tied to the Gemma4 + 1536-d shape. The capture layer is informational metadata; the engine taps + whatever layer the transpiler exposes as probe_hidden.""" + model_key = (model_id or "").lower() + if model_key and "parakeet" not in model_key: + return False + + out_dir = Path(output_dir) + checkpoint, source = _load_checkpoint( + out_dir, + file_names=("parakeet_handoff_probe.pt", "probe.pt"), + zip_names=("parakeet_handoff_kit.zip",), + inner_names=( + "parakeet_handoff_kit/parakeet_handoff_probe.pt", + "parakeet_handoff_kit/parakeet_handoff_kit/parakeet_handoff_probe.pt", + "parakeet_handoff_probe.pt", + ), + ) + if checkpoint is None: + return False + state = _state_dict_from_checkpoint(checkpoint) + layer = int(os.environ.get("CACTUS_PARAKEET_PROBE_LAYER", "17")) + _write_probe_bundle(state, out_dir, source, layer=layer, fmt="cactus_handoff_probe_parakeet") return True diff --git a/python/cactus/transpile/tdt_runtime.py b/python/cactus/transpile/tdt_runtime.py index 01275d7ec..ee97a7ab3 100644 --- a/python/cactus/transpile/tdt_runtime.py +++ b/python/cactus/transpile/tdt_runtime.py @@ -658,6 +658,46 @@ def forward(self, input_features: torch.Tensor) -> torch.Tensor: return x +# Encoder layer (0-indexed) whose output the cloud-handoff probe was trained on. +# parakeet-tdt-0.6b-v2 has 24 conformer layers; the probe taps layer 17 (NOT the +# final layer). Override with CACTUS_PARAKEET_PROBE_LAYER if a probe is retrained. +PARAKEET_PROBE_CAPTURE_LAYER = int(os.environ.get("CACTUS_PARAKEET_PROBE_LAYER", "17")) + + +class ParakeetTDTEncoderProbe(nn.Module): + """Wraps the Parakeet encoder so it returns (encoder_hidden_states, + probe_hidden), where probe_hidden is the raw output of conformer layer + ``capture_index``. Mirrors the Gemma4 ``capture_layer_index`` mechanism so + the C++ engine can run the cloud-handoff probe on ASR activations. + + Reuses the encoder's own ``pre_encode`` and ``layers`` submodules, so traced + parameter names match ``model.encoder`` exactly and weight mapping is + unchanged.""" + + def __init__(self, encoder: "ParakeetTDTEncoder", capture_index: int): + super().__init__() + self.pre_encode = encoder.pre_encode + self.layers = encoder.layers + num_layers = len(self.layers) + if not (0 <= capture_index < num_layers): + clamped = max(0, num_layers - 1) + print( + f"tdt_runtime: probe capture layer {capture_index} out of range for " + f"{num_layers}-layer encoder; clamping to {clamped}" + ) + capture_index = clamped + self.capture_index = capture_index + + def forward(self, input_features: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + x = self.pre_encode(input_features) + probe_hidden = x + for index, layer in enumerate(self.layers): + x = layer(x) + if index == self.capture_index: + probe_hidden = x + return x, probe_hidden + + class ParakeetTDTANESelfAttention(nn.Module): def __init__(self, attn: "ParakeetTDTSelfAttention", seq_len: int): super().__init__() @@ -1003,13 +1043,19 @@ def build_parakeet_tdt_component_specs( ane_seq_len = int(model.encoder.pre_encode(npu_input_features).shape[1]) ane_encoder = ParakeetTDTANEEncoder(model.encoder, ane_seq_len) + # CPU encoder additionally emits probe_hidden (layer-17 activations) for the + # cloud-handoff probe. The NPU/ANE encoder keeps its single "encoded" output + # (the offline NPU emit traces npu_module directly and ignores output_keys), + # so probe-based handoff is only available on the CPU graph path. + probe_encoder = ParakeetTDTEncoderProbe(model.encoder, PARAKEET_PROBE_CAPTURE_LAYER) + return [ ComponentModuleSpec( component="audio_encoder", - module=model.encoder, + module=probe_encoder, example_inputs=(input_features,), input_keys=("input_features",), - output_keys=("encoder_hidden_states",), + output_keys=("encoder_hidden_states", "probe_hidden"), graph_meta={**common_graph_meta, "component": "audio_encoder"}, metadata={"family": "parakeet_tdt", "task": "tdt_transcription"}, npu_module=ane_encoder, From fc1c6fdadca58612e562a2d8b5853c2daeef5e79 Mon Sep 17 00:00:00 2001 From: Karen Mosoyan Date: Tue, 16 Jun 2026 18:11:16 -0700 Subject: [PATCH 2/6] Parakeet handoff probe on the NPU/ANE bundle + simpler probe detection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The probe previously rode only the CPU encoder graph; the ANE encoder emitted just "encoded", so probe-based handoff was unavailable on NPU bundles. Make the ANE path carry the probe too. Transpile (tdt_runtime.py): ParakeetTDTANEEncoder now returns (encoded, probe_hidden), capturing conformer layer 17 like the CPU wrapper. npu/audio.py: emit_audio_encoder_mlpackage detects the 2-output encoder and names the extra CoreML output "probe_hidden" (single-output encoders unchanged). Engine: NPUEncoder gains encode_with_secondary (default no-op); ANEEncoder implements it, reading both "encoded" and "probe_hidden" from one ANE prediction. transcribe_parakeet_tdt accumulates probe_hidden across NPU chunks into the handoff rollout, with graceful fallback to plain encode (and a warning) when the mlpackage predates the probe output. Probe detection (handoff_probe.py): detect parakeet_handoff_probe.pt the same way as Gemma4 — a plain file in the bundle output dir or the convert working directory — instead of hunting the kit folder/zip. Validated: the real 24-layer ParakeetTDTANEEncoder exports to an int8 CoreML package with outputs encoded[1,375,1024] and probe_hidden[1,375,1024]; engine builds clean (ANE path included). On-ANE-device runtime read not exercised here. --- cactus-engine/src/engine.h | 14 ++++ cactus-engine/src/model.cpp | 40 ++++++++++- cactus-engine/src/npu_ane.h | 8 +++ cactus-engine/src/npu_ane.mm | 91 ++++++++++++++++++++++++++ python/cactus/convert/handoff_probe.py | 35 +++++----- python/cactus/transpile/npu/audio.py | 15 ++++- python/cactus/transpile/tdt_runtime.py | 18 +++-- 7 files changed, 193 insertions(+), 28 deletions(-) diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index 604f730f5..4a06951ff 100644 --- a/cactus-engine/src/engine.h +++ b/cactus-engine/src/engine.h @@ -39,6 +39,20 @@ class NPUEncoder { const std::vector& shape, const std::string& input_name = "x", const std::string& output_name = "") = 0; + // Run one prediction and copy two named outputs (e.g. "encoded" plus the + // cloud-handoff probe's "probe_hidden"). Returns elements written to + // primary_output, or 0 if the model does not expose secondary_name (caller + // should fall back to encode()). Default: unsupported. + virtual size_t encode_with_secondary( + const __fp16* input, __fp16* primary_output, __fp16* secondary_output, + const std::vector& shape, + const std::string& input_name = "x", + const std::string& primary_name = "encoded", + const std::string& secondary_name = "probe_hidden") { + (void)input; (void)primary_output; (void)secondary_output; (void)shape; + (void)input_name; (void)primary_name; (void)secondary_name; + return 0; + } virtual bool is_available() const = 0; virtual std::vector get_input_shape() const = 0; virtual std::vector get_output_shape() const = 0; diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 45a72990d..c2d2ae3ee 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -2955,6 +2955,14 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a const size_t num_chunks = (copy_frames + window_frames - 1) / window_frames; const size_t total_hidden_T = num_chunks * window_hidden; npu_hidden_storage.assign(total_hidden_T * hidden_dim_npu, __fp16(0)); + // If a handoff probe is packaged and the ANE encoder exposes a + // probe_hidden output, capture layer-17 activations alongside the + // encoded output (one ANE run produces both). + const bool want_probe = handoff_probe_loaded_ && + static_cast(handoff_probe_feat_dim_) == hidden_dim_npu; + std::vector<__fp16> npu_probe_storage; + if (want_probe) npu_probe_storage.assign(total_hidden_T * hidden_dim_npu, __fp16(0)); + bool probe_complete = want_probe; std::vector<__fp16> input_fp16(chunk_input_elems); bool all_ok = num_chunks > 0; for (size_t c = 0; c < num_chunks && all_ok; ++c) { @@ -2969,8 +2977,22 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a } } __fp16* out_ptr = npu_hidden_storage.data() + c * chunk_output_elems; - size_t written = npu_audio_encoder_->encode( - input_fp16.data(), out_ptr, in_shape, "x", "encoded"); + size_t written = 0; + if (probe_complete) { + __fp16* probe_ptr = npu_probe_storage.data() + c * chunk_output_elems; + written = npu_audio_encoder_->encode_with_secondary( + input_fp16.data(), out_ptr, probe_ptr, in_shape, "x", "encoded", "probe_hidden"); + if (written == 0) { + // ANE model has no probe_hidden output (older bundle); fall + // back to plain encode and skip probe-based handoff. + probe_complete = false; + written = npu_audio_encoder_->encode( + input_fp16.data(), out_ptr, in_shape, "x", "encoded"); + } + } else { + written = npu_audio_encoder_->encode( + input_fp16.data(), out_ptr, in_shape, "x", "encoded"); + } if (written == 0) { all_ok = false; break; } } if (all_ok) { @@ -2980,6 +3002,20 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a if (npu_hidden_T > total_hidden_T) npu_hidden_T = total_hidden_T; CACTUS_LOG_INFO("model", "Parakeet audio encoder ran on NPU (" << num_chunks << " chunks, " << npu_hidden_T << " valid hidden frames)"); + if (want_probe && probe_complete) { + const size_t probe_elems = npu_hidden_T * hidden_dim_npu; + handoff_probe_hidden_.reserve(handoff_probe_hidden_.size() + probe_elems); + const __fp16* pp = npu_probe_storage.data(); + for (size_t i = 0; i < probe_elems; ++i) { + handoff_probe_hidden_.push_back(static_cast(pp[i])); + } + CACTUS_LOG_INFO("cloud_handoff", "Captured " << npu_hidden_T + << " NPU probe frames for the handoff probe"); + } else if (want_probe) { + CACTUS_LOG_WARN("cloud_handoff", "NPU audio encoder has no probe_hidden output; " + "probe-based handoff disabled for this transcription (reconvert with --npu on the " + "probe-enabled branch)"); + } } else { CACTUS_LOG_WARN("model", "NPU audio encoder chunk failed; falling back to CPU graph"); } diff --git a/cactus-engine/src/npu_ane.h b/cactus-engine/src/npu_ane.h index 7594960bf..7d61c5d2c 100644 --- a/cactus-engine/src/npu_ane.h +++ b/cactus-engine/src/npu_ane.h @@ -38,6 +38,14 @@ class ANEEncoder : public NPUEncoder { const std::string& input_name = "x", const std::string& output_name = "") override; + size_t encode_with_secondary(const __fp16* input, + __fp16* primary_output, + __fp16* secondary_output, + const std::vector& shape, + const std::string& input_name = "x", + const std::string& primary_name = "encoded", + const std::string& secondary_name = "probe_hidden") override; + bool is_available() const override; std::vector get_input_shape() const override; diff --git a/cactus-engine/src/npu_ane.mm b/cactus-engine/src/npu_ane.mm index 99638e2e0..20eea5bff 100644 --- a/cactus-engine/src/npu_ane.mm +++ b/cactus-engine/src/npu_ane.mm @@ -229,6 +229,9 @@ - (MLMultiArray*)predictWithInput:(NSString*)inputName data:(const __fp16*)data shape:(NSArray*)shape outputName:(NSString*)outputName; +- (id)predictProviderWithInput:(NSString*)inputName + data:(const __fp16*)data + shape:(NSArray*)shape; - (BOOL)preallocateMultiInputBuffersWithOutputName:(NSString*)outputName; - (size_t)predictMultiInputWithDict:(NSMutableDictionary*)inputDict outputData:(__fp16*)output @@ -482,6 +485,54 @@ - (MLMultiArray*)predictWithInput:(NSString*)inputName return outputFeature.multiArrayValue; } +// Like predictWithInput, but returns the whole prediction so callers can read +// more than one named output (e.g. "encoded" + the handoff probe "probe_hidden") +// from a single ANE run. Uses a fresh prediction (no cached output backings). +- (id)predictProviderWithInput:(NSString*)inputName + data:(const __fp16*)data + shape:(NSArray*)shape { + if (!_model) return nil; + + NSError* error = nil; + MLMultiArrayDataType expectedDataType = MLMultiArrayDataTypeFloat16; + MLFeatureDescription* inputDesc = _modelDescription.inputDescriptionsByName[inputName]; + if (inputDesc && inputDesc.multiArrayConstraint) { + expectedDataType = inputDesc.multiArrayConstraint.dataType; + } + + MLMultiArray* inputArray = [[MLMultiArray alloc] initWithShape:shape + dataType:expectedDataType + error:&error]; + if (error || !inputArray) { + CACTUS_LOG_ERROR("npu", "ANE create input array failed (probe path)"); + return nil; + } + + NSUInteger totalElements = 1; + for (NSNumber* dim in shape) totalElements *= [dim unsignedIntegerValue]; + if (expectedDataType == MLMultiArrayDataTypeFloat16) { + memcpy(inputArray.dataPointer, data, totalElements * sizeof(__fp16)); + } else { + float* inputPtr = (float*)inputArray.dataPointer; + for (NSUInteger i = 0; i < totalElements; i++) inputPtr[i] = (float)data[i]; + } + + MLFeatureValue* inputFeature = [MLFeatureValue featureValueWithMultiArray:inputArray]; + id inputProvider = + [[MLDictionaryFeatureProvider alloc] initWithDictionary:@{inputName: inputFeature} error:&error]; + if (error || !inputProvider) { + CACTUS_LOG_ERROR("npu", "ANE create feature provider failed (probe path)"); + return nil; + } + + id outputProvider = [_model predictionFromFeatures:inputProvider error:&error]; + if (error) { + CACTUS_LOG_ERROR("npu", "ANE prediction failed (probe path): " << [[error localizedDescription] UTF8String]); + return nil; + } + return outputProvider; +} + static size_t copyMLArrayToFP16(MLMultiArray* array, __fp16* output) { size_t count = array.count; if (!array || !output || count == 0) return 0; @@ -899,6 +950,46 @@ - (size_t)predictMultiInputWithDict:(NSMutableDictionary& shape, + const std::string& input_name, + const std::string& primary_name, + const std::string& secondary_name) { + if (!impl_ || !input || !primary_output || !secondary_output) return 0; + + @autoreleasepool { + CactusANEImpl* impl = (__bridge CactusANEImpl*)impl_; + + NSMutableArray* shapeArray = [NSMutableArray arrayWithCapacity:shape.size()]; + for (int dim : shape) [shapeArray addObject:@(dim)]; + + NSString* inName = impl.cachedInputName; + if (!inName) inName = [NSString stringWithUTF8String:input_name.c_str()]; + + id provider = [impl predictProviderWithInput:inName + data:input + shape:shapeArray]; + if (!provider) return 0; + + NSString* secName = [NSString stringWithUTF8String:secondary_name.c_str()]; + MLFeatureValue* secFeat = [provider featureValueForName:secName]; + if (!secFeat || !secFeat.multiArrayValue) { + // Model does not expose the probe output; caller falls back to encode(). + return 0; + } + NSString* priName = [NSString stringWithUTF8String:primary_name.c_str()]; + MLFeatureValue* priFeat = [provider featureValueForName:priName]; + if (!priFeat || !priFeat.multiArrayValue) return 0; + + copyMLArrayToFP16(secFeat.multiArrayValue, secondary_output); + return copyMLArrayToFP16(priFeat.multiArrayValue, primary_output); + } + + return 0; +} + bool ANEEncoder::is_available() const { if (!impl_) return false; CactusANEImpl* impl = (__bridge CactusANEImpl*)impl_; diff --git a/python/cactus/convert/handoff_probe.py b/python/cactus/convert/handoff_probe.py index 3d82472a5..acbba9b05 100644 --- a/python/cactus/convert/handoff_probe.py +++ b/python/cactus/convert/handoff_probe.py @@ -27,17 +27,11 @@ def _candidate_probe_files(output_dir: Path, names: tuple[str, ...]) -> list[Path]: - cwd = Path.cwd() - roots = [output_dir, cwd, Path.home() / "Downloads"] - out: list[Path] = [] - for root in roots: - for name in names: - out.append(root / name) - # also look one release-dir deep (e.g. v10p6_probe_release/, the kit dir) - out.append(root / "v10p6_probe_release" / "global_attn_probe_v10p6.pt") - out.append(root / "parakeet_handoff_kit" / "parakeet_handoff_probe.pt") - out.append(root / "parakeet_handoff_kit" / "parakeet_handoff_kit" / "parakeet_handoff_probe.pt") - return out + # Same convention as the Gemma4 probe: look for the checkpoint by plain + # filename in the bundle output dir and the dir the convert command runs in + # (cwd), plus ~/Downloads as a convenience. Names may be relative paths. + roots = [output_dir, Path.cwd(), Path.home() / "Downloads"] + return [root / name for root in roots for name in names] def _candidate_probe_zips(output_dir: Path, zip_names: tuple[str, ...]) -> list[Path]: @@ -149,7 +143,8 @@ def export_gemma4_handoff_probe(output_dir: str | Path, *, model_id: str | None out_dir = Path(output_dir) checkpoint, source = _load_checkpoint( out_dir, - file_names=("probe.pt", "global_attn_probe_v10p6.pt"), + file_names=("probe.pt", "global_attn_probe_v10p6.pt", + "v10p6_probe_release/global_attn_probe_v10p6.pt"), zip_names=("v10p6_probe_release.zip",), inner_names=( "v10p6_probe_release/global_attn_probe_v10p6.pt", @@ -170,7 +165,11 @@ def export_parakeet_handoff_probe(output_dir: str | Path, *, model_id: str | Non Reads the probe dims from the checkpoint tensors (feat_dim=1024, t_h=32, h1=128, h2=64 for parakeet-tdt-0.6b-v2), so it is not tied to the Gemma4 1536-d shape. The capture layer is informational metadata; the engine taps - whatever layer the transpiler exposes as probe_hidden.""" + whatever layer the transpiler exposes as probe_hidden. + + Detected like the Gemma4 probe: a plain ``parakeet_handoff_probe.pt`` in the + bundle output dir or the convert working directory (drop the file next to + where you run ``cactus convert``).""" model_key = (model_id or "").lower() if model_key and "parakeet" not in model_key: return False @@ -178,13 +177,9 @@ def export_parakeet_handoff_probe(output_dir: str | Path, *, model_id: str | Non out_dir = Path(output_dir) checkpoint, source = _load_checkpoint( out_dir, - file_names=("parakeet_handoff_probe.pt", "probe.pt"), - zip_names=("parakeet_handoff_kit.zip",), - inner_names=( - "parakeet_handoff_kit/parakeet_handoff_probe.pt", - "parakeet_handoff_kit/parakeet_handoff_kit/parakeet_handoff_probe.pt", - "parakeet_handoff_probe.pt", - ), + file_names=("parakeet_handoff_probe.pt",), + zip_names=(), + inner_names=(), ) if checkpoint is None: return False diff --git a/python/cactus/transpile/npu/audio.py b/python/cactus/transpile/npu/audio.py index a60e4f690..7fc38936a 100644 --- a/python/cactus/transpile/npu/audio.py +++ b/python/cactus/transpile/npu/audio.py @@ -69,6 +69,16 @@ def emit_audio_encoder_mlpackage( wrapper = AudioEncoderWrapper(audio_module, baked_inputs) wrapper.eval() + # Encoders that expose a cloud-handoff probe return (encoded, probe_hidden); + # name the extra output so the engine can read it from the same prediction. + with torch.no_grad(): + sample_out = wrapper(example_input) + n_outputs = len(sample_out) if isinstance(sample_out, (tuple, list)) else 1 + output_names = [output_name] + (["probe_hidden"] if n_outputs == 2 else []) + if n_outputs > 2: + print(f"npu.audio: unexpected {n_outputs} encoder outputs; emitting only {output_name}") + output_names = [output_name] + try: with torch.no_grad(): exported = torch.export.export(wrapper, (example_input,)) @@ -88,7 +98,7 @@ def emit_audio_encoder_mlpackage( mlmodel = ct.convert( exported, inputs=[ct.TensorType(name=input_name, shape=tuple(example_input.shape))], - outputs=[ct.TensorType(name=output_name)], + outputs=[ct.TensorType(name=name) for name in output_names], compute_precision=ct.precision.FLOAT16, convert_to="mlprogram", minimum_deployment_target=target_attr, @@ -112,5 +122,6 @@ def emit_audio_encoder_mlpackage( print(f"npu.audio: mlpackage save failed ({type(exc).__name__}: {exc})") return None - print(f"npu.audio: wrote {out_path} (input_shape={tuple(example_input.shape)})") + print(f"npu.audio: wrote {out_path} (input_shape={tuple(example_input.shape)}, " + f"outputs={output_names})") return filename diff --git a/python/cactus/transpile/tdt_runtime.py b/python/cactus/transpile/tdt_runtime.py index ee97a7ab3..dd7c45433 100644 --- a/python/cactus/transpile/tdt_runtime.py +++ b/python/cactus/transpile/tdt_runtime.py @@ -777,7 +777,8 @@ def forward(self, x: torch.Tensor, conv_mask: torch.Tensor, key_bias: torch.Tens class ParakeetTDTANEEncoder(nn.Module): - def __init__(self, encoder: "ParakeetTDTEncoder", seq_len: int): + def __init__(self, encoder: "ParakeetTDTEncoder", seq_len: int, + capture_index: int = PARAKEET_PROBE_CAPTURE_LAYER): super().__init__() self.pre_encode = encoder.pre_encode self.layers = nn.ModuleList( @@ -785,6 +786,10 @@ def __init__(self, encoder: "ParakeetTDTEncoder", seq_len: int): ) factor = max(1, int(encoder.layers[0].config.subsampling_factor)) self._subsample_stages = max(0, factor.bit_length() - 1) + num_layers = len(self.layers) + if not (0 <= capture_index < num_layers): + capture_index = max(0, num_layers - 1) + self.capture_index = capture_index def _frame_masks(self, input_features: torch.Tensor, seq_len: int) -> tuple[torch.Tensor, torch.Tensor]: valid = (input_features.abs().sum(-1) > 0).to(input_features.dtype).sum(-1, keepdim=True) @@ -795,12 +800,17 @@ def _frame_masks(self, input_features: torch.Tensor, seq_len: int) -> tuple[torc key_bias = (1.0 - enc_valid).unsqueeze(1) * -1e4 return enc_valid, key_bias - def forward(self, input_features: torch.Tensor) -> torch.Tensor: + def forward(self, input_features: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: x = self.pre_encode(input_features) conv_mask, key_bias = self._frame_masks(input_features, x.shape[1]) - for layer in self.layers: + probe_hidden = x + for index, layer in enumerate(self.layers): x = layer(x, conv_mask, key_bias) - return x + if index == self.capture_index: + probe_hidden = x + # Emit encoded + probe_hidden so the ANE/NPU bundle can run the + # cloud-handoff probe (layer-17 activations), matching the CPU graph. + return x, probe_hidden class ParakeetTDTDecoderCell(nn.Module): From 5d848b79eec8a987894c514a3c49ee5471c1947c Mon Sep 17 00:00:00 2001 From: Karen Mosoyan Date: Wed, 17 Jun 2026 12:53:55 -0700 Subject: [PATCH 3/6] Parakeet ANE: version-aware rel-pos flip; ship probe for v2 only parakeet-tdt-0.6b v2 and v3 use opposite relative-position sign conventions in their checkpoints. The ANE conformer attention's `.flip(1)` on the rel-pos key table is correct for v3 (the published model) but wrong for v2 (yields blank NPU transcriptions). Detect the version from the model id/source (parakeet_checkpoint_version) and flip only for v3+; v2 uses the corrected no-flip path. Override with CACTUS_PARAKEET_ANE_FLIP=0/1. The cloud-handoff probe was trained on v2 layer-17 activations only, so: - expose probe_hidden (CPU graph + ANE 2-output) and ship handoff_probe.bin for v2 only; v3 (and other versions) get the standard single-output encoder. - export_parakeet_handoff_probe now gates on v2 in the model id. Validated end-to-end on the ANE: v2 transcribes correctly + probe computes a calibrated p_wrong (0.070, matching the kit reference); v3 transcribes correctly with no probe. CPU graph paths for both also transcribe correctly. --- python/cactus/convert/handoff_probe.py | 9 ++- python/cactus/transpile/tdt_runtime.py | 82 ++++++++++++++++++++------ 2 files changed, 71 insertions(+), 20 deletions(-) diff --git a/python/cactus/convert/handoff_probe.py b/python/cactus/convert/handoff_probe.py index acbba9b05..9637de3f7 100644 --- a/python/cactus/convert/handoff_probe.py +++ b/python/cactus/convert/handoff_probe.py @@ -169,10 +169,17 @@ def export_parakeet_handoff_probe(output_dir: str | Path, *, model_id: str | Non Detected like the Gemma4 probe: a plain ``parakeet_handoff_probe.pt`` in the bundle output dir or the convert working directory (drop the file next to - where you run ``cactus convert``).""" + where you run ``cactus convert``). + + The probe was trained on parakeet-tdt-0.6b-v2 layer-17 activations only, so + it is packaged for v2 and NOT for v3 (or other versions), whose encoder + activations differ and would give uncalibrated p_wrong.""" model_key = (model_id or "").lower() if model_key and "parakeet" not in model_key: return False + if model_key and "v2" not in model_key: + # Probe only ships with the v2 checkpoint it was trained on. + return False out_dir = Path(output_dir) checkpoint, source = _load_checkpoint( diff --git a/python/cactus/transpile/tdt_runtime.py b/python/cactus/transpile/tdt_runtime.py index dd7c45433..5d78bcd1b 100644 --- a/python/cactus/transpile/tdt_runtime.py +++ b/python/cactus/transpile/tdt_runtime.py @@ -664,6 +664,28 @@ def forward(self, input_features: torch.Tensor) -> torch.Tensor: PARAKEET_PROBE_CAPTURE_LAYER = int(os.environ.get("CACTUS_PARAKEET_PROBE_LAYER", "17")) +def parakeet_checkpoint_version(model_source: str | None) -> int: + """Best-effort Parakeet-TDT version from the model id/source path. + Defaults to 3 (the published/supported model) when unidentifiable.""" + source = (model_source or "").lower() + match = re.search(r"[-_]v(\d+)(?:[-_/.]|$)", source) + return int(match.group(1)) if match else 3 + + +def parakeet_ane_flip_relpos(model_source: str | None) -> bool: + """Whether the ANE conformer attention flips its relative-position key table. + + parakeet-tdt-0.6b-v2 and v3 use OPPOSITE relative-position sign conventions: + the `.flip(1)` is correct for v3 (the published/supported model) but wrong + for v2 (yields blank NPU transcriptions). Default to v3 behaviour (flip) for + anything not clearly identified as v2. Override with CACTUS_PARAKEET_ANE_FLIP=0/1. + """ + env = os.environ.get("CACTUS_PARAKEET_ANE_FLIP") + if env is not None: + return env.strip().lower() not in ("0", "false", "no", "off", "") + return parakeet_checkpoint_version(model_source) != 2 # v2 -> no flip; v3+ -> flip + + class ParakeetTDTEncoderProbe(nn.Module): """Wraps the Parakeet encoder so it returns (encoder_hidden_states, probe_hidden), where probe_hidden is the raw output of conformer layer @@ -699,7 +721,7 @@ def forward(self, input_features: torch.Tensor) -> tuple[torch.Tensor, torch.Ten class ParakeetTDTANESelfAttention(nn.Module): - def __init__(self, attn: "ParakeetTDTSelfAttention", seq_len: int): + def __init__(self, attn: "ParakeetTDTSelfAttention", seq_len: int, flip_relpos: bool = True): super().__init__() self.linear_q = attn.linear_q self.linear_k = attn.linear_k @@ -719,7 +741,11 @@ def __init__(self, attn: "ParakeetTDTSelfAttention", seq_len: int): with torch.no_grad(): rel_k = attn.linear_pos(rel_pos).view( 1, 2 * int(seq_len) - 1, self.config.attention_heads, self.config.attention_head_dim - ).flip(1) + ) + # The relative-position sign convention differs by checkpoint: v3 + # needs the flip, v2 does not (see parakeet_ane_flip_relpos). + if flip_relpos: + rel_k = rel_k.flip(1) self.register_buffer("_rel_k", rel_k.detach(), persistent=False) def forward(self, x: torch.Tensor, key_bias: torch.Tensor | None = None) -> torch.Tensor: @@ -755,11 +781,11 @@ def forward(self, x: torch.Tensor, key_bias: torch.Tensor | None = None) -> torc class ParakeetTDTANEEncoderLayer(nn.Module): - def __init__(self, layer: "ParakeetTDTEncoderLayer", seq_len: int): + def __init__(self, layer: "ParakeetTDTEncoderLayer", seq_len: int, flip_relpos: bool = True): super().__init__() self.feed_forward1 = layer.feed_forward1 self.feed_forward2 = layer.feed_forward2 - self.self_attn = ParakeetTDTANESelfAttention(layer.self_attn, seq_len) + self.self_attn = ParakeetTDTANESelfAttention(layer.self_attn, seq_len, flip_relpos=flip_relpos) self.conv = layer.conv self.norm_feed_forward1 = layer.norm_feed_forward1 self.norm_self_att = layer.norm_self_att @@ -778,11 +804,13 @@ def forward(self, x: torch.Tensor, conv_mask: torch.Tensor, key_bias: torch.Tens class ParakeetTDTANEEncoder(nn.Module): def __init__(self, encoder: "ParakeetTDTEncoder", seq_len: int, - capture_index: int = PARAKEET_PROBE_CAPTURE_LAYER): + capture_index: int = PARAKEET_PROBE_CAPTURE_LAYER, + flip_relpos: bool = True, + emit_probe_hidden: bool = True): super().__init__() self.pre_encode = encoder.pre_encode self.layers = nn.ModuleList( - [ParakeetTDTANEEncoderLayer(layer, seq_len) for layer in encoder.layers] + [ParakeetTDTANEEncoderLayer(layer, seq_len, flip_relpos=flip_relpos) for layer in encoder.layers] ) factor = max(1, int(encoder.layers[0].config.subsampling_factor)) self._subsample_stages = max(0, factor.bit_length() - 1) @@ -790,6 +818,7 @@ def __init__(self, encoder: "ParakeetTDTEncoder", seq_len: int, if not (0 <= capture_index < num_layers): capture_index = max(0, num_layers - 1) self.capture_index = capture_index + self.emit_probe_hidden = emit_probe_hidden def _frame_masks(self, input_features: torch.Tensor, seq_len: int) -> tuple[torch.Tensor, torch.Tensor]: valid = (input_features.abs().sum(-1) > 0).to(input_features.dtype).sum(-1, keepdim=True) @@ -800,7 +829,7 @@ def _frame_masks(self, input_features: torch.Tensor, seq_len: int) -> tuple[torc key_bias = (1.0 - enc_valid).unsqueeze(1) * -1e4 return enc_valid, key_bias - def forward(self, input_features: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def forward(self, input_features: torch.Tensor): x = self.pre_encode(input_features) conv_mask, key_bias = self._frame_masks(input_features, x.shape[1]) probe_hidden = x @@ -808,9 +837,12 @@ def forward(self, input_features: torch.Tensor) -> tuple[torch.Tensor, torch.Ten x = layer(x, conv_mask, key_bias) if index == self.capture_index: probe_hidden = x - # Emit encoded + probe_hidden so the ANE/NPU bundle can run the - # cloud-handoff probe (layer-17 activations), matching the CPU graph. - return x, probe_hidden + # When the cloud-handoff probe ships (v2 only), emit encoder + + # probe_hidden so the ANE/NPU bundle can run the probe (layer-17 + # activations). Otherwise emit only the encoded output (v3/standard). + if self.emit_probe_hidden: + return x, probe_hidden + return x class ParakeetTDTDecoderCell(nn.Module): @@ -1051,21 +1083,33 @@ def build_parakeet_tdt_component_specs( npu_input_features[:, :valid_frames, :] = input_features[:, :valid_frames, :] with torch.no_grad(): ane_seq_len = int(model.encoder.pre_encode(npu_input_features).shape[1]) - ane_encoder = ParakeetTDTANEEncoder(model.encoder, ane_seq_len) - - # CPU encoder additionally emits probe_hidden (layer-17 activations) for the - # cloud-handoff probe. The NPU/ANE encoder keeps its single "encoded" output - # (the offline NPU emit traces npu_module directly and ignores output_keys), - # so probe-based handoff is only available on the CPU graph path. - probe_encoder = ParakeetTDTEncoderProbe(model.encoder, PARAKEET_PROBE_CAPTURE_LAYER) + model_source = getattr(model.config, "model_source", None) + version = parakeet_checkpoint_version(model_source) + # v2 and v3 use opposite relative-position sign conventions; the ANE + # encoder flips its rel-pos table only for v3+ (correct), not v2. + ane_flip = parakeet_ane_flip_relpos(model_source) + # The cloud-handoff probe was trained on v2 layer-17 activations only, so + # expose probe_hidden (CPU graph + ANE) for v2 and ship the standard + # single-output encoder for v3/other versions. + include_probe = (version == 2) + + if include_probe: + encoder_module = ParakeetTDTEncoderProbe(model.encoder, PARAKEET_PROBE_CAPTURE_LAYER) + encoder_output_keys = ("encoder_hidden_states", "probe_hidden") + else: + encoder_module = model.encoder + encoder_output_keys = ("encoder_hidden_states",) + ane_encoder = ParakeetTDTANEEncoder( + model.encoder, ane_seq_len, flip_relpos=ane_flip, emit_probe_hidden=include_probe + ) return [ ComponentModuleSpec( component="audio_encoder", - module=probe_encoder, + module=encoder_module, example_inputs=(input_features,), input_keys=("input_features",), - output_keys=("encoder_hidden_states", "probe_hidden"), + output_keys=encoder_output_keys, graph_meta={**common_graph_meta, "component": "audio_encoder"}, metadata={"family": "parakeet_tdt", "task": "tdt_transcription"}, npu_module=ane_encoder, From 4a3431d74f09054ff154eac9f03f33916743ca7d Mon Sep 17 00:00:00 2001 From: Karen Mosoyan Date: Wed, 17 Jun 2026 13:38:46 -0700 Subject: [PATCH 4/6] added stuff back --- cactus-engine/src/engine.h | 4 - cactus-engine/src/model.cpp | 36 +---- cactus-engine/src/npu_ane.mm | 10 +- cactus-engine/src/transcribe.cpp | 8 +- cactus-engine/tests/transcribe.cpp | 13 +- python/cactus/convert/handoff_probe.py | 192 ++++++++++++------------- python/cactus/transpile/npu/audio.py | 10 +- python/cactus/transpile/tdt_runtime.py | 62 +------- 8 files changed, 111 insertions(+), 224 deletions(-) diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index 4a06951ff..b3ca50545 100644 --- a/cactus-engine/src/engine.h +++ b/cactus-engine/src/engine.h @@ -39,10 +39,6 @@ class NPUEncoder { const std::vector& shape, const std::string& input_name = "x", const std::string& output_name = "") = 0; - // Run one prediction and copy two named outputs (e.g. "encoded" plus the - // cloud-handoff probe's "probe_hidden"). Returns elements written to - // primary_output, or 0 if the model does not expose secondary_name (caller - // should fall back to encode()). Default: unsupported. virtual size_t encode_with_secondary( const __fp16* input, __fp16* primary_output, __fp16* secondary_output, const std::vector& shape, diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index c2d2ae3ee..1068f9e88 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -354,9 +354,6 @@ void Model::maybe_capture_handoff_probe_hidden(const Component& comp) { if (rows == 0) return; const auto* data = static_cast(comp.graph->get_output(node)); if (!data) return; - // Append every row of probe_hidden. Text decode steps emit a single row - // (one token), so this preserves prior behaviour; the Parakeet audio encoder - // emits all T time frames in one pass, which the probe then pools over. for (size_t row = 0; row < rows; ++row) { size_t base = row * static_cast(handoff_probe_feat_dim_); for (size_t i = 0; i < handoff_probe_feat_dim_; ++i) { @@ -641,21 +638,11 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, } } if (load_handoff_probe()) { - // The probe reads a "probe_hidden" graph output. For Parakeet ASR it is - // exposed by the audio_encoder (layer-17 activations); for text models - // by the decoder_step. Disable the probe if the bundle predates it. - bool exposes_probe_hidden = false; - const char* reconvert_hint = - "decoder_step does not expose probe_hidden; reconvert Gemma4 to enable probe-based handoff"; - if (config_.model_type == Config::ModelType::PARAKEET_TDT) { - exposes_probe_hidden = audio_encoder_ && output_index(*audio_encoder_, "probe_hidden") >= 0; - reconvert_hint = - "audio_encoder does not expose probe_hidden; reconvert Parakeet to enable probe-based handoff"; - } else { - exposes_probe_hidden = decoder_ && output_index(*decoder_, "probe_hidden") >= 0; - } - if (!exposes_probe_hidden) { - CACTUS_LOG_WARN("cloud_handoff", "Handoff probe is packaged, but " << reconvert_hint); + const Component* probe_comp = (config_.model_type == Config::ModelType::PARAKEET_TDT) + ? audio_encoder_ : decoder_; + if (!probe_comp || output_index(*probe_comp, "probe_hidden") < 0) { + CACTUS_LOG_WARN("cloud_handoff", "Handoff probe is packaged, but probe_hidden is not exposed; " + "reconvert to enable probe-based handoff"); handoff_probe_loaded_ = false; } } @@ -2955,9 +2942,6 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a const size_t num_chunks = (copy_frames + window_frames - 1) / window_frames; const size_t total_hidden_T = num_chunks * window_hidden; npu_hidden_storage.assign(total_hidden_T * hidden_dim_npu, __fp16(0)); - // If a handoff probe is packaged and the ANE encoder exposes a - // probe_hidden output, capture layer-17 activations alongside the - // encoded output (one ANE run produces both). const bool want_probe = handoff_probe_loaded_ && static_cast(handoff_probe_feat_dim_) == hidden_dim_npu; std::vector<__fp16> npu_probe_storage; @@ -2983,8 +2967,6 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a written = npu_audio_encoder_->encode_with_secondary( input_fp16.data(), out_ptr, probe_ptr, in_shape, "x", "encoded", "probe_hidden"); if (written == 0) { - // ANE model has no probe_hidden output (older bundle); fall - // back to plain encode and skip probe-based handoff. probe_complete = false; written = npu_audio_encoder_->encode( input_fp16.data(), out_ptr, in_shape, "x", "encoded"); @@ -3012,9 +2994,8 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a CACTUS_LOG_INFO("cloud_handoff", "Captured " << npu_hidden_T << " NPU probe frames for the handoff probe"); } else if (want_probe) { - CACTUS_LOG_WARN("cloud_handoff", "NPU audio encoder has no probe_hidden output; " - "probe-based handoff disabled for this transcription (reconvert with --npu on the " - "probe-enabled branch)"); + CACTUS_LOG_WARN("cloud_handoff", "NPU audio encoder lacks probe_hidden; " + "skipping probe-based handoff for this transcription"); } } else { CACTUS_LOG_WARN("model", "NPU audio encoder chunk failed; falling back to CPU graph"); @@ -3023,9 +3004,6 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a } if (!used_npu) { audio_enc->graph->execute(); - // Capture layer-17 activations for the cloud-handoff probe. Only the CPU - // graph exposes probe_hidden; the NPU encoder emits only "encoded", so - // probe-based handoff is unavailable on the NPU path. maybe_capture_handoff_probe_hidden(*audio_enc); } diff --git a/cactus-engine/src/npu_ane.mm b/cactus-engine/src/npu_ane.mm index 20eea5bff..e2d2ea412 100644 --- a/cactus-engine/src/npu_ane.mm +++ b/cactus-engine/src/npu_ane.mm @@ -485,9 +485,6 @@ - (MLMultiArray*)predictWithInput:(NSString*)inputName return outputFeature.multiArrayValue; } -// Like predictWithInput, but returns the whole prediction so callers can read -// more than one named output (e.g. "encoded" + the handoff probe "probe_hidden") -// from a single ANE run. Uses a fresh prediction (no cached output backings). - (id)predictProviderWithInput:(NSString*)inputName data:(const __fp16*)data shape:(NSArray*)shape { @@ -504,7 +501,7 @@ - (MLMultiArray*)predictWithInput:(NSString*)inputName dataType:expectedDataType error:&error]; if (error || !inputArray) { - CACTUS_LOG_ERROR("npu", "ANE create input array failed (probe path)"); + CACTUS_LOG_ERROR("npu", "ANE create input array failed: " << (error ? [[error localizedDescription] UTF8String] : "")); return nil; } @@ -521,13 +518,13 @@ - (MLMultiArray*)predictWithInput:(NSString*)inputName id inputProvider = [[MLDictionaryFeatureProvider alloc] initWithDictionary:@{inputName: inputFeature} error:&error]; if (error || !inputProvider) { - CACTUS_LOG_ERROR("npu", "ANE create feature provider failed (probe path)"); + CACTUS_LOG_ERROR("npu", "ANE create feature provider failed: " << (error ? [[error localizedDescription] UTF8String] : "")); return nil; } id outputProvider = [_model predictionFromFeatures:inputProvider error:&error]; if (error) { - CACTUS_LOG_ERROR("npu", "ANE prediction failed (probe path): " << [[error localizedDescription] UTF8String]); + CACTUS_LOG_ERROR("npu", "ANE prediction failed: " << [[error localizedDescription] UTF8String]); return nil; } return outputProvider; @@ -976,7 +973,6 @@ - (size_t)predictMultiInputWithDict:(NSMutableDictionary(1, static_cast(handle->model->get_config().num_mel_bins)); std::vector audio_features; - std::vector handoff_audio_samples; // 16k mono, pre-emphasis-free, for cloud handoff + std::vector handoff_audio_samples; if (is_parakeet) { auto cfg = cactus::audio::get_parakeet_spectrogram_config(); const size_t waveform_samples = audio_samples.size(); @@ -427,10 +427,6 @@ int cactus_transcribe( : 0.0f; float confidence = 1.0f - mean_entropy; - // Cloud-handoff decision for Parakeet, driven by the handoff probe's - // p_wrong (confidence = 1 - p_wrong). Mirrors the Gemma4 path in - // complete.cpp. TDT decoding produces no token entropy, so without the - // probe confidence stays 1.0 and no handoff occurs. bool cloud_handoff_used = false; std::string handoff_reason; float reported_threshold = -1.0f; @@ -461,7 +457,7 @@ int cactus_transcribe( confidence = std::max(0.0f, std::min(1.0f, 1.0f - p_wrong)); CACTUS_LOG_DEBUG("cloud_handoff", "Parakeet handoff probe p_wrong=" << p_wrong << " confidence=" << confidence); - if (confidence < threshold && !handoff_audio_samples.empty()) { + if (confidence < threshold) { CACTUS_LOG_WARN("cloud_handoff", "Cloud transcription handoff triggered: p_wrong=" << p_wrong << " confidence=" << confidence << " threshold=" << threshold); std::vector pcm16(handoff_audio_samples.size()); diff --git a/cactus-engine/tests/transcribe.cpp b/cactus-engine/tests/transcribe.cpp index 084207df3..84d22b0e9 100644 --- a/cactus-engine/tests/transcribe.cpp +++ b/cactus-engine/tests/transcribe.cpp @@ -104,18 +104,7 @@ static int transcribe_file(cactus_model_t model, const std::string& audio_path, const std::string& model_path, const std::string& language) { std::string prompt = whisper_prompt(model_path, language); std::vector response_buffer(RESPONSE_BUFFER_SIZE, 0); - // Cloud-handoff knobs are env-overridable so this binary can exercise the - // Parakeet handoff-probe path (CACTUS_TEST_AUTO_HANDOFF=1, optional - // CACTUS_TEST_CONFIDENCE_THRESHOLD / CACTUS_TEST_CLOUD_TIMEOUT_MS). - const bool auto_handoff = std::getenv("CACTUS_TEST_AUTO_HANDOFF") != nullptr; - const char* thr = std::getenv("CACTUS_TEST_CONFIDENCE_THRESHOLD"); - const char* cto = std::getenv("CACTUS_TEST_CLOUD_TIMEOUT_MS"); - std::ostringstream opts; - opts << "{\"max_tokens\":500,\"language\":\"" << language << "\"" - << ",\"auto_handoff\":" << (auto_handoff ? "true" : "false") - << ",\"confidence_threshold\":" << (thr ? thr : "-1") - << ",\"cloud_timeout_ms\":" << (cto ? cto : "15000") << "}"; - std::string options = opts.str(); + std::string options = "{\"max_tokens\":500,\"auto_handoff\":false,\"language\":\"" + language + "\"}"; auto start = std::chrono::steady_clock::now(); int rc = cactus_transcribe( diff --git a/python/cactus/convert/handoff_probe.py b/python/cactus/convert/handoff_probe.py index 9637de3f7..e9348c23a 100644 --- a/python/cactus/convert/handoff_probe.py +++ b/python/cactus/convert/handoff_probe.py @@ -3,7 +3,6 @@ import io import hashlib import json -import os import struct import zipfile from pathlib import Path @@ -26,51 +25,61 @@ ) -def _candidate_probe_files(output_dir: Path, names: tuple[str, ...]) -> list[Path]: - # Same convention as the Gemma4 probe: look for the checkpoint by plain - # filename in the bundle output dir and the dir the convert command runs in - # (cwd), plus ~/Downloads as a convenience. Names may be relative paths. - roots = [output_dir, Path.cwd(), Path.home() / "Downloads"] - return [root / name for root in roots for name in names] +def _candidate_probe_files(output_dir: Path) -> list[Path]: + cwd = Path.cwd() + return [ + output_dir / "probe.pt", + output_dir / "global_attn_probe_v10p6.pt", + cwd / "probe.pt", + cwd / "v10p6_probe_release" / "global_attn_probe_v10p6.pt", + Path.home() / "Downloads" / "probe.pt", + Path.home() / "Downloads" / "v10p6_probe_release" / "global_attn_probe_v10p6.pt", + ] -def _candidate_probe_zips(output_dir: Path, zip_names: tuple[str, ...]) -> list[Path]: +def _candidate_probe_zips(output_dir: Path) -> list[Path]: cwd = Path.cwd() - return [root / name for root in (output_dir, cwd, Path.home() / "Downloads") for name in zip_names] + return [ + output_dir / "v10p6_probe_release.zip", + cwd / "v10p6_probe_release.zip", + Path.home() / "Downloads" / "v10p6_probe_release.zip", + ] -def _load_checkpoint_from_zip(zip_path: Path, inner_names: tuple[str, ...]) -> Any: +def _load_checkpoint_from_zip(zip_path: Path) -> Any: import torch with zipfile.ZipFile(zip_path) as zf: - available = set(zf.namelist()) - for name in inner_names: - if name in available: + names = set(zf.namelist()) + for name in ( + "v10p6_probe_release/global_attn_probe_v10p6.pt", + "global_attn_probe_v10p6.pt", + "probe.pt", + ): + if name in names: with zf.open(name) as f: - return torch.load(io.BytesIO(f.read()), map_location="cpu", weights_only=False) + return torch.load(io.BytesIO(f.read()), map_location="cpu") raise FileNotFoundError(f"no probe checkpoint found in {zip_path}") def _state_dict_from_checkpoint(checkpoint: Any) -> dict[str, Any]: if isinstance(checkpoint, dict): - for key in ("model_state", "state_dict"): + for key in ("state_dict", "model_state"): value = checkpoint.get(key) if isinstance(value, dict): return value return checkpoint -def _load_checkpoint( - output_dir: Path, file_names: tuple[str, ...], zip_names: tuple[str, ...], inner_names: tuple[str, ...] -) -> tuple[Any, str] | tuple[None, None]: +def _load_checkpoint(output_dir: Path) -> tuple[Any, str] | tuple[None, None]: import torch - for path in _candidate_probe_files(output_dir, file_names): + for path in _candidate_probe_files(output_dir): if path.exists(): - return torch.load(path, map_location="cpu", weights_only=False), str(path) - for path in _candidate_probe_zips(output_dir, zip_names): + return torch.load(path, map_location="cpu"), str(path) + for path in _candidate_probe_zips(output_dir): if path.exists(): - return _load_checkpoint_from_zip(path, inner_names), str(path) + return _load_checkpoint_from_zip(path), str(path) return None, None @@ -84,113 +93,92 @@ def _tensor_hash(state: dict[str, Any]) -> str: return digest.hexdigest() -def _probe_dims(state: dict[str, Any]) -> tuple[int, int, int, int]: - """Read (feat_dim, t_h, h1, h2) from tensor shapes and validate the - architecture the C++ loader (Model::load_handoff_probe) assumes.""" +def export_gemma4_handoff_probe(output_dir: str | Path, *, model_id: str | None = None) -> bool: + """Package the Gemma4 v10p6 cloud-handoff probe into a C++-readable bundle file.""" + model_key = (model_id or "").lower() + if model_key and "gemma-4" not in model_key and "gemma4" not in model_key: + return False + + out_dir = Path(output_dir) + checkpoint, source = _load_checkpoint(out_dir) + if checkpoint is None: + return False + + state = _state_dict_from_checkpoint(checkpoint) missing = [key for key in _ORDERED_KEYS if key not in state] if missing: raise RuntimeError(f"handoff probe checkpoint missing tensors: {', '.join(missing)}") - feat_dim = int(state["norm.weight"].shape[0]) - t_h = int(state["proj.weight"].shape[0]) - h1 = int(state["head.0.weight"].shape[0]) - h2 = int(state["head.2.weight"].shape[0]) - expect = { - "norm.weight": (feat_dim,), "norm.bias": (feat_dim,), - "proj.weight": (t_h, feat_dim), "proj.bias": (t_h,), - "attn_query": (t_h,), - "head.0.weight": (h1, t_h), "head.0.bias": (h1,), - "head.2.weight": (h2, h1), "head.2.bias": (h2,), - "head.4.weight": (1, h2), "head.4.bias": (1,), - } - for key, shape in expect.items(): - got = tuple(int(d) for d in state[key].shape) - if got != shape: - raise RuntimeError(f"handoff probe tensor {key}: expected {shape}, got {got}") - return feat_dim, t_h, h1, h2 - -def _write_probe_bundle(state: dict[str, Any], out_dir: Path, source: str, *, layer: int, fmt: str) -> dict: - feat_dim, t_h, h1, h2 = _probe_dims(state) out_dir.mkdir(parents=True, exist_ok=True) probe_path = out_dir / "handoff_probe.bin" with probe_path.open("wb") as f: f.write(_PROBE_MAGIC) - f.write(struct.pack(" bool: - """Package the Gemma4 v10p6 cloud-handoff probe into a C++-readable bundle file.""" +def export_parakeet_handoff_probe(output_dir: str | Path, *, model_id: str | None = None) -> bool: + """Package the Parakeet-TDT v2 cloud-handoff probe into a C++-readable bundle file.""" model_key = (model_id or "").lower() - if model_key and "gemma-4" not in model_key and "gemma4" not in model_key: + if model_key and ("parakeet" not in model_key or "v2" not in model_key): return False out_dir = Path(output_dir) - checkpoint, source = _load_checkpoint( - out_dir, - file_names=("probe.pt", "global_attn_probe_v10p6.pt", - "v10p6_probe_release/global_attn_probe_v10p6.pt"), - zip_names=("v10p6_probe_release.zip",), - inner_names=( - "v10p6_probe_release/global_attn_probe_v10p6.pt", - "global_attn_probe_v10p6.pt", - "probe.pt", - ), - ) + checkpoint = source = None + for path in ( + out_dir / "parakeet_handoff_probe.pt", + Path.cwd() / "parakeet_handoff_probe.pt", + Path.home() / "Downloads" / "parakeet_handoff_probe.pt", + ): + if path.exists(): + import torch + + checkpoint, source = torch.load(path, map_location="cpu"), str(path) + break if checkpoint is None: return False - state = _state_dict_from_checkpoint(checkpoint) - _write_probe_bundle(state, out_dir, source, layer=28, fmt="cactus_handoff_probe_v10p6") - return True - -def export_parakeet_handoff_probe(output_dir: str | Path, *, model_id: str | None = None) -> bool: - """Package the Parakeet-TDT cloud-handoff probe into a C++-readable bundle. - - Reads the probe dims from the checkpoint tensors (feat_dim=1024, t_h=32, - h1=128, h2=64 for parakeet-tdt-0.6b-v2), so it is not tied to the Gemma4 - 1536-d shape. The capture layer is informational metadata; the engine taps - whatever layer the transpiler exposes as probe_hidden. + state = _state_dict_from_checkpoint(checkpoint) + missing = [key for key in _ORDERED_KEYS if key not in state] + if missing: + raise RuntimeError(f"handoff probe checkpoint missing tensors: {', '.join(missing)}") - Detected like the Gemma4 probe: a plain ``parakeet_handoff_probe.pt`` in the - bundle output dir or the convert working directory (drop the file next to - where you run ``cactus convert``). + feat_dim = int(state["norm.weight"].shape[0]) + t_h = int(state["proj.weight"].shape[0]) + h1 = int(state["head.0.weight"].shape[0]) + h2 = int(state["head.2.weight"].shape[0]) - The probe was trained on parakeet-tdt-0.6b-v2 layer-17 activations only, so - it is packaged for v2 and NOT for v3 (or other versions), whose encoder - activations differ and would give uncalibrated p_wrong.""" - model_key = (model_id or "").lower() - if model_key and "parakeet" not in model_key: - return False - if model_key and "v2" not in model_key: - # Probe only ships with the v2 checkpoint it was trained on. - return False + out_dir.mkdir(parents=True, exist_ok=True) + probe_path = out_dir / "handoff_probe.bin" + with probe_path.open("wb") as f: + f.write(_PROBE_MAGIC) + f.write(struct.pack(" 2: - print(f"npu.audio: unexpected {n_outputs} encoder outputs; emitting only {output_name}") - output_names = [output_name] + output_names = [output_name] + if isinstance(sample_out, (tuple, list)) and len(sample_out) == 2: + output_names.append("probe_hidden") try: with torch.no_grad(): diff --git a/python/cactus/transpile/tdt_runtime.py b/python/cactus/transpile/tdt_runtime.py index 5d78bcd1b..8231bc4e8 100644 --- a/python/cactus/transpile/tdt_runtime.py +++ b/python/cactus/transpile/tdt_runtime.py @@ -658,56 +658,19 @@ def forward(self, input_features: torch.Tensor) -> torch.Tensor: return x -# Encoder layer (0-indexed) whose output the cloud-handoff probe was trained on. -# parakeet-tdt-0.6b-v2 has 24 conformer layers; the probe taps layer 17 (NOT the -# final layer). Override with CACTUS_PARAKEET_PROBE_LAYER if a probe is retrained. -PARAKEET_PROBE_CAPTURE_LAYER = int(os.environ.get("CACTUS_PARAKEET_PROBE_LAYER", "17")) +PARAKEET_PROBE_CAPTURE_LAYER = 17 def parakeet_checkpoint_version(model_source: str | None) -> int: - """Best-effort Parakeet-TDT version from the model id/source path. - Defaults to 3 (the published/supported model) when unidentifiable.""" - source = (model_source or "").lower() - match = re.search(r"[-_]v(\d+)(?:[-_/.]|$)", source) + match = re.search(r"[-_]v(\d+)(?:[-_/.]|$)", (model_source or "").lower()) return int(match.group(1)) if match else 3 -def parakeet_ane_flip_relpos(model_source: str | None) -> bool: - """Whether the ANE conformer attention flips its relative-position key table. - - parakeet-tdt-0.6b-v2 and v3 use OPPOSITE relative-position sign conventions: - the `.flip(1)` is correct for v3 (the published/supported model) but wrong - for v2 (yields blank NPU transcriptions). Default to v3 behaviour (flip) for - anything not clearly identified as v2. Override with CACTUS_PARAKEET_ANE_FLIP=0/1. - """ - env = os.environ.get("CACTUS_PARAKEET_ANE_FLIP") - if env is not None: - return env.strip().lower() not in ("0", "false", "no", "off", "") - return parakeet_checkpoint_version(model_source) != 2 # v2 -> no flip; v3+ -> flip - - class ParakeetTDTEncoderProbe(nn.Module): - """Wraps the Parakeet encoder so it returns (encoder_hidden_states, - probe_hidden), where probe_hidden is the raw output of conformer layer - ``capture_index``. Mirrors the Gemma4 ``capture_layer_index`` mechanism so - the C++ engine can run the cloud-handoff probe on ASR activations. - - Reuses the encoder's own ``pre_encode`` and ``layers`` submodules, so traced - parameter names match ``model.encoder`` exactly and weight mapping is - unchanged.""" - def __init__(self, encoder: "ParakeetTDTEncoder", capture_index: int): super().__init__() self.pre_encode = encoder.pre_encode self.layers = encoder.layers - num_layers = len(self.layers) - if not (0 <= capture_index < num_layers): - clamped = max(0, num_layers - 1) - print( - f"tdt_runtime: probe capture layer {capture_index} out of range for " - f"{num_layers}-layer encoder; clamping to {clamped}" - ) - capture_index = clamped self.capture_index = capture_index def forward(self, input_features: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: @@ -742,8 +705,6 @@ def __init__(self, attn: "ParakeetTDTSelfAttention", seq_len: int, flip_relpos: rel_k = attn.linear_pos(rel_pos).view( 1, 2 * int(seq_len) - 1, self.config.attention_heads, self.config.attention_head_dim ) - # The relative-position sign convention differs by checkpoint: v3 - # needs the flip, v2 does not (see parakeet_ane_flip_relpos). if flip_relpos: rel_k = rel_k.flip(1) self.register_buffer("_rel_k", rel_k.detach(), persistent=False) @@ -814,9 +775,6 @@ def __init__(self, encoder: "ParakeetTDTEncoder", seq_len: int, ) factor = max(1, int(encoder.layers[0].config.subsampling_factor)) self._subsample_stages = max(0, factor.bit_length() - 1) - num_layers = len(self.layers) - if not (0 <= capture_index < num_layers): - capture_index = max(0, num_layers - 1) self.capture_index = capture_index self.emit_probe_hidden = emit_probe_hidden @@ -837,9 +795,6 @@ def forward(self, input_features: torch.Tensor): x = layer(x, conv_mask, key_bias) if index == self.capture_index: probe_hidden = x - # When the cloud-handoff probe ships (v2 only), emit encoder + - # probe_hidden so the ANE/NPU bundle can run the probe (layer-17 - # activations). Otherwise emit only the encoded output (v3/standard). if self.emit_probe_hidden: return x, probe_hidden return x @@ -1083,15 +1038,8 @@ def build_parakeet_tdt_component_specs( npu_input_features[:, :valid_frames, :] = input_features[:, :valid_frames, :] with torch.no_grad(): ane_seq_len = int(model.encoder.pre_encode(npu_input_features).shape[1]) - model_source = getattr(model.config, "model_source", None) - version = parakeet_checkpoint_version(model_source) - # v2 and v3 use opposite relative-position sign conventions; the ANE - # encoder flips its rel-pos table only for v3+ (correct), not v2. - ane_flip = parakeet_ane_flip_relpos(model_source) - # The cloud-handoff probe was trained on v2 layer-17 activations only, so - # expose probe_hidden (CPU graph + ANE) for v2 and ship the standard - # single-output encoder for v3/other versions. - include_probe = (version == 2) + version = parakeet_checkpoint_version(getattr(model.config, "model_source", None)) + include_probe = version == 2 if include_probe: encoder_module = ParakeetTDTEncoderProbe(model.encoder, PARAKEET_PROBE_CAPTURE_LAYER) @@ -1100,7 +1048,7 @@ def build_parakeet_tdt_component_specs( encoder_module = model.encoder encoder_output_keys = ("encoder_hidden_states",) ane_encoder = ParakeetTDTANEEncoder( - model.encoder, ane_seq_len, flip_relpos=ane_flip, emit_probe_hidden=include_probe + model.encoder, ane_seq_len, flip_relpos=version != 2, emit_probe_hidden=include_probe ) return [ From d15ba432d4ab0f7e4d4a4d4f5bdccf0caff0910d Mon Sep 17 00:00:00 2001 From: Karen Mosoyan Date: Wed, 17 Jun 2026 23:01:55 -0700 Subject: [PATCH 5/6] removed parakeet v2 --- cactus-engine/src/engine.h | 12 +--- cactus-engine/src/model.cpp | 41 ++++-------- cactus-engine/src/npu_ane.h | 8 --- cactus-engine/src/npu_ane.mm | 87 -------------------------- python/cactus/convert/handoff_probe.py | 18 +++--- python/cactus/transpile/npu/audio.py | 11 +--- python/cactus/transpile/tdt_runtime.py | 66 +++---------------- 7 files changed, 35 insertions(+), 208 deletions(-) diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index b3ca50545..eda731a2c 100644 --- a/cactus-engine/src/engine.h +++ b/cactus-engine/src/engine.h @@ -39,16 +39,6 @@ class NPUEncoder { const std::vector& shape, const std::string& input_name = "x", const std::string& output_name = "") = 0; - virtual size_t encode_with_secondary( - const __fp16* input, __fp16* primary_output, __fp16* secondary_output, - const std::vector& shape, - const std::string& input_name = "x", - const std::string& primary_name = "encoded", - const std::string& secondary_name = "probe_hidden") { - (void)input; (void)primary_output; (void)secondary_output; (void)shape; - (void)input_name; (void)primary_name; (void)secondary_name; - return 0; - } virtual bool is_available() const = 0; virtual std::vector get_input_shape() const = 0; virtual std::vector get_output_shape() const = 0; @@ -848,7 +838,7 @@ class Model { int output_index(const Component& comp, const std::string& name) const; uint32_t argmax_last_logits(float* out_uncertainty = nullptr); bool load_handoff_probe(); - void maybe_capture_handoff_probe_hidden(const Component& comp); + void maybe_capture_handoff_probe_hidden(const Component& comp, const std::string& output_name = "probe_hidden"); void run_vision_encoder(const std::string& image_path); void run_vision_encoder_lfm2_vl(const std::string& image_path); void encode_lfm2_vl_image_into_features(const std::string& image_path); diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 1068f9e88..f2392beb6 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -343,9 +343,9 @@ bool Model::has_handoff_probe_rollout() const { && handoff_probe_hidden_.size() >= static_cast(handoff_probe_feat_dim_); } -void Model::maybe_capture_handoff_probe_hidden(const Component& comp) { +void Model::maybe_capture_handoff_probe_hidden(const Component& comp, const std::string& output_name) { if (!handoff_probe_loaded_ || handoff_probe_feat_dim_ == 0) return; - int idx = output_index(comp, "probe_hidden"); + int idx = output_index(comp, output_name); if (idx < 0 || static_cast(idx) >= comp.output_node_ids.size()) return; size_t node = static_cast(comp.output_node_ids[idx]); const auto& desc = comp.graph->get_output_buffer(node); @@ -638,10 +638,11 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, } } if (load_handoff_probe()) { - const Component* probe_comp = (config_.model_type == Config::ModelType::PARAKEET_TDT) - ? audio_encoder_ : decoder_; - if (!probe_comp || output_index(*probe_comp, "probe_hidden") < 0) { - CACTUS_LOG_WARN("cloud_handoff", "Handoff probe is packaged, but probe_hidden is not exposed; " + const bool is_parakeet = config_.model_type == Config::ModelType::PARAKEET_TDT; + const Component* probe_comp = is_parakeet ? audio_encoder_ : decoder_; + const char* probe_output = is_parakeet ? "encoder_hidden_states" : "probe_hidden"; + if (!probe_comp || output_index(*probe_comp, probe_output) < 0) { + CACTUS_LOG_WARN("cloud_handoff", "Handoff probe is packaged, but its probe input is not exposed; " "reconvert to enable probe-based handoff"); handoff_probe_loaded_ = false; } @@ -2944,9 +2945,6 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a npu_hidden_storage.assign(total_hidden_T * hidden_dim_npu, __fp16(0)); const bool want_probe = handoff_probe_loaded_ && static_cast(handoff_probe_feat_dim_) == hidden_dim_npu; - std::vector<__fp16> npu_probe_storage; - if (want_probe) npu_probe_storage.assign(total_hidden_T * hidden_dim_npu, __fp16(0)); - bool probe_complete = want_probe; std::vector<__fp16> input_fp16(chunk_input_elems); bool all_ok = num_chunks > 0; for (size_t c = 0; c < num_chunks && all_ok; ++c) { @@ -2961,20 +2959,8 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a } } __fp16* out_ptr = npu_hidden_storage.data() + c * chunk_output_elems; - size_t written = 0; - if (probe_complete) { - __fp16* probe_ptr = npu_probe_storage.data() + c * chunk_output_elems; - written = npu_audio_encoder_->encode_with_secondary( - input_fp16.data(), out_ptr, probe_ptr, in_shape, "x", "encoded", "probe_hidden"); - if (written == 0) { - probe_complete = false; - written = npu_audio_encoder_->encode( - input_fp16.data(), out_ptr, in_shape, "x", "encoded"); - } - } else { - written = npu_audio_encoder_->encode( - input_fp16.data(), out_ptr, in_shape, "x", "encoded"); - } + size_t written = npu_audio_encoder_->encode( + input_fp16.data(), out_ptr, in_shape, "x", "encoded"); if (written == 0) { all_ok = false; break; } } if (all_ok) { @@ -2984,18 +2970,15 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a if (npu_hidden_T > total_hidden_T) npu_hidden_T = total_hidden_T; CACTUS_LOG_INFO("model", "Parakeet audio encoder ran on NPU (" << num_chunks << " chunks, " << npu_hidden_T << " valid hidden frames)"); - if (want_probe && probe_complete) { + if (want_probe) { const size_t probe_elems = npu_hidden_T * hidden_dim_npu; handoff_probe_hidden_.reserve(handoff_probe_hidden_.size() + probe_elems); - const __fp16* pp = npu_probe_storage.data(); + const __fp16* pp = npu_hidden_storage.data(); for (size_t i = 0; i < probe_elems; ++i) { handoff_probe_hidden_.push_back(static_cast(pp[i])); } CACTUS_LOG_INFO("cloud_handoff", "Captured " << npu_hidden_T << " NPU probe frames for the handoff probe"); - } else if (want_probe) { - CACTUS_LOG_WARN("cloud_handoff", "NPU audio encoder lacks probe_hidden; " - "skipping probe-based handoff for this transcription"); } } else { CACTUS_LOG_WARN("model", "NPU audio encoder chunk failed; falling back to CPU graph"); @@ -3004,7 +2987,7 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a } if (!used_npu) { audio_enc->graph->execute(); - maybe_capture_handoff_probe_hidden(*audio_enc); + maybe_capture_handoff_probe_hidden(*audio_enc, "encoder_hidden_states"); } int hidden_idx = output_index(*audio_enc, "encoder_hidden_states"); diff --git a/cactus-engine/src/npu_ane.h b/cactus-engine/src/npu_ane.h index 7d61c5d2c..7594960bf 100644 --- a/cactus-engine/src/npu_ane.h +++ b/cactus-engine/src/npu_ane.h @@ -38,14 +38,6 @@ class ANEEncoder : public NPUEncoder { const std::string& input_name = "x", const std::string& output_name = "") override; - size_t encode_with_secondary(const __fp16* input, - __fp16* primary_output, - __fp16* secondary_output, - const std::vector& shape, - const std::string& input_name = "x", - const std::string& primary_name = "encoded", - const std::string& secondary_name = "probe_hidden") override; - bool is_available() const override; std::vector get_input_shape() const override; diff --git a/cactus-engine/src/npu_ane.mm b/cactus-engine/src/npu_ane.mm index e2d2ea412..99638e2e0 100644 --- a/cactus-engine/src/npu_ane.mm +++ b/cactus-engine/src/npu_ane.mm @@ -229,9 +229,6 @@ - (MLMultiArray*)predictWithInput:(NSString*)inputName data:(const __fp16*)data shape:(NSArray*)shape outputName:(NSString*)outputName; -- (id)predictProviderWithInput:(NSString*)inputName - data:(const __fp16*)data - shape:(NSArray*)shape; - (BOOL)preallocateMultiInputBuffersWithOutputName:(NSString*)outputName; - (size_t)predictMultiInputWithDict:(NSMutableDictionary*)inputDict outputData:(__fp16*)output @@ -485,51 +482,6 @@ - (MLMultiArray*)predictWithInput:(NSString*)inputName return outputFeature.multiArrayValue; } -- (id)predictProviderWithInput:(NSString*)inputName - data:(const __fp16*)data - shape:(NSArray*)shape { - if (!_model) return nil; - - NSError* error = nil; - MLMultiArrayDataType expectedDataType = MLMultiArrayDataTypeFloat16; - MLFeatureDescription* inputDesc = _modelDescription.inputDescriptionsByName[inputName]; - if (inputDesc && inputDesc.multiArrayConstraint) { - expectedDataType = inputDesc.multiArrayConstraint.dataType; - } - - MLMultiArray* inputArray = [[MLMultiArray alloc] initWithShape:shape - dataType:expectedDataType - error:&error]; - if (error || !inputArray) { - CACTUS_LOG_ERROR("npu", "ANE create input array failed: " << (error ? [[error localizedDescription] UTF8String] : "")); - return nil; - } - - NSUInteger totalElements = 1; - for (NSNumber* dim in shape) totalElements *= [dim unsignedIntegerValue]; - if (expectedDataType == MLMultiArrayDataTypeFloat16) { - memcpy(inputArray.dataPointer, data, totalElements * sizeof(__fp16)); - } else { - float* inputPtr = (float*)inputArray.dataPointer; - for (NSUInteger i = 0; i < totalElements; i++) inputPtr[i] = (float)data[i]; - } - - MLFeatureValue* inputFeature = [MLFeatureValue featureValueWithMultiArray:inputArray]; - id inputProvider = - [[MLDictionaryFeatureProvider alloc] initWithDictionary:@{inputName: inputFeature} error:&error]; - if (error || !inputProvider) { - CACTUS_LOG_ERROR("npu", "ANE create feature provider failed: " << (error ? [[error localizedDescription] UTF8String] : "")); - return nil; - } - - id outputProvider = [_model predictionFromFeatures:inputProvider error:&error]; - if (error) { - CACTUS_LOG_ERROR("npu", "ANE prediction failed: " << [[error localizedDescription] UTF8String]); - return nil; - } - return outputProvider; -} - static size_t copyMLArrayToFP16(MLMultiArray* array, __fp16* output) { size_t count = array.count; if (!array || !output || count == 0) return 0; @@ -947,45 +899,6 @@ - (size_t)predictMultiInputWithDict:(NSMutableDictionary& shape, - const std::string& input_name, - const std::string& primary_name, - const std::string& secondary_name) { - if (!impl_ || !input || !primary_output || !secondary_output) return 0; - - @autoreleasepool { - CactusANEImpl* impl = (__bridge CactusANEImpl*)impl_; - - NSMutableArray* shapeArray = [NSMutableArray arrayWithCapacity:shape.size()]; - for (int dim : shape) [shapeArray addObject:@(dim)]; - - NSString* inName = impl.cachedInputName; - if (!inName) inName = [NSString stringWithUTF8String:input_name.c_str()]; - - id provider = [impl predictProviderWithInput:inName - data:input - shape:shapeArray]; - if (!provider) return 0; - - NSString* secName = [NSString stringWithUTF8String:secondary_name.c_str()]; - MLFeatureValue* secFeat = [provider featureValueForName:secName]; - if (!secFeat || !secFeat.multiArrayValue) { - return 0; - } - NSString* priName = [NSString stringWithUTF8String:primary_name.c_str()]; - MLFeatureValue* priFeat = [provider featureValueForName:priName]; - if (!priFeat || !priFeat.multiArrayValue) return 0; - - copyMLArrayToFP16(secFeat.multiArrayValue, secondary_output); - return copyMLArrayToFP16(priFeat.multiArrayValue, primary_output); - } - - return 0; -} - bool ANEEncoder::is_available() const { if (!impl_) return false; CactusANEImpl* impl = (__bridge CactusANEImpl*)impl_; diff --git a/python/cactus/convert/handoff_probe.py b/python/cactus/convert/handoff_probe.py index e9348c23a..873e55f70 100644 --- a/python/cactus/convert/handoff_probe.py +++ b/python/cactus/convert/handoff_probe.py @@ -131,19 +131,21 @@ def export_gemma4_handoff_probe(output_dir: str | Path, *, model_id: str | None return True +_PARAKEET_PROBE_FILE = "parakeet_v3_probe.pt" +_PARAKEET_PROBE_LAYER = 23 + + def export_parakeet_handoff_probe(output_dir: str | Path, *, model_id: str | None = None) -> bool: - """Package the Parakeet-TDT v2 cloud-handoff probe into a C++-readable bundle file.""" + """Package the Parakeet-TDT-v3 cloud-handoff probe into a C++-readable bundle file.""" model_key = (model_id or "").lower() - if model_key and ("parakeet" not in model_key or "v2" not in model_key): + if model_key and "parakeet" not in model_key: return False + layer = _PARAKEET_PROBE_LAYER out_dir = Path(output_dir) checkpoint = source = None - for path in ( - out_dir / "parakeet_handoff_probe.pt", - Path.cwd() / "parakeet_handoff_probe.pt", - Path.home() / "Downloads" / "parakeet_handoff_probe.pt", - ): + for path in (out_dir / _PARAKEET_PROBE_FILE, Path.cwd() / _PARAKEET_PROBE_FILE, + Path.home() / "Downloads" / _PARAKEET_PROBE_FILE): if path.exists(): import torch @@ -174,7 +176,7 @@ def export_parakeet_handoff_probe(output_dir: str | Path, *, model_id: str | Non metadata = { "format": "cactus_handoff_probe_parakeet", "source": source, - "layer": 17, + "layer": layer, "feat_dim": feat_dim, "t_h": t_h, "output": str(probe_path.name), diff --git a/python/cactus/transpile/npu/audio.py b/python/cactus/transpile/npu/audio.py index 4df29db44..a60e4f690 100644 --- a/python/cactus/transpile/npu/audio.py +++ b/python/cactus/transpile/npu/audio.py @@ -69,12 +69,6 @@ def emit_audio_encoder_mlpackage( wrapper = AudioEncoderWrapper(audio_module, baked_inputs) wrapper.eval() - with torch.no_grad(): - sample_out = wrapper(example_input) - output_names = [output_name] - if isinstance(sample_out, (tuple, list)) and len(sample_out) == 2: - output_names.append("probe_hidden") - try: with torch.no_grad(): exported = torch.export.export(wrapper, (example_input,)) @@ -94,7 +88,7 @@ def emit_audio_encoder_mlpackage( mlmodel = ct.convert( exported, inputs=[ct.TensorType(name=input_name, shape=tuple(example_input.shape))], - outputs=[ct.TensorType(name=name) for name in output_names], + outputs=[ct.TensorType(name=output_name)], compute_precision=ct.precision.FLOAT16, convert_to="mlprogram", minimum_deployment_target=target_attr, @@ -118,6 +112,5 @@ def emit_audio_encoder_mlpackage( print(f"npu.audio: mlpackage save failed ({type(exc).__name__}: {exc})") return None - print(f"npu.audio: wrote {out_path} (input_shape={tuple(example_input.shape)}, " - f"outputs={output_names})") + print(f"npu.audio: wrote {out_path} (input_shape={tuple(example_input.shape)})") return filename diff --git a/python/cactus/transpile/tdt_runtime.py b/python/cactus/transpile/tdt_runtime.py index 8231bc4e8..948e4f301 100644 --- a/python/cactus/transpile/tdt_runtime.py +++ b/python/cactus/transpile/tdt_runtime.py @@ -658,33 +658,8 @@ def forward(self, input_features: torch.Tensor) -> torch.Tensor: return x -PARAKEET_PROBE_CAPTURE_LAYER = 17 - - -def parakeet_checkpoint_version(model_source: str | None) -> int: - match = re.search(r"[-_]v(\d+)(?:[-_/.]|$)", (model_source or "").lower()) - return int(match.group(1)) if match else 3 - - -class ParakeetTDTEncoderProbe(nn.Module): - def __init__(self, encoder: "ParakeetTDTEncoder", capture_index: int): - super().__init__() - self.pre_encode = encoder.pre_encode - self.layers = encoder.layers - self.capture_index = capture_index - - def forward(self, input_features: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - x = self.pre_encode(input_features) - probe_hidden = x - for index, layer in enumerate(self.layers): - x = layer(x) - if index == self.capture_index: - probe_hidden = x - return x, probe_hidden - - class ParakeetTDTANESelfAttention(nn.Module): - def __init__(self, attn: "ParakeetTDTSelfAttention", seq_len: int, flip_relpos: bool = True): + def __init__(self, attn: "ParakeetTDTSelfAttention", seq_len: int): super().__init__() self.linear_q = attn.linear_q self.linear_k = attn.linear_k @@ -704,9 +679,7 @@ def __init__(self, attn: "ParakeetTDTSelfAttention", seq_len: int, flip_relpos: with torch.no_grad(): rel_k = attn.linear_pos(rel_pos).view( 1, 2 * int(seq_len) - 1, self.config.attention_heads, self.config.attention_head_dim - ) - if flip_relpos: - rel_k = rel_k.flip(1) + ).flip(1) self.register_buffer("_rel_k", rel_k.detach(), persistent=False) def forward(self, x: torch.Tensor, key_bias: torch.Tensor | None = None) -> torch.Tensor: @@ -742,11 +715,11 @@ def forward(self, x: torch.Tensor, key_bias: torch.Tensor | None = None) -> torc class ParakeetTDTANEEncoderLayer(nn.Module): - def __init__(self, layer: "ParakeetTDTEncoderLayer", seq_len: int, flip_relpos: bool = True): + def __init__(self, layer: "ParakeetTDTEncoderLayer", seq_len: int): super().__init__() self.feed_forward1 = layer.feed_forward1 self.feed_forward2 = layer.feed_forward2 - self.self_attn = ParakeetTDTANESelfAttention(layer.self_attn, seq_len, flip_relpos=flip_relpos) + self.self_attn = ParakeetTDTANESelfAttention(layer.self_attn, seq_len) self.conv = layer.conv self.norm_feed_forward1 = layer.norm_feed_forward1 self.norm_self_att = layer.norm_self_att @@ -764,19 +737,14 @@ def forward(self, x: torch.Tensor, conv_mask: torch.Tensor, key_bias: torch.Tens class ParakeetTDTANEEncoder(nn.Module): - def __init__(self, encoder: "ParakeetTDTEncoder", seq_len: int, - capture_index: int = PARAKEET_PROBE_CAPTURE_LAYER, - flip_relpos: bool = True, - emit_probe_hidden: bool = True): + def __init__(self, encoder: "ParakeetTDTEncoder", seq_len: int): super().__init__() self.pre_encode = encoder.pre_encode self.layers = nn.ModuleList( - [ParakeetTDTANEEncoderLayer(layer, seq_len, flip_relpos=flip_relpos) for layer in encoder.layers] + [ParakeetTDTANEEncoderLayer(layer, seq_len) for layer in encoder.layers] ) factor = max(1, int(encoder.layers[0].config.subsampling_factor)) self._subsample_stages = max(0, factor.bit_length() - 1) - self.capture_index = capture_index - self.emit_probe_hidden = emit_probe_hidden def _frame_masks(self, input_features: torch.Tensor, seq_len: int) -> tuple[torch.Tensor, torch.Tensor]: valid = (input_features.abs().sum(-1) > 0).to(input_features.dtype).sum(-1, keepdim=True) @@ -790,13 +758,8 @@ def _frame_masks(self, input_features: torch.Tensor, seq_len: int) -> tuple[torc def forward(self, input_features: torch.Tensor): x = self.pre_encode(input_features) conv_mask, key_bias = self._frame_masks(input_features, x.shape[1]) - probe_hidden = x - for index, layer in enumerate(self.layers): + for layer in self.layers: x = layer(x, conv_mask, key_bias) - if index == self.capture_index: - probe_hidden = x - if self.emit_probe_hidden: - return x, probe_hidden return x @@ -1038,18 +1001,9 @@ def build_parakeet_tdt_component_specs( npu_input_features[:, :valid_frames, :] = input_features[:, :valid_frames, :] with torch.no_grad(): ane_seq_len = int(model.encoder.pre_encode(npu_input_features).shape[1]) - version = parakeet_checkpoint_version(getattr(model.config, "model_source", None)) - include_probe = version == 2 - - if include_probe: - encoder_module = ParakeetTDTEncoderProbe(model.encoder, PARAKEET_PROBE_CAPTURE_LAYER) - encoder_output_keys = ("encoder_hidden_states", "probe_hidden") - else: - encoder_module = model.encoder - encoder_output_keys = ("encoder_hidden_states",) - ane_encoder = ParakeetTDTANEEncoder( - model.encoder, ane_seq_len, flip_relpos=version != 2, emit_probe_hidden=include_probe - ) + encoder_module = model.encoder + encoder_output_keys = ("encoder_hidden_states",) + ane_encoder = ParakeetTDTANEEncoder(model.encoder, ane_seq_len) return [ ComponentModuleSpec( From 8f940ebc063cf28c87b5234f7fc7cacc5e61b01f Mon Sep 17 00:00:00 2001 From: jakmro Date: Thu, 18 Jun 2026 10:42:04 +0200 Subject: [PATCH 6/6] Tidy Parakeet handoff probe: dedup exporters, guard probe capture Signed-off-by: jakmro --- cactus-engine/src/model.cpp | 4 +- python/cactus/convert/handoff_probe.py | 97 +++++++++++++------------- python/cactus/transpile/tdt_runtime.py | 8 +-- 3 files changed, 53 insertions(+), 56 deletions(-) diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index f2392beb6..25a2cd871 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -333,7 +333,7 @@ bool Model::load_handoff_probe() { handoff_probe_hidden_.clear(); handoff_probe_loaded_ = true; - CACTUS_LOG_INFO("cloud_handoff", "Loaded Gemma4 v10p6 handoff probe from " << path); + CACTUS_LOG_INFO("cloud_handoff", "Loaded handoff probe from " << path); return true; } @@ -350,6 +350,8 @@ void Model::maybe_capture_handoff_probe_hidden(const Component& comp, const std: size_t node = static_cast(comp.output_node_ids[idx]); const auto& desc = comp.graph->get_output_buffer(node); if (desc.total_size < handoff_probe_feat_dim_) return; + if (!desc.shape.empty() && + static_cast(desc.shape.back()) != static_cast(handoff_probe_feat_dim_)) return; size_t rows = desc.total_size / handoff_probe_feat_dim_; if (rows == 0) return; const auto* data = static_cast(comp.graph->get_output(node)); diff --git a/python/cactus/convert/handoff_probe.py b/python/cactus/convert/handoff_probe.py index 873e55f70..763c67f53 100644 --- a/python/cactus/convert/handoff_probe.py +++ b/python/cactus/convert/handoff_probe.py @@ -93,59 +93,77 @@ def _tensor_hash(state: dict[str, Any]) -> str: return digest.hexdigest() -def export_gemma4_handoff_probe(output_dir: str | Path, *, model_id: str | None = None) -> bool: - """Package the Gemma4 v10p6 cloud-handoff probe into a C++-readable bundle file.""" - model_key = (model_id or "").lower() - if model_key and "gemma-4" not in model_key and "gemma4" not in model_key: - return False - - out_dir = Path(output_dir) - checkpoint, source = _load_checkpoint(out_dir) - if checkpoint is None: - return False - +def _probe_state_dict(checkpoint: Any) -> dict[str, Any]: state = _state_dict_from_checkpoint(checkpoint) missing = [key for key in _ORDERED_KEYS if key not in state] if missing: raise RuntimeError(f"handoff probe checkpoint missing tensors: {', '.join(missing)}") + return state + +def _write_probe_bundle(out_dir: Path, state: dict[str, Any], *, feat_dim: int, t_h: int, + h1: int, h2: int, fmt: str, layer: int, source: str | None) -> bool: out_dir.mkdir(parents=True, exist_ok=True) probe_path = out_dir / "handoff_probe.bin" with probe_path.open("wb") as f: f.write(_PROBE_MAGIC) - f.write(struct.pack(" bool: + """Package the Gemma4 v10p6 cloud-handoff probe into a C++-readable bundle file.""" + model_key = (model_id or "").lower() + if model_key and "gemma-4" not in model_key and "gemma4" not in model_key: + return False + + out_dir = Path(output_dir) + checkpoint, source = _load_checkpoint(out_dir) + if checkpoint is None: + return False + + state = _probe_state_dict(checkpoint) + return _write_probe_bundle( + out_dir, state, feat_dim=1536, t_h=32, h1=128, h2=64, + fmt="cactus_handoff_probe_v10p6", layer=28, source=source, + ) + + _PARAKEET_PROBE_FILE = "parakeet_v3_probe.pt" _PARAKEET_PROBE_LAYER = 23 +def _candidate_parakeet_probe_files(output_dir: Path) -> list[Path]: + return [ + output_dir / _PARAKEET_PROBE_FILE, + Path.cwd() / _PARAKEET_PROBE_FILE, + Path.home() / "Downloads" / _PARAKEET_PROBE_FILE, + ] + + def export_parakeet_handoff_probe(output_dir: str | Path, *, model_id: str | None = None) -> bool: """Package the Parakeet-TDT-v3 cloud-handoff probe into a C++-readable bundle file.""" model_key = (model_id or "").lower() if model_key and "parakeet" not in model_key: return False - layer = _PARAKEET_PROBE_LAYER out_dir = Path(output_dir) checkpoint = source = None - for path in (out_dir / _PARAKEET_PROBE_FILE, Path.cwd() / _PARAKEET_PROBE_FILE, - Path.home() / "Downloads" / _PARAKEET_PROBE_FILE): + for path in _candidate_parakeet_probe_files(out_dir): if path.exists(): import torch @@ -154,33 +172,12 @@ def export_parakeet_handoff_probe(output_dir: str | Path, *, model_id: str | Non if checkpoint is None: return False - state = _state_dict_from_checkpoint(checkpoint) - missing = [key for key in _ORDERED_KEYS if key not in state] - if missing: - raise RuntimeError(f"handoff probe checkpoint missing tensors: {', '.join(missing)}") - - feat_dim = int(state["norm.weight"].shape[0]) - t_h = int(state["proj.weight"].shape[0]) - h1 = int(state["head.0.weight"].shape[0]) - h2 = int(state["head.2.weight"].shape[0]) - - out_dir.mkdir(parents=True, exist_ok=True) - probe_path = out_dir / "handoff_probe.bin" - with probe_path.open("wb") as f: - f.write(_PROBE_MAGIC) - f.write(struct.pack(" tuple[torc key_bias = (1.0 - enc_valid).unsqueeze(1) * -1e4 return enc_valid, key_bias - def forward(self, input_features: torch.Tensor): + def forward(self, input_features: torch.Tensor) -> torch.Tensor: x = self.pre_encode(input_features) conv_mask, key_bias = self._frame_masks(input_features, x.shape[1]) for layer in self.layers: @@ -1001,17 +1001,15 @@ def build_parakeet_tdt_component_specs( npu_input_features[:, :valid_frames, :] = input_features[:, :valid_frames, :] with torch.no_grad(): ane_seq_len = int(model.encoder.pre_encode(npu_input_features).shape[1]) - encoder_module = model.encoder - encoder_output_keys = ("encoder_hidden_states",) ane_encoder = ParakeetTDTANEEncoder(model.encoder, ane_seq_len) return [ ComponentModuleSpec( component="audio_encoder", - module=encoder_module, + module=model.encoder, example_inputs=(input_features,), input_keys=("input_features",), - output_keys=encoder_output_keys, + output_keys=("encoder_hidden_states",), graph_meta={**common_graph_meta, "component": "audio_encoder"}, metadata={"family": "parakeet_tdt", "task": "tdt_transcription"}, npu_module=ane_encoder,