diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 1f19daaff..2fc60d1e4 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -4757,6 +4757,21 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a write_typed_buffer(feat_buf, feat_desc.precision, transposed.data(), transposed.size() * sizeof(float), Precision::FP32); + const size_t subsampling = std::max(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(audio_enc->runtime_input_node_ids[mask_idx])); + std::vector 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 Model::transcribe_parakeet_tdt(const std::vector& a }; } - size_t commit_to = T; - if (stream) { - size_t valid_hidden = T; - if (expected_frames > 0) - valid_hidden = std::min(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) * diff --git a/cactus-engine/src/stream.cpp b/cactus-engine/src/stream.cpp index 182d81e4e..62e99a632 100644 --- a/cactus-engine/src/stream.cpp +++ b/cactus-engine/src/stream.cpp @@ -125,7 +125,7 @@ std::vector window_features(std::vector window, size_t mel_bins) { const size_t waveform_samples = window.size(); cactus::audio::apply_preemphasis(window, 0.97f); std::vector 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; diff --git a/cactus-engine/src/transcribe.cpp b/cactus-engine/src/transcribe.cpp index 90c4cf5ca..077dbe232 100644 --- a/cactus-engine/src/transcribe.cpp +++ b/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; diff --git a/cactus-graph/src/io.cpp b/cactus-graph/src/io.cpp index 664ab2677..ad4fbf5f0 100644 --- a/cactus-graph/src/io.cpp +++ b/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> MappedFileRegistry::cache_; @@ -931,8 +929,6 @@ void MappedFileRegistry::clear() { cache_.clear(); } -} // namespace GraphFile - // MappedFile implementation MappedFile::MappedFile(const std::string& filename) diff --git a/python/cactus/convert/tests/test_audio_buckets.py b/python/cactus/convert/tests/test_audio_buckets.py index 00b6f6bee..723ee03e5 100644 --- a/python/cactus/convert/tests/test_audio_buckets.py +++ b/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) diff --git a/python/cactus/convert/tests/test_parakeet_capture.py b/python/cactus/convert/tests/test_parakeet_capture.py new file mode 100644 index 000000000..cc55e919a --- /dev/null +++ b/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) diff --git a/python/cactus/transpile/audio_preprocess.py b/python/cactus/transpile/audio_preprocess.py index 6956a5038..805f6d940 100644 --- a/python/cactus/transpile/audio_preprocess.py +++ b/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: diff --git a/python/cactus/transpile/fusion/conv.py b/python/cactus/transpile/fusion/conv.py index 265501f9a..ec893c566 100644 --- a/python/cactus/transpile/fusion/conv.py +++ b/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, diff --git a/python/cactus/transpile/hf_model.py b/python/cactus/transpile/hf_model.py index 67280412c..d6c579485 100644 --- a/python/cactus/transpile/hf_model.py +++ b/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") diff --git a/python/cactus/transpile/lower.py b/python/cactus/transpile/lower.py index fab23d6ad..8896291d2 100644 --- a/python/cactus/transpile/lower.py +++ b/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, diff --git a/python/cactus/transpile/ops.py b/python/cactus/transpile/ops.py index 5e3ee5436..4dcdd243a 100644 --- a/python/cactus/transpile/ops.py +++ b/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", ), diff --git a/python/cactus/transpile/optimize_graph.py b/python/cactus/transpile/optimize_graph.py index e66d80881..8aefb82ab 100644 --- a/python/cactus/transpile/optimize_graph.py +++ b/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), } diff --git a/python/cactus/transpile/tdt_runtime.py b/python/cactus/transpile/tdt_runtime.py index 07c44b63c..2519552af 100644 --- a/python/cactus/transpile/tdt_runtime.py +++ b/python/cactus/transpile/tdt_runtime.py @@ -457,7 +457,7 @@ def __init__(self, config: ParakeetTDTConfig, prefix: str, state_dict: dict[str, self.pos_bias_v = nn.Parameter(state_dict[f"{prefix}.pos_bias_v"].clone()) self.config = config - def forward(self, x: torch.Tensor) -> torch.Tensor: + def forward(self, x: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor: batch, seq_len, _ = x.shape num_heads = self.config.attention_heads head_dim = self.config.attention_head_dim @@ -488,6 +488,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: rel_k_heads, scale=self.config.attention_scale, ) + rel_bias = rel_bias + attn_mask attn = F.scaled_dot_product_attention( q_u_heads, k_heads, @@ -551,12 +552,11 @@ def __init__(self, config: ParakeetTDTConfig, prefix: str, state_dict: dict[str, ) self.config = config - def forward(self, x: torch.Tensor, pad_mask: torch.Tensor | None = None) -> torch.Tensor: + def forward(self, x: torch.Tensor, pad_mask: torch.Tensor) -> torch.Tensor: x = x.transpose(1, 2) x = self.pointwise_conv1(x) x = F.glu(x, dim=1) - if pad_mask is not None: - x = x * pad_mask + x = x * pad_mask x = self.depthwise_conv(x) x = self.batch_norm(x) x = _apply_activation(x, self.config.encoder_hidden_act) @@ -583,10 +583,10 @@ def __init__(self, config: ParakeetTDTConfig, layer_index: int, state_dict: dict module.bias.data.copy_(state_dict[f"{prefix}.{name}.bias"].to(dtype=module.bias.dtype)) self.config = config - def forward(self, x: torch.Tensor) -> torch.Tensor: + def forward(self, x: torch.Tensor, conv_mask: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor: x = x + 0.5 * self.feed_forward1(self.norm_feed_forward1(x), activation=self.config.encoder_hidden_act) - x = x + self.self_attn(self.norm_self_att(x)) - x = x + self.conv(self.norm_conv(x)) + x = x + self.self_attn(self.norm_self_att(x), attn_mask) + x = x + self.conv(self.norm_conv(x), conv_mask) x = x + 0.5 * self.feed_forward2(self.norm_feed_forward2(x), activation=self.config.encoder_hidden_act) return self.norm_out(x) @@ -639,11 +639,15 @@ def __init__(self, config: ParakeetTDTConfig, state_dict: dict[str, torch.Tensor self.layers = nn.ModuleList( [ParakeetTDTEncoderLayer(config, index, state_dict) for index in range(config.num_layers)] ) + self.config = config - def forward(self, input_features: torch.Tensor) -> torch.Tensor: + def forward(self, input_features: torch.Tensor, input_features_mask: torch.Tensor) -> torch.Tensor: x = self.pre_encode(input_features) + pad_mask = input_features_mask[:, :: self.config.subsampling_factor] + conv_mask = pad_mask.unsqueeze(1) + attn_mask = ((pad_mask - 1.0) * 10000.0).unsqueeze(1).unsqueeze(1) for layer in self.layers: - x = layer(x) + x = layer(x, conv_mask, attn_mask) return x @@ -766,7 +770,7 @@ def __init__(self, config: ParakeetTDTConfig, state_dict: dict[str, torch.Tensor self.decoder_step = ParakeetTDTDecoderStep(config, state_dict) def forward(self, input_features: torch.Tensor) -> torch.Tensor: - return self.encoder(input_features) + return self.encoder(input_features, torch.ones_like(input_features[:, :, 0])) def initial_decoder_state(self, *, batch_size: int, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, ...]: state: list[torch.Tensor] = [] @@ -777,7 +781,7 @@ def initial_decoder_state(self, *, batch_size: int, device: torch.device, dtype: def greedy_decode_token_ids(self, input_features: torch.Tensor) -> list[int]: with torch.no_grad(): - encoder_hidden = self.encoder(input_features) + encoder_hidden = self.encoder(input_features, torch.ones_like(input_features[:, :, 0])) batch = int(encoder_hidden.shape[0]) if batch != 1: raise ValueError("Parakeet TDT local greedy decode currently expects batch size 1") @@ -844,8 +848,13 @@ def build_parakeet_tdt_component_specs( named_tensors: dict[str, torch.Tensor], weights_dir: str | None = None, ) -> list[ComponentModuleSpec]: - input_features = named_tensors["input_features"] - example_hidden = model.encoder(input_features) + features = named_tensors["input_features"] + input_features = torch.randn( + (int(features.shape[0]), 3000, int(features.shape[2])), + device=features.device, + dtype=features.dtype, + ) + example_hidden = model.encoder(input_features, torch.ones_like(input_features[:, :, 0])) batch_size = int(example_hidden.shape[0]) initial_states = model.initial_decoder_state( batch_size=batch_size, @@ -876,12 +885,13 @@ def build_parakeet_tdt_component_specs( encoder_specs: list[ComponentModuleSpec] = [] for frames in audio_bucket_frames(full_frames): component = "audio_encoder" if frames == full_frames else f"audio_encoder_{frames}" + bucket_features = input_features[:, :frames, :].contiguous() encoder_specs.append( ComponentModuleSpec( component=component, module=model.encoder, - example_inputs=(input_features[:, :frames, :].contiguous(),), - input_keys=("input_features",), + example_inputs=(bucket_features, torch.ones_like(bucket_features[:, :, 0])), + input_keys=("input_features", "input_features_mask"), output_keys=("encoder_hidden_states",), graph_meta={**common_graph_meta, "component": component}, metadata={