diff --git a/python/cactus/convert/tests/test_audio_buckets.py b/python/cactus/convert/tests/test_audio_buckets.py index 723ee03e5..e08b3876c 100644 --- a/python/cactus/convert/tests/test_audio_buckets.py +++ b/python/cactus/convert/tests/test_audio_buckets.py @@ -1,5 +1,6 @@ from cactus.transpile.audio_preprocess import audio_bucket_frames from cactus.transpile.audio_preprocess import audio_bucket_seconds +from cactus.transpile.audio_preprocess import audio_capture_frames def test_fibonacci_ladder_capped_at_model_window(): @@ -26,8 +27,18 @@ def test_buckets_are_sorted_and_unique(): assert len(buckets) == len(set(buckets)) -def test_default_ladder_is_fibonacci_seconds(): - assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13, 21) +def test_default_ladder_seconds(): + assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 300, 377) + + +def test_capture_window_is_420_seconds(): + assert audio_capture_frames() == 42000 + + +def test_env_override_keeps_capture_window(monkeypatch): + monkeypatch.setenv("CACTUS_TRANSPILER_AUDIO_BUCKETS", "30,60") + assert audio_capture_frames() == 42000 + assert audio_bucket_frames(audio_capture_frames()) == [3000, 6000, 42000] def test_env_override_replaces_ladder(monkeypatch): @@ -43,9 +54,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, 21) + assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 300, 377) 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, 21) + assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 300, 377) diff --git a/python/cactus/convert/tests/test_parakeet_capture.py b/python/cactus/convert/tests/test_parakeet_capture.py index cc55e919a..eac8ea0cd 100644 --- a/python/cactus/convert/tests/test_parakeet_capture.py +++ b/python/cactus/convert/tests/test_parakeet_capture.py @@ -19,9 +19,11 @@ def test_capture_window_is_independent_of_validation_audio(): 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 [int(spec.metadata["audio_frames"]) for spec in encoders] == [ + 100, 200, 300, 500, 800, 1300, 2100, 3400, 5500, 8900, 14400, 23300, 30000, 37700, 42000, + ] assert encoders[-1].component == "audio_encoder" - assert tuple(encoders[-1].example_inputs[0].shape) == (1, 3000, 128) + assert tuple(encoders[-1].example_inputs[0].shape) == (1, 42000, 128) assert torch.all(features == 2.0) assert not torch.all(encoders[-1].example_inputs[0] == 2.0) diff --git a/python/cactus/transpile/audio_preprocess.py b/python/cactus/transpile/audio_preprocess.py index 805f6d940..8cf039ad4 100644 --- a/python/cactus/transpile/audio_preprocess.py +++ b/python/cactus/transpile/audio_preprocess.py @@ -21,7 +21,8 @@ _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, 21) +_DEFAULT_AUDIO_BUCKET_SECONDS = (1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 300, 377) +_AUDIO_CAPTURE_SECONDS = 420 def audio_duration_limit_seconds() -> float: @@ -47,6 +48,10 @@ def audio_bucket_seconds() -> tuple[int, ...]: return positive +def audio_capture_frames() -> int: + return _AUDIO_CAPTURE_SECONDS * (_PARAKEET_SAMPLE_RATE // _PARAKEET_HOP_LENGTH) + + def audio_bucket_frames(max_frames: int) -> list[int]: if max_frames <= 0: return [] diff --git a/python/cactus/transpile/tdt_runtime.py b/python/cactus/transpile/tdt_runtime.py index 2519552af..70b02d485 100644 --- a/python/cactus/transpile/tdt_runtime.py +++ b/python/cactus/transpile/tdt_runtime.py @@ -14,6 +14,7 @@ from cactus.transpile.component_pipeline import ComponentModuleSpec from cactus.transpile.audio_preprocess import audio_bucket_frames +from cactus.transpile.audio_preprocess import audio_capture_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 @@ -850,11 +851,12 @@ def build_parakeet_tdt_component_specs( ) -> list[ComponentModuleSpec]: features = named_tensors["input_features"] input_features = torch.randn( - (int(features.shape[0]), 3000, int(features.shape[2])), + (int(features.shape[0]), audio_capture_frames(), int(features.shape[2])), device=features.device, dtype=features.dtype, ) - example_hidden = model.encoder(input_features, torch.ones_like(input_features[:, :, 0])) + example_features = input_features[:, : audio_bucket_frames(int(input_features.shape[1]))[0], :] + example_hidden = model.encoder(example_features, torch.ones_like(example_features[:, :, 0])) batch_size = int(example_hidden.shape[0]) initial_states = model.initial_decoder_state( batch_size=batch_size, diff --git a/python/cactus/transpiler/Converter/input_processor.py b/python/cactus/transpiler/Converter/input_processor.py index 617b5d42f..36eb4b75f 100644 --- a/python/cactus/transpiler/Converter/input_processor.py +++ b/python/cactus/transpiler/Converter/input_processor.py @@ -95,6 +95,7 @@ def parakeet_processor(model_id: str, configs: dict[str, dict[str, Any]], model_ "google/gemma-4-E2B": gemma4_processor, "openai/whisper-tiny": default_processor, "nvidia/parakeet-tdt-0.6b-v3": parakeet_processor, + "nvidia/parakeet-tdt-0.6b-v2": parakeet_processor, "LiquidAI/LFM2-VL-3B": default_processor, "Qwen/Qwen2.5-0.5B": default_processor, "LiquidAI/LFM2.5-8B-A1B": text_tokenizer_processor, diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index 691b2a653..8ee4fd98d 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -586,6 +586,7 @@ def decoder_kv_state(producer: str = "decoder_prefill_chunk") -> StateContract: "openai/whisper-tiny": WHISPER_PROFILE, "openai/whisper-small": WHISPER_PROFILE, "nvidia/parakeet-tdt-0.6b-v3": PARAKEET_PROFILE, + "nvidia/parakeet-tdt-0.6b-v2": PARAKEET_PROFILE, "LiquidAI/LFM2-VL-450M": LFM_VLM_PROFILE, "LiquidAI/LFM2-VL-3B": LFM_VLM_PROFILE, "Qwen/Qwen2.5-0.5B": QWEN2_5_0_5B_PROFILE,