From f095c5bc00c2a7fa9f80eb2559e47b7ea1094614 Mon Sep 17 00:00:00 2001 From: jakmro Date: Tue, 8 Sep 2026 12:04:03 +0200 Subject: [PATCH] Fix padded-frame leakage and feature mismatch in the Parakeet encoder The bucketed Parakeet encoder zero-padded normalized log-mel features up to the bucket length but the exported graph had no length mask, so padded frames leaked into real frames through attention and the depthwise conv, and the decoder also walked the padded tail. Short clips lost their last word and gained hallucinated words (MOB-12663). Separately, the engine built the mel filterbank with the HTK scale and no normalization while NeMo uses Slaney/Slaney, and the full encoder was traced to the transpiler's 20 s example asset rather than the configured audio limit. - The encoder takes a required second input, input_features_mask [batch, frames], like the Gemma4 audio encoder. The graph subsamples it by the subsampling factor and applies it the way NeMo does: an additive -10000 key mask on attention and a multiply after the GLU in the conv module. The conv-module fusion accepts that multiply so the fused lowering is kept. - The engine fills the mask up to the copied frame count, decodes only the valid encoder frames in both the one-shot and streaming paths, and rejects bundles without the input. - Parakeet features use the Slaney mel scale and norm; the 4-bit engine now scores the same WER as the official fp32 checkpoint on the CoreApp eval clips. - The Parakeet example is padded to CACTUS_TRANSPILER_MAX_AUDIO_SECONDS (30 s by default) and the default bucket ladder becomes 1, 2, 3, 5, 8, 13, 21 s plus the full window. - io.cpp: define MappedFileRegistry outside the nested GraphFile namespace so the tree compiles again. Claude-Session: https://claude.ai/code/session_01EZYnugf8VMdNDtnK8stdpU --- cactus-engine/src/model.cpp | 24 +++++--- cactus-engine/src/stream.cpp | 2 +- cactus-engine/src/transcribe.cpp | 4 +- cactus-graph/src/io.cpp | 4 -- .../convert/tests/test_audio_buckets.py | 7 ++- .../convert/tests/test_parakeet_capture.py | 60 +++++++++++++++++++ python/cactus/transpile/audio_preprocess.py | 2 +- python/cactus/transpile/fusion/conv.py | 18 ++++++ python/cactus/transpile/hf_model.py | 24 +++++++- python/cactus/transpile/lower.py | 8 +++ python/cactus/transpile/ops.py | 1 + python/cactus/transpile/optimize_graph.py | 3 + python/cactus/transpile/tdt_runtime.py | 40 ++++++++----- 13 files changed, 161 insertions(+), 36 deletions(-) create mode 100644 python/cactus/convert/tests/test_parakeet_capture.py 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={