Fix padded-frame leakage and feature normalization in Parakeet encoder

cactus-engine/src/model.cpp

@@ -4757,6 +4757,21 @@ std::vector<uint32_t> Model::transcribe_parakeet_tdt(const std::vector<float>& a
     write_typed_buffer(feat_buf, feat_desc.precision, transposed.data(),
                        transposed.size() * sizeof(float), Precision::FP32);
 
+    const size_t subsampling = std::max<uint32_t>(1, config_.subsampling_factor);
+    const size_t valid_encoder_frames = (copy_frames + subsampling - 1) / subsampling;
+    int mask_idx = input_index(*audio_enc, "input_features_mask");
+    if (mask_idx < 0) {
+        CACTUS_LOG_ERROR("model", "audio_encoder has no input_features_mask input");
+        return emitted;
+    }
+    auto& mask_buf = audio_enc->input_buffers[mask_idx];
+    const auto& mask_desc = audio_enc->graph->get_output_buffer(
+        static_cast<size_t>(audio_enc->runtime_input_node_ids[mask_idx]));
+    std::vector<float> mask(mask_desc.total_size, 0.0f);
+    std::fill(mask.begin(), mask.begin() + std::min(copy_frames, mask.size()), 1.0f);
+    write_typed_buffer(mask_buf, mask_desc.precision, mask.data(),
+                       mask.size() * sizeof(float), Precision::FP32);
+
     if (should_stop && should_stop->load()) return emitted;
     audio_enc->graph->execute();
     maybe_capture_handoff_probe_hidden(*audio_enc, "encoder_hidden_states");
@@ -4851,13 +4866,8 @@ std::vector<uint32_t> Model::transcribe_parakeet_tdt(const std::vector<float>& a
         };
     }
 
-    size_t commit_to = T;
-    if (stream) {
-        size_t valid_hidden = T;
-        if (expected_frames > 0)
-            valid_hidden = std::min<size_t>(T, (copy_frames * T) / expected_frames);
-        commit_to = (end_frame > 0) ? std::min(end_frame, valid_hidden) : valid_hidden;
-    }
+    size_t commit_to = std::min(T, valid_encoder_frames);
+    if (stream && end_frame > 0) commit_to = std::min(end_frame, commit_to);
     Tokenizer* stream_tok = stream ? get_tokenizer() : nullptr;
     const auto& tdt_vocab_bias = get_vocab_bias();
     const float frame_sec = (160.0f / 16000.0f) *

cactus-engine/src/stream.cpp

@@ -125,7 +125,7 @@ std::vector<float> window_features(std::vector<float> window, size_t mel_bins) {
     const size_t waveform_samples = window.size();
     cactus::audio::apply_preemphasis(window, 0.97f);
     std::vector<float> features = cactus::audio::compute_spectrogram_graph(
-        window, cfg, mel_bins, 0.0f, 8000.0f, cactus::audio::WHISPER_SAMPLE_RATE, 0, 0);
+        window, cfg, mel_bins, 0.0f, 8000.0f, cactus::audio::WHISPER_SAMPLE_RATE, 1, 2);
     cactus::audio::normalize_parakeet_log_mel(features, mel_bins);
     size_t valid_frames = waveform_samples / cfg.hop_length;
     if (valid_frames == 0) valid_frames = 1;

cactus-engine/src/transcribe.cpp

@@ -109,7 +109,7 @@ CACTUS_FFI_EXPORT int cactus_preprocess_audio_features(
             cactus::audio::apply_preemphasis(audio_samples, 0.97f);
             features = cactus::audio::compute_spectrogram_graph(
                 audio_samples, cfg, bins, 0.0f, 8000.0f,
-                cactus::audio::WHISPER_SAMPLE_RATE, 0, 0);
+                cactus::audio::WHISPER_SAMPLE_RATE, 1, 2);
             cactus::audio::normalize_parakeet_log_mel(features, bins);
             size_t valid_frames = waveform_samples / cfg.hop_length;
             if (valid_frames == 0) valid_frames = 1;
@@ -251,7 +251,7 @@ int cactus_transcribe(
             cactus::audio::apply_preemphasis(audio_samples, 0.97f);
             audio_features = cactus::audio::compute_spectrogram_graph(
                 audio_samples, cfg, mel_bins, 0.0f, 8000.0f,
-                cactus::audio::WHISPER_SAMPLE_RATE, 0, 0);
+                cactus::audio::WHISPER_SAMPLE_RATE, 1, 2);
             cactus::audio::normalize_parakeet_log_mel(audio_features, mel_bins);
             size_t valid_frames = waveform_samples / cfg.hop_length;
             if (valid_frames == 0) valid_frames = 1;

cactus-graph/src/io.cpp

@@ -896,8 +896,6 @@ void save_node(CactusGraph& graph, size_t node_id, const std::string& filename)
 
 // MappedFileRegistry implementation
 
-namespace GraphFile {
-
 std::mutex MappedFileRegistry::mutex_;
 std::unordered_map<std::string, std::weak_ptr<MappedFile>> MappedFileRegistry::cache_;
 
@@ -931,8 +929,6 @@ void MappedFileRegistry::clear() {
     cache_.clear();
 }
 
-} // namespace GraphFile
-
 // MappedFile implementation
 
 MappedFile::MappedFile(const std::string& filename)

python/cactus/convert/tests/test_audio_buckets.py

@@ -4,6 +4,7 @@
 
 def test_fibonacci_ladder_capped_at_model_window():
     assert audio_bucket_frames(2003) == [100, 200, 300, 500, 800, 1300, 2003]
+    assert audio_bucket_frames(3000) == [100, 200, 300, 500, 800, 1300, 2100, 3000]
 
 
 def test_window_below_ladder_yields_single_bucket():
@@ -26,7 +27,7 @@ def test_buckets_are_sorted_and_unique():
 
 
 def test_default_ladder_is_fibonacci_seconds():
-    assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13)
+    assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13, 21)
 
 
 def test_env_override_replaces_ladder(monkeypatch):
@@ -42,9 +43,9 @@ def test_env_override_is_sorted_and_deduplicated(monkeypatch):
 
 def test_malformed_env_override_falls_back_to_default(monkeypatch):
     monkeypatch.setenv("CACTUS_TRANSPILER_AUDIO_BUCKETS", "not-a-number")
-    assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13)
+    assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13, 21)
 
 
 def test_empty_env_override_falls_back_to_default(monkeypatch):
     monkeypatch.setenv("CACTUS_TRANSPILER_AUDIO_BUCKETS", "0,-3")
-    assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13)
+    assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13, 21)

python/cactus/convert/tests/test_parakeet_capture.py

@@ -0,0 +1,60 @@
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import numpy as np
+import torch
+
+from cactus.transpile.hf_model import _run_parakeet_tdt_component_decode
+from cactus.transpile.tdt_runtime import build_parakeet_tdt_component_specs
+
+
+def test_capture_window_is_independent_of_validation_audio():
+    features = torch.full((1, 2003, 128), 2.0)
+    model = SimpleNamespace(
+        config=SimpleNamespace(blank_id=0, predictor_num_layers=0),
+        encoder=lambda features, mask: features[:, ::8, :4],
+        initial_decoder_state=lambda **kwargs: (),
+        decoder_step=object(),
+    )
+    with patch.dict("os.environ", {"CACTUS_TRANSPILER_AUDIO_BUCKETS": ""}):
+        specs = build_parakeet_tdt_component_specs(model, named_tensors={"input_features": features})
+    encoders = specs[:-1]
+    assert [int(spec.metadata["audio_frames"]) for spec in encoders] == [100, 200, 300, 500, 800, 1300, 2100, 3000]
+    assert encoders[-1].component == "audio_encoder"
+    assert tuple(encoders[-1].example_inputs[0].shape) == (1, 3000, 128)
+    assert torch.all(features == 2.0)
+    assert not torch.all(encoders[-1].example_inputs[0] == 2.0)
+
+
+def test_validation_preserves_real_features_and_masks_padding():
+    features = torch.randn(1, 2003, 128)
+    prepared = SimpleNamespace(names=("input_features",), tensors=(features,))
+    encoder = SimpleNamespace(metadata={"audio_frames": "2100"})
+    components = {
+        "audio_encoder": SimpleNamespace(metadata={"audio_frames": "3000"}),
+        "audio_encoder_2100": encoder,
+        "decoder": object(),
+    }
+    model = SimpleNamespace(
+        config=SimpleNamespace(subsampling_factor=8, predictor_num_layers=0),
+        initial_decoder_state=lambda **kwargs: (),
+        decode_token_ids=lambda tokens: "test",
+    )
+
+    def execute(selected, *, initial_store):
+        assert selected == [encoder]
+        padded = initial_store["input_features"]
+        mask = initial_store["input_features_mask"]
+        assert tuple(padded.shape) == (1, 2100, 128)
+        assert torch.equal(padded[:, :2003], features)
+        assert torch.all(padded[:, 2003:] == 0)
+        assert torch.all(mask[:, :2003] == 1)
+        assert torch.all(mask[:, 2003:] == 0)
+        return {"encoder_hidden_states": np.zeros((1, 263, 4))}, {}
+
+    with (
+        patch("cactus.transpile.hf_model.execute_component_pipeline", side_effect=execute),
+        patch("cactus.transpile.hf_model.greedy_decode_parakeet_tdt_token_ids", return_value=[]) as decode,
+    ):
+        _run_parakeet_tdt_component_decode(component_graphs=components, model=model, prepared=prepared)
+    assert decode.call_args.kwargs["encoder_hidden_states"].shape == (1, 251, 4)

python/cactus/transpile/audio_preprocess.py

@@ -21,7 +21,7 @@
 _PARAKEET_PREEMPHASIS = 0.97
 _PARAKEET_LOG_FLOOR = np.float32(2**-24)
 _DEFAULT_MAX_AUDIO_SECONDS = 30.0
-_DEFAULT_AUDIO_BUCKET_SECONDS = (1, 2, 3, 5, 8, 13)
+_DEFAULT_AUDIO_BUCKET_SECONDS = (1, 2, 3, 5, 8, 13, 21)
 
 
 def audio_duration_limit_seconds() -> float:

python/cactus/transpile/fusion/conv.py

@@ -22,6 +22,7 @@ class ConvModuleMatch:
     batch_norm_running_var_value_id: str
     pointwise2_weight_value_id: str
     pointwise2_bias_value_id: str | None
+    mask_value_id: str | None
     eps: float
     depthwise_kernel_size: int
     depthwise_padding: int
@@ -51,6 +52,21 @@ def match_conv_module(graph: IRGraph, node: IRNode) -> ConvModuleMatch | None:
         return None
 
     glu = producer(graph, strip_passthrough(graph, depthwise.inputs[0]))
+    mask_multiply: IRNode | None = None
+    mask_value_id: str | None = None
+    if glu is not None and glu.op == "multiply" and len(glu.inputs) == 2:
+        for index, operand in enumerate(glu.inputs):
+            operand_node = producer(graph, strip_passthrough(graph, operand))
+            if operand_node is not None and operand_node.op == "glu":
+                mask_multiply = glu
+                mask_value_id = glu.inputs[1 - index]
+                mask_value = graph.values.get(mask_value_id)
+                if mask_value is None or mask_value.shape is None:
+                    return None
+                if len(mask_value.shape) != 3 or mask_value.shape[1] != 1:
+                    return None
+                glu = operand_node
+                break
     if glu is None or glu.op != "glu" or len(glu.inputs) != 1:
         return None
     if int(glu.attrs.get("axis", -1)) not in {1, -2}:
@@ -83,6 +99,7 @@ def match_conv_module(graph: IRGraph, node: IRNode) -> ConvModuleMatch | None:
         batch_norm_running_var_value_id=batch_norm.inputs[4],
         pointwise2_weight_value_id=pointwise2.inputs[1],
         pointwise2_bias_value_id=pointwise2.inputs[2] if len(pointwise2.inputs) > 2 else None,
+        mask_value_id=mask_value_id,
         eps=float(batch_norm.attrs.get("eps", 1e-5)),
         depthwise_kernel_size=depthwise_kernel_size,
         depthwise_padding=int(depthwise.attrs.get("padding", 0)),
@@ -92,6 +109,7 @@ def match_conv_module(graph: IRGraph, node: IRNode) -> ConvModuleMatch | None:
             silu,
             batch_norm,
             depthwise,
+            mask_multiply,
             glu,
             pointwise1,
             input_permute,

python/cactus/transpile/hf_model.py

@@ -555,11 +555,29 @@ def _run_parakeet_tdt_component_decode(
     model,
     prepared: PreparedInputs,
 ) -> dict[str, object]:
+    input_features = _named_tensor_store(prepared)["input_features"]
+    active_frames = int(input_features.shape[1])
+    encoders = [
+        component for name, component in component_graphs.items()
+        if name.startswith("audio_encoder")
+        and int(component.metadata["audio_frames"]) >= active_frames
+    ]
+    if not encoders:
+        raise ValueError("Parakeet validation audio exceeds the 30-second capture window")
+    encoder = min(encoders, key=lambda component: int(component.metadata["audio_frames"]))
+    capacity_frames = int(encoder.metadata["audio_frames"])
+    padded_features = input_features.new_zeros(
+        (int(input_features.shape[0]), capacity_frames, int(input_features.shape[2]))
+    )
+    padded_features[:, :active_frames, :] = input_features
+    mask = input_features.new_zeros((int(input_features.shape[0]), capacity_frames))
+    mask[:, :active_frames] = 1
     store, _ = execute_component_pipeline(
-        [component_graphs["audio_encoder"]],
-        initial_store=_named_tensor_store(prepared),
+        [encoder],
+        initial_store={"input_features": padded_features, "input_features_mask": mask},
     )
-    encoder_hidden_states = np.asarray(store["encoder_hidden_states"])
+    valid_frames = (active_frames + model.config.subsampling_factor - 1) // model.config.subsampling_factor
+    encoder_hidden_states = np.asarray(store["encoder_hidden_states"])[:, :valid_frames, :]
     batch_size = int(encoder_hidden_states.shape[0])
     if batch_size != 1:
         raise ValueError("Parakeet TDT component decode currently expects batch size 1")

python/cactus/transpile/lower.py

@@ -1908,6 +1908,11 @@ def _lower_ir_node(g: Graph, node: IRNode, env: dict[str, Any], ir: IRGraph) ->
         pointwise2_bias = None
         if bool(node.attrs.get("has_pointwise2_bias", False)):
             pointwise2_bias = _tensor(env, node.inputs[input_index])
+            input_index += 1
+        mask_ncl = None
+        if bool(node.attrs.get("has_mask", False)):
+            mask_ncl = _tensor(env, node.inputs[input_index])
+            input_index += 1
 
         if len(x_nlc.shape) != 3:
             raise NotImplementedError(f"conv_module expects rank-3 NLC input, got {x_nlc.shape}")
@@ -1921,6 +1926,9 @@ def _lower_ir_node(g: Graph, node: IRNode, env: dict[str, Any], ir: IRGraph) ->
 
         current = g.conv1d_pointwise(x_nlc, pointwise1_weight, bias=pointwise1_bias)
         current = g.glu(current, axis=-1)
+        if mask_ncl is not None:
+            mask_nlc = g.reshape(mask_ncl, (int(mask_ncl.shape[0]), int(mask_ncl.shape[2]), 1))
+            current = _lower_binary_op(g, current, mask_nlc, "multiply")
         current = g.conv1d_same_depthwise_k9(current, depthwise_weight, bias=depthwise_bias)
         current = g.batch_norm(
             current,

python/cactus/transpile/ops.py

@@ -215,6 +215,7 @@ class OpSchema:
             "has_pointwise1_bias",
             "has_depthwise_bias",
             "has_pointwise2_bias",
+            "has_mask",
             "depthwise_kernel_size",
             "depthwise_padding",
         ),

python/cactus/transpile/optimize_graph.py

@@ -1206,6 +1206,8 @@ def fuse_conv_modules(graph: IRGraph) -> bool:
         )
         if match.pointwise2_bias_value_id is not None:
             inputs.append(match.pointwise2_bias_value_id)
+        if match.mask_value_id is not None:
+            inputs.append(match.mask_value_id)
 
         node.op = "conv_module"
         node.inputs = inputs
@@ -1214,6 +1216,7 @@ def fuse_conv_modules(graph: IRGraph) -> bool:
             "has_pointwise1_bias": bool(match.pointwise1_bias_value_id is not None),
             "has_depthwise_bias": bool(match.depthwise_bias_value_id is not None),
             "has_pointwise2_bias": bool(match.pointwise2_bias_value_id is not None),
+            "has_mask": bool(match.mask_value_id is not None),
             "depthwise_kernel_size": int(match.depthwise_kernel_size),
             "depthwise_padding": int(match.depthwise_padding),
         }