diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index 604f730f5..eda731a2c 100644 --- a/cactus-engine/src/engine.h +++ b/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); diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index a955b871f..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; } @@ -343,21 +343,24 @@ 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); 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; - 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)); + 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 +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 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) { @@ -2935,6 +2945,8 @@ 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)); + const bool want_probe = handoff_probe_loaded_ && + static_cast(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,6 +2972,16 @@ 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) { + 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(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"); } @@ -2967,6 +2989,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, "encoder_hidden_states"); } 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..11a5822b7 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; 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 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 +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; 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..763c67f53 100644 --- a/python/cactus/convert/handoff_probe.py +++ b/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(" 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, + )