Expand Parakeet buckets

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)

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)
 

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 []

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,

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,

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,