Karen/parakeet probe

cactus-engine/src/engine.h

@@ -838,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);

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;
 }
 
@@ -343,21 +343,24 @@ bool Model::has_handoff_probe_rollout() const {
         && handoff_probe_hidden_.size() >= static_cast<size_t>(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<size_t>(idx) >= comp.output_node_ids.size()) return;
     size_t node = static_cast<size_t>(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<size_t>(desc.shape.back()) != static_cast<size_t>(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<const uint8_t*>(comp.graph->get_output(node));
     if (!data) return;
-    size_t base = row * static_cast<size_t>(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));
+    for (size_t row = 0; row < rows; ++row) {
+        size_t base = row * static_cast<size_t>(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 +639,15 @@ 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()) {
+        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;
+        }
     }
 
     initialized_ = true;
@@ -2883,6 +2891,8 @@ std::vector<uint32_t> Model::transcribe_parakeet_tdt(const std::vector<float>& a
     std::vector<uint32_t> 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) {
@@ -2935,6 +2945,8 @@ std::vector<uint32_t> Model::transcribe_parakeet_tdt(const std::vector<float>& 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));
+            const bool want_probe = handoff_probe_loaded_ &&
+                static_cast<size_t>(handoff_probe_feat_dim_) == hidden_dim_npu;
             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) {
@@ -2960,13 +2972,24 @@ std::vector<uint32_t> Model::transcribe_parakeet_tdt(const std::vector<float>& 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) {
+                    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_hidden_storage.data();
+                    for (size_t i = 0; i < probe_elems; ++i) {
+                        handoff_probe_hidden_.push_back(static_cast<float>(pp[i]));
+                    }
+                    CACTUS_LOG_INFO("cloud_handoff", "Captured " << npu_hidden_T
+                                    << " NPU probe frames for the handoff probe");
+                }
             } else {
                 CACTUS_LOG_WARN("model", "NPU audio encoder chunk failed; falling back to CPU graph");
             }
         }
     }
     if (!used_npu) {
         audio_enc->graph->execute();
+        maybe_capture_handoff_probe_hidden(*audio_enc, "encoder_hidden_states");
     }
 
     int hidden_idx = output_index(*audio_enc, "encoder_hidden_states");

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 <algorithm>
 #include <chrono>
 #include <cctype>
+#include <cmath>
 #include <cstring>
 #include <mutex>
 
@@ -237,9 +239,11 @@ int cactus_transcribe(
 
         const size_t mel_bins = std::max<size_t>(1, static_cast<size_t>(handle->model->get_config().num_mel_bins));
         std::vector<float> audio_features;
+        std::vector<float> handoff_audio_samples;
         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,72 @@ int cactus_transcribe(
             : 0.0f;
         float confidence = 1.0f - mean_entropy;
 
+        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) {
+                        CACTUS_LOG_WARN("cloud_handoff", "Cloud transcription handoff triggered: p_wrong="
+                            << p_wrong << " confidence=" << confidence << " threshold=" << threshold);
+                        std::vector<int16_t> 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<int16_t>(std::lround(clamped * 32767.0f));
+                        }
+                        std::vector<uint8_t> wav = cactus::ffi::cloud_build_wav(
+                            reinterpret_cast<const uint8_t*>(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<long>(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<std::chrono::microseconds>(end_time - start_time).count() / 1000.0;
@@ -436,7 +506,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;

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

python/cactus/convert/handoff_probe.py

@@ -93,39 +93,91 @@ 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("<IIIII", 1, 1536, 32, 128, 64))
+        f.write(struct.pack("<IIIII", 1, feat_dim, t_h, h1, h2))
         for key in _ORDERED_KEYS:
             tensor = state[key].detach().cpu().contiguous().float().numpy()
             f.write(tensor.tobytes(order="C"))
 
     metadata = {
-        "format": "cactus_handoff_probe_v10p6",
+        "format": fmt,
         "source": source,
-        "layer": 28,
-        "feat_dim": 1536,
-        "t_h": 32,
-        "output": str(probe_path.name),
+        "layer": layer,
+        "feat_dim": feat_dim,
+        "t_h": t_h,
+        "output": probe_path.name,
         "tensor_sha256": _tensor_hash(state),
     }
     (out_dir / "handoff_probe.json").write_text(json.dumps(metadata, indent=2) + "\n")
     return True
+
+
+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 = _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
+
+    out_dir = Path(output_dir)
+    checkpoint = source = None
+    for path in _candidate_parakeet_probe_files(out_dir):
+        if path.exists():
+            import torch
+
+            checkpoint, source = torch.load(path, map_location="cpu"), str(path)
+            break
+    if checkpoint is None:
+        return False
+
+    state = _probe_state_dict(checkpoint)
+    return _write_probe_bundle(
+        out_dir, state,
+        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]),
+        fmt="cactus_handoff_probe_parakeet", layer=_PARAKEET_PROBE_LAYER, source=source,
+    )