Repository navigation
Karen/parakeet probe - #739
Merged
Merged
Conversation
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.
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.
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.
Signed-off-by: jakmro <kubamroz124@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.