From ea6061d3a5b32d6a8188612bc500024bcba275b4 Mon Sep 17 00:00:00 2001 From: jakmro Date: Mon, 17 Aug 2026 13:32:38 +0200 Subject: [PATCH 1/3] feat: implement audio bucket selection logic and corresponding tests --- cactus-engine/src/engine.h | 3 + cactus-engine/src/model.cpp | 31 +++++++- cactus-engine/tests/test_bucket.cpp | 78 +++++++++++++++++++ .../convert/tests/test_audio_buckets.py | 50 ++++++++++++ python/cactus/transpile/audio_preprocess.py | 28 +++++++ python/cactus/transpile/tdt_runtime.py | 31 +++++--- 6 files changed, 211 insertions(+), 10 deletions(-) create mode 100644 cactus-engine/tests/test_bucket.cpp create mode 100644 python/cactus/convert/tests/test_audio_buckets.py diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index 0ee887c3d..d592be929 100644 --- a/cactus-engine/src/engine.h +++ b/cactus-engine/src/engine.h @@ -669,6 +669,8 @@ class ToolCallConstrainer { void feed_gemma_char(char ch); }; +size_t select_bucket(const std::vector& capacities, size_t required_frames); + class Model { public: Model(); @@ -1008,6 +1010,7 @@ class Model { int output_index(const Component& comp, const std::string& name) const; uint32_t argmax_last_logits(float* out_uncertainty = nullptr, float repetition_penalty = 1.0f); bool load_handoff_probe(); + Component* select_audio_encoder(size_t required_frames); 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); diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index d7f1213c5..5383c89bf 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -4619,6 +4619,31 @@ std::vector Model::transcribe_whisper_seq2seq( should_stop); } +size_t select_bucket(const std::vector& capacities, size_t required_frames) { + size_t best = capacities.size(); + for (size_t i = 0; i < capacities.size(); ++i) { + const size_t capacity = capacities[i]; + if (capacity == 0 || capacity < required_frames) continue; + if (best == capacities.size() || capacity < capacities[best]) best = i; + } + return best; +} + +Model::Component* Model::select_audio_encoder(size_t required_frames) { + std::vector candidates; + std::vector capacities; + for (auto& [name, comp] : components_) { + if (name.rfind("audio_encoder", 0) != 0) continue; + auto it = comp.metadata.find("audio_frames"); + if (it == comp.metadata.end()) continue; + candidates.push_back(&comp); + capacities.push_back(parse_size_or_zero(it->second)); + } + const size_t index = select_bucket(capacities, required_frames); + if (index < candidates.size()) return candidates[index]; + return components_.count("audio_encoder") ? &components_.at("audio_encoder") : nullptr; +} + std::vector Model::transcribe_parakeet_tdt(const std::vector& audio_features, ParakeetTdtStreamState* stream, bool is_final, size_t end_frame, @@ -4628,7 +4653,11 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a reset_handoff_probe_rollout(); - Component* audio_enc = components_.count("audio_encoder") ? &components_.at("audio_encoder") : nullptr; + const size_t config_mels = static_cast(config_.num_mel_bins); + const size_t required_frames = config_mels > 0 + ? audio_features.size() / config_mels + : std::numeric_limits::max(); + Component* audio_enc = select_audio_encoder(required_frames); Component* dec = components_.count("decoder") ? &components_.at("decoder") : nullptr; if (!audio_enc || !dec) { CACTUS_LOG_ERROR("model", "Parakeet TDT bundle missing audio_encoder or decoder component"); diff --git a/cactus-engine/tests/test_bucket.cpp b/cactus-engine/tests/test_bucket.cpp new file mode 100644 index 000000000..951393782 --- /dev/null +++ b/cactus-engine/tests/test_bucket.cpp @@ -0,0 +1,78 @@ +#include "test_utils.h" +#include "../src/engine.h" + +#include +#include + +using namespace TestUtils; +using cactus::engine::select_bucket; + +namespace { + +const std::vector kFibonacci = {100, 200, 300, 500, 800, 1300, 2003}; + +bool picks_exact_fit() { + return select_bucket(kFibonacci, 300) == 2; +} + +bool picks_smallest_that_fits() { + return select_bucket(kFibonacci, 201) == 2 + && select_bucket(kFibonacci, 301) == 3 + && select_bucket(kFibonacci, 801) == 5; +} + +bool picks_first_bucket_for_tiny_input() { + return select_bucket(kFibonacci, 1) == 0 + && select_bucket(kFibonacci, 100) == 0; +} + +bool reports_none_when_above_every_bucket() { + return select_bucket(kFibonacci, 2004) == kFibonacci.size(); +} + +bool reports_none_for_empty_ladder() { + return select_bucket({}, 100) == 0; +} + +bool ignores_zero_capacities() { + const std::vector with_zero = {0, 500, 0, 200}; + return select_bucket(with_zero, 150) == 3 + && select_bucket(with_zero, 300) == 1; +} + +bool is_independent_of_declaration_order() { + const std::vector shuffled = {2003, 300, 100, 1300, 200, 800, 500}; + return shuffled[select_bucket(shuffled, 250)] == 300 + && shuffled[select_bucket(shuffled, 900)] == 1300; +} + +bool handles_zero_required_frames() { + return select_bucket(kFibonacci, 0) == 0; +} + +bool unknown_required_frames_falls_back() { + return select_bucket(kFibonacci, std::numeric_limits::max()) == kFibonacci.size(); +} + +bool handles_duplicate_capacities() { + const std::vector duplicated = {300, 300, 800}; + return duplicated[select_bucket(duplicated, 250)] == 300; +} + +} + +int main() { + TestRunner runner("Audio Encoder Bucket Selection"); + runner.run_test("Exact fit", picks_exact_fit()); + runner.run_test("Smallest that fits", picks_smallest_that_fits()); + runner.run_test("Tiny input takes first bucket", picks_first_bucket_for_tiny_input()); + runner.run_test("None when above every bucket", reports_none_when_above_every_bucket()); + runner.run_test("None for empty ladder", reports_none_for_empty_ladder()); + runner.run_test("Zero capacities ignored", ignores_zero_capacities()); + runner.run_test("Order independent", is_independent_of_declaration_order()); + runner.run_test("Zero required frames", handles_zero_required_frames()); + runner.run_test("Unknown required frames falls back", unknown_required_frames_falls_back()); + runner.run_test("Duplicate capacities", handles_duplicate_capacities()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/python/cactus/convert/tests/test_audio_buckets.py b/python/cactus/convert/tests/test_audio_buckets.py new file mode 100644 index 000000000..00b6f6bee --- /dev/null +++ b/python/cactus/convert/tests/test_audio_buckets.py @@ -0,0 +1,50 @@ +from cactus.transpile.audio_preprocess import audio_bucket_frames +from cactus.transpile.audio_preprocess import audio_bucket_seconds + + +def test_fibonacci_ladder_capped_at_model_window(): + assert audio_bucket_frames(2003) == [100, 200, 300, 500, 800, 1300, 2003] + + +def test_window_below_ladder_yields_single_bucket(): + assert audio_bucket_frames(80) == [80] + + +def test_window_equal_to_a_rung_is_not_duplicated(): + assert audio_bucket_frames(800) == [100, 200, 300, 500, 800] + + +def test_non_positive_window_yields_no_buckets(): + assert audio_bucket_frames(0) == [] + assert audio_bucket_frames(-1) == [] + + +def test_buckets_are_sorted_and_unique(): + buckets = audio_bucket_frames(2003) + assert buckets == sorted(buckets) + assert len(buckets) == len(set(buckets)) + + +def test_default_ladder_is_fibonacci_seconds(): + assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13) + + +def test_env_override_replaces_ladder(monkeypatch): + monkeypatch.setenv("CACTUS_TRANSPILER_AUDIO_BUCKETS", "2,4,6") + assert audio_bucket_seconds() == (2, 4, 6) + assert audio_bucket_frames(2003) == [200, 400, 600, 2003] + + +def test_env_override_is_sorted_and_deduplicated(monkeypatch): + monkeypatch.setenv("CACTUS_TRANSPILER_AUDIO_BUCKETS", "6, 2, 4, 2") + assert audio_bucket_seconds() == (2, 4, 6) + + +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) + + +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) diff --git a/python/cactus/transpile/audio_preprocess.py b/python/cactus/transpile/audio_preprocess.py index 5e6d0aeea..6956a5038 100644 --- a/python/cactus/transpile/audio_preprocess.py +++ b/python/cactus/transpile/audio_preprocess.py @@ -21,6 +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) def audio_duration_limit_seconds() -> float: @@ -32,6 +33,33 @@ def audio_duration_limit_seconds() -> float: return max(0.0, value) +def audio_bucket_seconds() -> tuple[int, ...]: + raw = os.environ.get("CACTUS_TRANSPILER_AUDIO_BUCKETS") + if not raw: + return _DEFAULT_AUDIO_BUCKET_SECONDS + try: + values = sorted({int(part) for part in raw.split(",")}) + except (TypeError, ValueError): + return _DEFAULT_AUDIO_BUCKET_SECONDS + positive = tuple(value for value in values if value > 0) + if not positive: + return _DEFAULT_AUDIO_BUCKET_SECONDS + return positive + + +def audio_bucket_frames(max_frames: int) -> list[int]: + if max_frames <= 0: + return [] + frames_per_second = _PARAKEET_SAMPLE_RATE // _PARAKEET_HOP_LENGTH + buckets = [] + for seconds in audio_bucket_seconds(): + frames = seconds * frames_per_second + if frames < max_frames: + buckets.append(frames) + buckets.append(max_frames) + return buckets + + def normalize_audio_samples(samples: np.ndarray) -> np.ndarray: if samples.dtype == np.uint8: return ((samples.astype(np.float32) - 128.0) / 128.0).astype(np.float32) diff --git a/python/cactus/transpile/tdt_runtime.py b/python/cactus/transpile/tdt_runtime.py index 9f0b25dfb..07c44b63c 100644 --- a/python/cactus/transpile/tdt_runtime.py +++ b/python/cactus/transpile/tdt_runtime.py @@ -13,6 +13,7 @@ import torch.nn.functional as F from cactus.transpile.component_pipeline import ComponentModuleSpec +from cactus.transpile.audio_preprocess import audio_bucket_frames from cactus.transpile.audio_preprocess import prepare_native_parakeet_audio_features from cactus.transpile.model_profiles import add_tensor_aliases from cactus.transpile.model_profiles import PARAKEET_TDT_PROFILE @@ -871,16 +872,28 @@ def build_parakeet_tdt_component_specs( "adapter_family": "parakeet_tdt", } + full_frames = int(input_features.shape[1]) + encoder_specs: list[ComponentModuleSpec] = [] + for frames in audio_bucket_frames(full_frames): + component = "audio_encoder" if frames == full_frames else f"audio_encoder_{frames}" + encoder_specs.append( + ComponentModuleSpec( + component=component, + module=model.encoder, + example_inputs=(input_features[:, :frames, :].contiguous(),), + input_keys=("input_features",), + output_keys=("encoder_hidden_states",), + graph_meta={**common_graph_meta, "component": component}, + metadata={ + "family": "parakeet_tdt", + "task": "tdt_transcription", + "audio_frames": str(frames), + }, + ) + ) + return [ - ComponentModuleSpec( - component="audio_encoder", - module=model.encoder, - example_inputs=(input_features,), - input_keys=("input_features",), - output_keys=("encoder_hidden_states",), - graph_meta={**common_graph_meta, "component": "audio_encoder"}, - metadata={"family": "parakeet_tdt", "task": "tdt_transcription"}, - ), + *encoder_specs, ComponentModuleSpec( component="decoder", module=model.decoder_step, From 1b16ec82b3d033651fea7382b9373329930e447d Mon Sep 17 00:00:00 2001 From: jakmro Date: Mon, 17 Aug 2026 15:11:35 +0200 Subject: [PATCH 2/3] refactor: streamline main test runner by removing redundant test cases --- cactus-engine/tests/test_llm.cpp | 6 ------ 1 file changed, 6 deletions(-) diff --git a/cactus-engine/tests/test_llm.cpp b/cactus-engine/tests/test_llm.cpp index 1bfab8069..4a56c690b 100644 --- a/cactus-engine/tests/test_llm.cpp +++ b/cactus-engine/tests/test_llm.cpp @@ -740,23 +740,17 @@ bool test_decode_batch_throughput() { int main() { TestUtils::TestRunner runner("LLM Tests"); - runner.run_test("1k_context", test_1k_context()); runner.run_test("streaming", test_streaming()); runner.run_test("prefill", test_prefill()); - runner.run_test("prefill_idempotent_reuse", test_prefill_idempotent_reuse()); - runner.run_test("prefill_prefix_extension_reuse", test_prefill_prefix_extension_reuse()); runner.run_test("prefill_invalidated_on_message_change", test_prefill_invalidated_on_message_change()); runner.run_test("chunked_prefill_padding", test_chunked_prefill_padding()); runner.run_test("tool_calls", test_tool_call()); runner.run_test("tool_multiple_tool_call_invocations", test_multiple_tool_call_invocations()); - runner.run_test("tool_calls_with_three_tools", test_tool_call_with_three_tools()); runner.run_test("tool_constraint_clear_releases_bias", test_tool_constraint_clear_releases_bias()); runner.run_test("partition_thinking_response", test_partition_thinking_response()); runner.run_test("prompt_retains_thinking", test_prompt_gemma4_retains_thinking()); runner.run_test("complete_thinking_api_clean", test_complete_gemma4_thinking_api_clean()); runner.run_test("multiturn_thinking_persist", test_multiturn_thinking_persist()); - runner.run_test("multiturn_turn2_distinct", test_multiturn_turn2_distinct()); - runner.run_test("decode_batch", test_decode_batch()); runner.run_test("generate_batch_ragged", test_generate_batch_ragged()); runner.run_test("decode_batch_throughput", test_decode_batch_throughput()); runner.run_test("batch_distinct4_matches_single", test_batch_distinct4_matches_single()); From 66d2e06a53530f711e5a1ad7cd816f2f4e1e56c0 Mon Sep 17 00:00:00 2001 From: jakmro Date: Mon, 17 Aug 2026 16:49:55 +0200 Subject: [PATCH 3/3] fix --- python/cactus/py.typed | 0 python/pyproject.toml | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) create mode 100644 python/cactus/py.typed diff --git a/python/cactus/py.typed b/python/cactus/py.typed new file mode 100644 index 000000000..e69de29bb diff --git a/python/pyproject.toml b/python/pyproject.toml index 7f3a893c1..4bc52d98c 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -73,7 +73,7 @@ exclude = ["cactus.code*"] [tool.setuptools.package-data] "cactus.bindings" = ["lib/*.dylib", "lib/*.so"] -"cactus" = ["bin/run", "bin/transcribe", "assets/*", "code/**/*"] +"cactus" = ["bin/run", "bin/transcribe", "assets/*", "code/**/*", "py.typed"] "cactus.convert" = ["assets/*/*"] [tool.pytest.ini_options]