diff --git a/README.md b/README.md index 85a2742f9..0a83d15cd 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,7 @@ multi-component export for pipelines. | **Multimodal** | Gemma 3/4, Phi-4MM (vision + audio + LoRA), Nemotron Parse, LLaVA, InternVL2, Mage-VL (image + streaming video), MiniCPM-V 4.6, Qwen2.5-VL, Qwen3-VL, Qwen3.5/3.6-VL, Pixtral | | **Encoder-only** | BERT, RoBERTa, ALBERT, DeBERTa, DistilBERT, ELECTRA, XLNet | | **Encoder-Decoder** | BART, T5/mT5, Marian, M2M-100, Pegasus, BigBird-Pegasus | -| **Speech-to-Text** | Whisper, Moonshine, FastConformer-RNNT, FunASR, GLM-ASR, Qwen3-ASR, SenseVoice | +| **Speech-to-Text** | Whisper, Moonshine, Moonshine Streaming, FastConformer-RNNT, FunASR, GLM-ASR, Qwen3-ASR, SenseVoice | | **Audio** | Wav2Vec2, HuBERT, WavLM, SpeechT5 | | **Vision** | ViT, BEiT, DeiT, DINOv2, Swin, CLIP, SigLIP | | **Diffusion** | Stable Diffusion (UNet + VAE + ControlNet), Flux, SD3, DiT, QwenImage / Qwen-Image-Edit-2509, HunyuanDiT, CogVideoX | diff --git a/scripts/generate_golden.py b/scripts/generate_golden.py index e2cda9270..50bbdf10d 100644 --- a/scripts/generate_golden.py +++ b/scripts/generate_golden.py @@ -680,6 +680,39 @@ def _generate_image_to_text(case: TestCase, json_path: Path, device: str) -> Non ) +@contextlib.contextmanager +def non_mutating_encoder_context(model: object): + """Stop a seq2seq decoder from mutating ``encoder_hidden_states`` in place. + + Some speech decoders add an absolute position table to the encoder output + with ``+=`` (Moonshine Streaming does). Under cached cross-attention the + mutated tensor is never read again, so generation is unaffected — but a + golden reference must not *depend* on that. Cloning the argument on entry + makes the reference provably free of encoder-context accumulation, so the + committed golden describes "encoder context is added exactly once", which + is the semantics the exported ONNX decoder implements. + + A no-op for models that expose no ``model.decoder`` or never mutate. + """ + decoder = getattr(getattr(model, "model", None), "decoder", None) + forward = getattr(decoder, "forward", None) + if forward is None: + yield + return + + def guarded(*args, **kwargs): + states = kwargs.get("encoder_hidden_states") + if states is not None and hasattr(states, "clone"): + kwargs["encoder_hidden_states"] = states.clone() + return forward(*args, **kwargs) + + decoder.forward = guarded + try: + yield + finally: + decoder.forward = forward + + def _generate_speech_to_text(case: TestCase, json_path: Path, device: str) -> None: """Generate golden data for an encoder-decoder speech recognition model.""" import librosa @@ -690,12 +723,15 @@ def _generate_speech_to_text(case: TestCase, json_path: Path, device: str) -> No model = transformers.AutoModelForSpeechSeq2Seq.from_pretrained( case.model_id, + revision=case.revision, device_map=device, trust_remote_code=case.trust_remote_code, ) model.eval() processor = transformers.AutoProcessor.from_pretrained( - case.model_id, trust_remote_code=case.trust_remote_code + case.model_id, + revision=case.revision, + trust_remote_code=case.trust_remote_code, ) # Load and preprocess audio @@ -712,7 +748,7 @@ def _generate_speech_to_text(case: TestCase, json_path: Path, device: str) -> No decoder_input_ids = torch.tensor( [[decoder_start_id]], dtype=torch.long, device=model_device ) - with torch.no_grad(): + with torch.no_grad(), non_mutating_encoder_context(model): outputs = model( **processed, decoder_input_ids=decoder_input_ids, @@ -724,7 +760,7 @@ def _generate_speech_to_text(case: TestCase, json_path: Path, device: str) -> No # L5: greedy generation generated_ids = None if "L5" in case.level: - with torch.no_grad(): + with torch.no_grad(), non_mutating_encoder_context(model): gen = model.generate( **processed, max_new_tokens=case.generation_params.get("max_new_tokens", 50), diff --git a/src/mobius/__init__.py b/src/mobius/__init__.py index 781e1d084..bffa8a4bc 100644 --- a/src/mobius/__init__.py +++ b/src/mobius/__init__.py @@ -35,6 +35,7 @@ "MambaConfig", "MllamaConfig", "MoonshineConfig", + "MoonshineStreamingConfig", "ModelPackage", "ModelRegistration", "ModelRegistry", @@ -99,6 +100,7 @@ MllamaConfig, MMSConfig, MoonshineConfig, + MoonshineStreamingConfig, Sam2Config, SegformerConfig, SpeechToTextConfig, diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py index 3fd858ac0..09b3f5cd1 100644 --- a/src/mobius/_builder.py +++ b/src/mobius/_builder.py @@ -144,7 +144,7 @@ def build_from_module( module: An ``onnxscript.nn.Module`` whose ``forward`` signature is compatible with *task*. config: Architecture configuration. Its ``dtype`` controls model - precision and its optional ``validate`` method runs before build. + precision and its ``validate`` method runs before build. task: Task name or :class:`ModelTask` instance. execution_provider: Target for EP-aware optimizations. trace_optimization: Log optimization diagnostics when true. diff --git a/src/mobius/_configs/__init__.py b/src/mobius/_configs/__init__.py index 960751059..f80779bd4 100644 --- a/src/mobius/_configs/__init__.py +++ b/src/mobius/_configs/__init__.py @@ -59,6 +59,7 @@ MllamaConfig, MMSConfig, MoonshineConfig, + MoonshineStreamingConfig, MuseGlimmerConfig, NanoChatConfig, NemotronHConfig, @@ -149,6 +150,7 @@ "MiniMaxConfig", "MMSConfig", "MoonshineConfig", + "MoonshineStreamingConfig", "MuseGlimmerConfig", "NanoChatConfig", "NemotronParseConfig", diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index a1cc6e1f9..f2ad90233 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -4489,6 +4489,16 @@ class SpeechToTextConfig(ArchitectureConfig): decoder_start_token_id: int | None = None layer_norm_eps: float = 1e-5 + @property + def encoder_output_size(self) -> int: + """Channel width of ``encoder_hidden_states``. + + Defaults to the decoder width because most speech encoder-decoders share + one model dimension. Architectures whose encoder is narrower or wider + (e.g. Moonshine Streaming with a projection adapter) override this. + """ + return self.hidden_size + @dataclasses.dataclass class WhisperConfig(SpeechToTextConfig): @@ -4635,6 +4645,173 @@ def from_transformers(cls, config, parent_config=None) -> MoonshineConfig: return cls(**options) +def _sub_config_get(config, name: str, default): + """Read ``name`` from a sub-config that may be an object or a plain dict.""" + if isinstance(config, dict): + value = config.get(name, default) + else: + value = getattr(config, name, default) + return default if value is None else value + + +@dataclasses.dataclass +class MoonshineStreamingConfig(SpeechToTextConfig): + """Configuration for Moonshine Streaming raw-waveform encoder-decoder ASR. + + Moonshine Streaming replaces the offline Moonshine convolutional stem with a + fixed-length framing front end (``frame_ms`` frames of raw samples), per-frame + CMVN, learned asinh compression, and two *causal* strided convolutions. Its + encoder carries no rotary embedding; position information comes from the + causal stem plus per-layer asymmetric ``(left, right)`` sliding windows, which + bound the streaming lookahead. The decoder adds an absolute learned position + table (``pos_emb``) to the encoder output before cross-attention. + """ + + encoder_input_name: str = "input_values" + encoder_input_channels: int | None = None + encoder_uses_attention_mask: bool = True + decoder_uses_encoder_attention_mask: bool = True + encoder_hidden_size: int = DEFAULT_INT + encoder_intermediate_size: int = DEFAULT_INT + encoder_num_hidden_layers: int = DEFAULT_INT + encoder_num_attention_heads: int = DEFAULT_INT + encoder_num_key_value_heads: int = DEFAULT_INT + encoder_head_dim: int = DEFAULT_INT + encoder_hidden_act: str = "gelu" + decoder_hidden_act: str = "silu" + #: Q/K/V and output projection bias of the encoder attention. Upstream + #: gates all four encoder projections on the encoder sub-config's + #: ``attention_bias`` (the decoder's output projection stays bias-free). + encoder_attention_bias: bool = False + #: Per-layer ``(left_window, right_window)`` attention spans of the encoder. + #: ``left`` counts the query position itself; ``right`` is the strict + #: lookahead. ``right == 0`` means the layer is fully causal. + encoder_sliding_windows: tuple[tuple[int, int], ...] = ((16, 4),) + encoder_sample_rate: int = 16_000 + encoder_frame_ms: float = 5.0 + #: Epsilon of the per-frame cepstral mean/variance normalisation. + encoder_cmvn_eps: float = 1e-6 + + def __post_init__(self): + # A tuple keeps the config hashable and prevents accidental mutation of + # the per-layer window schedule shared by every encoder layer. + self.encoder_sliding_windows = tuple( + (int(left), int(right)) for left, right in self.encoder_sliding_windows + ) + if self.encoder_hidden_size == DEFAULT_INT: + self.encoder_hidden_size = self.hidden_size + if self.encoder_intermediate_size == DEFAULT_INT: + self.encoder_intermediate_size = self.intermediate_size + if self.encoder_num_attention_heads == DEFAULT_INT: + self.encoder_num_attention_heads = self.num_attention_heads + if self.encoder_num_key_value_heads == DEFAULT_INT: + self.encoder_num_key_value_heads = self.encoder_num_attention_heads + if self.encoder_num_hidden_layers == DEFAULT_INT: + self.encoder_num_hidden_layers = len(self.encoder_sliding_windows) + if self.encoder_head_dim == DEFAULT_INT: + self.encoder_head_dim = ( + self.encoder_hidden_size // self.encoder_num_attention_heads + ) + if len(self.encoder_sliding_windows) != self.encoder_num_hidden_layers: + raise ValueError( + "MoonshineStreamingConfig: encoder_sliding_windows has " + f"{len(self.encoder_sliding_windows)} entries but the encoder has " + f"{self.encoder_num_hidden_layers} layers." + ) + + @property + def encoder_output_size(self) -> int: + """Encoder width; the decoder projects it when it differs from its own.""" + return self.encoder_hidden_size + + @property + def frame_length(self) -> int: + """Raw samples per encoder frame (``sample_rate * frame_ms / 1000``). + + Upstream rounds to the nearest integer, so the audio length fed to the + encoder must be a multiple of this value; the processor's + ``pad_to_multiple_of`` enforces that. + """ + return round(self.encoder_sample_rate * self.encoder_frame_ms / 1000.0) + + @classmethod + def from_transformers(cls, config, parent_config=None) -> MoonshineStreamingConfig: + if config.model_type != "moonshine_streaming": + raise ValueError( + "MoonshineStreamingConfig expects model_type='moonshine_streaming', " + f"got '{config.model_type}'" + ) + + hidden_size = config.hidden_size + decoder_heads = config.num_attention_heads + # ``encoder_config`` is a nested MoonshineStreamingEncoderConfig on a + # trusted config object and a plain dict when the JSON is read directly. + encoder_config = getattr(config, "encoder_config", None) or {} + encoder_hidden_size = _sub_config_get(encoder_config, "hidden_size", hidden_size) + encoder_heads = _sub_config_get(encoder_config, "num_attention_heads", decoder_heads) + rope_parameters = getattr(config, "rope_parameters", None) or getattr( + config, "rope_scaling", None + ) + rope_parameters = rope_parameters or {} + windows = _sub_config_get( + encoder_config, + "sliding_windows", + ((16, 4), (16, 4), (16, 0), (16, 0), (16, 4), (16, 4)), + ) + + options = dict( + vocab_size=config.vocab_size, + hidden_size=hidden_size, + intermediate_size=config.intermediate_size, + num_hidden_layers=config.num_hidden_layers, + num_attention_heads=decoder_heads, + num_key_value_heads=getattr(config, "num_key_value_heads", decoder_heads), + head_dim=getattr(config, "head_dim", None) or hidden_size // decoder_heads, + hidden_act=getattr(config, "hidden_act", "silu"), + pad_token_id=getattr(config, "pad_token_id", 0), + tie_word_embeddings=getattr(config, "tie_word_embeddings", False), + attn_qkv_bias=getattr(config, "attention_bias", False), + attn_o_bias=False, + max_position_embeddings=getattr(config, "max_position_embeddings", 4096), + rope_type=rope_parameters.get("rope_type", "default"), + rope_theta=rope_parameters.get("rope_theta", 10_000.0), + rope_scaling=rope_parameters or None, + partial_rotary_factor=rope_parameters.get("partial_rotary_factor", 1.0), + rope_interleave=True, + mlp_bias=True, + encoder_hidden_size=encoder_hidden_size, + encoder_intermediate_size=_sub_config_get( + encoder_config, "intermediate_size", config.intermediate_size + ), + encoder_num_hidden_layers=_sub_config_get( + encoder_config, "num_hidden_layers", config.num_hidden_layers + ), + encoder_num_attention_heads=encoder_heads, + encoder_num_key_value_heads=_sub_config_get( + encoder_config, "num_key_value_heads", encoder_heads + ), + encoder_head_dim=_sub_config_get(encoder_config, "head_dim", None) + or encoder_hidden_size // encoder_heads, + encoder_hidden_act=_sub_config_get(encoder_config, "hidden_act", "gelu"), + encoder_attention_bias=bool( + _sub_config_get(encoder_config, "attention_bias", False) + ), + decoder_hidden_act=getattr(config, "hidden_act", "silu"), + encoder_sliding_windows=tuple(tuple(window) for window in windows), + encoder_sample_rate=_sub_config_get(encoder_config, "sample_rate", 16_000), + encoder_frame_ms=_sub_config_get(encoder_config, "frame_ms", 5.0), + decoder_start_token_id=getattr(config, "decoder_start_token_id", 1), + layer_norm_eps=getattr(config, "layer_norm_eps", 1e-5), + model_type="moonshine_streaming", + bos_token_id=getattr(config, "bos_token_id", 1), + eos_token_id=getattr(config, "eos_token_id", 2), + ) + resolved = _resolve_dtype(config) + if resolved is not None: + options["dtype"] = resolved + return cls(**options) + + def _conv_widths(config, defaults, hidden_size: int) -> tuple[int, ...]: """Per-layer channel widths of a wav2vec2-family convolutional feature encoder. diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index 51e0e5d4e..767085da7 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -40,6 +40,7 @@ MiniMaxConfig, MMSConfig, MoonshineConfig, + MoonshineStreamingConfig, MuseGlimmerConfig, NemotronParseConfig, ParakeetCTCConfig, @@ -115,6 +116,7 @@ Mistral4GGUFCausalLMModel, MoECausalLMModel, MoonshineForConditionalGeneration, + MoonshineStreamingForConditionalGeneration, NanoChatCausalLMModel, NemotronCausalLMModel, NemotronParseForConditionalGeneration, @@ -902,6 +904,11 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: task="speech-to-text", config_class=MoonshineConfig, ), + "moonshine_streaming": ModelRegistration( + MoonshineStreamingForConditionalGeneration, + task="speech-to-text", + config_class=MoonshineStreamingConfig, + ), # --- Encoder-only --- "albert": ModelRegistration(BertModel, task="feature-extraction"), "bert": ModelRegistration(BertModel, task="feature-extraction"), @@ -1361,6 +1368,7 @@ def _create_default_registry() -> ModelRegistry: # --- Speech --- "moonshine": "moonshine-ai/moonshine-tiny", + "moonshine_streaming": "moonshine-ai/moonshine-streaming-tiny", "whisper": "openai/whisper-tiny", "qwen3_asr": "Qwen/Qwen3-ASR-0.6B", "fun_asr": "justinchuby/Fun-ASR-Nano-2512", diff --git a/src/mobius/components/_whisper.py b/src/mobius/components/_whisper.py index 4c058d9ac..218d12ce0 100644 --- a/src/mobius/components/_whisper.py +++ b/src/mobius/components/_whisper.py @@ -23,7 +23,12 @@ class Conv1d(nn.Module): - """1D convolution layer.""" + """1D convolution layer. + + ``padding`` is symmetric when given as an ``int``. Pass a + ``(left, right)`` pair for asymmetric padding — ``(kernel_size - 1, 0)`` + makes the convolution causal, as used by streaming audio front ends. + """ def __init__( self, @@ -31,7 +36,7 @@ def __init__( out_channels: int, kernel_size: int, stride: int = 1, - padding: int = 0, + padding: int | tuple[int, int] = 0, bias: bool = True, groups: int = 1, ): @@ -40,7 +45,11 @@ def __init__( self.bias = nn.Parameter([out_channels]) if bias else None self._kernel_shape = [kernel_size] self._strides = [stride] - self._pads = [padding, padding] + if isinstance(padding, int): + self._pads = [padding, padding] + else: + left, right = padding + self._pads = [int(left), int(right)] self._groups = groups def forward(self, op: OpBuilder, x: ir.Value): diff --git a/src/mobius/components/_whisper_test.py b/src/mobius/components/_whisper_test.py index b8f6d5284..9016a2d4e 100644 --- a/src/mobius/components/_whisper_test.py +++ b/src/mobius/components/_whisper_test.py @@ -41,6 +41,16 @@ def test_positional_bias_and_grouped_weight_shape(self): assert biasless.bias is None assert list(grouped.weight.shape) == [4, 1, 3] + def test_symmetric_padding_from_int(self): + conv = Conv1d(8, 8, kernel_size=3, padding=1) + assert conv._pads == [1, 1] + + def test_asymmetric_padding_from_pair(self): + """A ``(left, right)`` pair makes the convolution causal.""" + conv = Conv1d(8, 8, kernel_size=5, stride=2, padding=(4, 0)) + assert conv._pads == [4, 0] + assert conv._strides == [2] + def test_forward_builds_graph(self): conv = Conv1d(80, 512, kernel_size=3, padding=1) builder, op, graph = create_test_builder() diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 44ed105ab..fc950e735 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -1913,11 +1913,6 @@ def _runtime_capability_warnings(pkg: ModelPackage) -> tuple[str, ...]: "onnxruntime-genai 0.15.2 does not support Mage-VL's patch_positions " "vision input and 1D decoder position_ids contract." ) - if getattr(config, "model_type", None) == "moonshine": - warnings.append( - "onnxruntime-genai 0.15.2 does not support Moonshine's variable-length " - "raw-waveform encoder." - ) decoder_key = "decoder" if "decoder" in pkg else "model" decoder_model = pkg.get(decoder_key) if decoder_model is not None: diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 8f55bd9ca..4661c3ff8 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -95,14 +95,6 @@ def _mock_decoder_model( return _mock_model(inputs=inputs, outputs=outputs) -def test_moonshine_runtime_limitation_is_advisory(tmp_path): - result = write_ort_genai_config(_make_fake_llm_pkg("moonshine"), str(tmp_path)) - - compatibility = json.loads(Path(result["runtime_compatibility"]).read_text()) - assert compatibility["runtime_validation_status"] == "unsupported-by-tested-runtime" - assert "variable-length raw-waveform encoder" in compatibility["warnings"][0] - - def _make_fake_llm_pkg(model_type: str = "qwen2"): """Build a minimal LLM-only ModelPackage with a fake config.""" import dataclasses diff --git a/src/mobius/integrations/transformers/_config_resolver_test.py b/src/mobius/integrations/transformers/_config_resolver_test.py index 84bc1e283..c583a6923 100644 --- a/src/mobius/integrations/transformers/_config_resolver_test.py +++ b/src/mobius/integrations/transformers/_config_resolver_test.py @@ -12,6 +12,7 @@ from mobius._configs import ( ArchitectureConfig, MoonshineConfig, + MoonshineStreamingConfig, QuantizationConfig, WhisperConfig, ) @@ -546,6 +547,94 @@ def test_moonshine_default_task(self): assert _default_task_for_model("moonshine") == "speech-to-text" +class TestMoonshineStreamingEncoderDecoder: + """Moonshine Streaming extraction keeps its encoder sub-config semantics.""" + + def _hf_config(self, encoder_config): + return type( + "FakeMoonshineStreamingConfig", + (), + { + "model_type": "moonshine_streaming", + "encoder_config": encoder_config, + "vocab_size": 32768, + "hidden_size": 320, + "intermediate_size": 1280, + "num_hidden_layers": 6, + "num_attention_heads": 8, + "num_key_value_heads": 8, + "head_dim": 40, + "hidden_act": "silu", + "max_position_embeddings": 4096, + "rope_parameters": { + "rope_type": "default", + "rope_theta": 10_000.0, + "partial_rotary_factor": 0.8, + }, + "attention_bias": False, + "pad_token_id": 0, + "bos_token_id": 1, + "eos_token_id": 2, + "decoder_start_token_id": 1, + "tie_word_embeddings": False, + }, + )() + + def _encoder_dict(self): + return { + "hidden_size": 320, + "intermediate_size": 1280, + "num_hidden_layers": 6, + "num_attention_heads": 8, + "num_key_value_heads": 8, + "head_dim": 40, + "hidden_act": "gelu", + "sample_rate": 16000, + "frame_ms": 5.0, + "sliding_windows": [[16, 4], [16, 4], [16, 0], [16, 0], [16, 4], [16, 4]], + } + + def test_routes_to_moonshine_streaming_config(self): + result = _config_from_hf(self._hf_config(self._encoder_dict())) + assert isinstance(result, MoonshineStreamingConfig) + assert result.encoder_input_name == "input_values" + assert result.encoder_uses_attention_mask is True + assert result.decoder_uses_encoder_attention_mask is True + assert result.tie_word_embeddings is False + + def test_encoder_sub_config_dict_and_object_agree(self): + """A dict sub-config and an attribute-style sub-config extract alike.""" + from_dict = _config_from_hf(self._hf_config(self._encoder_dict())) + encoder_object = type("FakeEncoderConfig", (), self._encoder_dict())() + from_object = _config_from_hf(self._hf_config(encoder_object)) + assert from_dict == from_object + + def test_streaming_specific_fields(self): + result = _config_from_hf(self._hf_config(self._encoder_dict())) + assert result.encoder_sliding_windows == ( + (16, 4), + (16, 4), + (16, 0), + (16, 0), + (16, 4), + (16, 4), + ) + assert result.encoder_hidden_act == "gelu" + assert result.decoder_hidden_act == "silu" + assert result.encoder_head_dim == 40 + assert result.encoder_sample_rate == 16000 + assert result.encoder_frame_ms == pytest.approx(5.0) + assert result.frame_length == 80 + assert result.partial_rotary_factor == pytest.approx(0.8) + assert result.rope_interleave is True + assert result.mlp_bias is True + assert result.attn_qkv_bias is False + assert result.attn_o_bias is False + + def test_default_task(self): + assert _default_task_for_model("moonshine_streaming") == "speech-to-text" + + # ── _dict_to_pretrained_config ────────────────────────────────────────── diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 40dfb2cda..a68ed2f5d 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -123,6 +123,7 @@ "MiniCPMCausalLMModel", "MiniCPMV46ForConditionalGeneration", "MoonshineForConditionalGeneration", + "MoonshineStreamingForConditionalGeneration", "MuseGlimmerForConditionalGeneration", "MuseGlimmerTextCausalLMModel", "MimiModel", @@ -352,6 +353,9 @@ Qwen2MoECausalLMModel, ) from mobius.models.moonshine import MoonshineForConditionalGeneration +from mobius.models.moonshine_streaming import ( + MoonshineStreamingForConditionalGeneration, +) from mobius.models.moshi import ( MoshiDepformerModel, MoshiTemporalModel, diff --git a/src/mobius/models/moonshine.py b/src/mobius/models/moonshine.py index 82e2b90f9..4cfde7479 100644 --- a/src/mobius/models/moonshine.py +++ b/src/mobius/models/moonshine.py @@ -74,7 +74,17 @@ def forward(self, op: OpBuilder, position_ids: ir.Value) -> tuple[ir.Value, ir.V class MoonshineAttention(nn.Module): - """Bias-free self/cross attention with optional partial interleaved RoPE.""" + """Self/cross attention with optional partial interleaved RoPE. + + ``hidden_size`` and ``head_dim`` default to the decoder-side values derived + from *config*. Encoder stacks whose width differs from the decoder (Moonshine + Streaming) pass them explicitly. + + Hugging Face bias-gates the Q/K/V projections on ``config.attention_bias`` + and keeps the decoder's output projection bias-free unconditionally, so the + two are separate switches here. Both default to ``False``, which is what + every published Moonshine checkpoint uses. + """ def __init__( self, @@ -83,22 +93,23 @@ def __init__( num_heads: int, num_key_value_heads: int, is_causal: bool = False, + hidden_size: int | None = None, + head_dim: int | None = None, + qkv_bias: bool = False, + o_bias: bool = False, ): super().__init__() + hidden_size = config.hidden_size if hidden_size is None else hidden_size self._num_heads = num_heads self._num_key_value_heads = num_key_value_heads - self._head_dim = config.hidden_size // num_heads + self._head_dim = hidden_size // num_heads if head_dim is None else head_dim self._rotary_dim = int(self._head_dim * config.partial_rotary_factor) self._scale = self._head_dim**-0.5 self._is_causal = is_causal - self.q_proj = Linear(config.hidden_size, num_heads * self._head_dim, bias=False) - self.k_proj = Linear( - config.hidden_size, num_key_value_heads * self._head_dim, bias=False - ) - self.v_proj = Linear( - config.hidden_size, num_key_value_heads * self._head_dim, bias=False - ) - self.o_proj = Linear(num_heads * self._head_dim, config.hidden_size, bias=False) + self.q_proj = Linear(hidden_size, num_heads * self._head_dim, bias=qkv_bias) + self.k_proj = Linear(hidden_size, num_key_value_heads * self._head_dim, bias=qkv_bias) + self.v_proj = Linear(hidden_size, num_key_value_heads * self._head_dim, bias=qkv_bias) + self.o_proj = Linear(num_heads * self._head_dim, hidden_size, bias=o_bias) def forward( self, @@ -199,6 +210,7 @@ def __init__(self, config: MoonshineConfig): config, num_heads=config.encoder_num_attention_heads, num_key_value_heads=config.encoder_num_key_value_heads, + qkv_bias=config.attn_qkv_bias, ) self.post_attention_layernorm = LayerNormNoBias( config.hidden_size, eps=config.layer_norm_eps @@ -239,6 +251,7 @@ def __init__(self, config: MoonshineConfig): num_heads=config.num_attention_heads, num_key_value_heads=config.num_key_value_heads, is_causal=True, + qkv_bias=config.attn_qkv_bias, ) self.post_attention_layernorm = LayerNormNoBias( config.hidden_size, eps=config.layer_norm_eps @@ -247,6 +260,7 @@ def __init__(self, config: MoonshineConfig): config, num_heads=config.num_attention_heads, num_key_value_heads=config.num_key_value_heads, + qkv_bias=config.attn_qkv_bias, ) self.final_layernorm = LayerNormNoBias(config.hidden_size, eps=config.layer_norm_eps) self.mlp = MoonshineDecoderMLP(config) diff --git a/src/mobius/models/moonshine_streaming.py b/src/mobius/models/moonshine_streaming.py new file mode 100644 index 000000000..520f0ee92 --- /dev/null +++ b/src/mobius/models/moonshine_streaming.py @@ -0,0 +1,425 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Moonshine Streaming raw-waveform encoder-decoder model for speech recognition. + +Replicates Hugging Face ``MoonshineStreamingForConditionalGeneration`` as separate +encoder and cached-decoder ONNX graphs. + +Compared with offline Moonshine the streaming variant changes the whole audio +front end and the encoder's positional scheme: + +* **Framing front end** — the waveform is reshaped into fixed ``frame_ms`` frames + (80 raw samples at 16 kHz / 5 ms), per-frame mean/RMS normalised (CMVN), + compressed with a learned ``asinh(exp(log_k) * x)`` gain, and projected by a + single bias-free linear layer. +* **Causal downsampling** — two left-padded stride-2 convolutions replace + Moonshine's centred convolution stem, so no future frame ever leaks backwards. +* **No encoder RoPE** — encoder self-attention is purely content based; ordering + comes from the causal stem plus per-layer asymmetric ``(left, right)`` sliding + windows. ``right`` is the strict lookahead in encoder frames and is ``0`` for + the fully causal layers, which is what bounds streaming latency. +* **Unit-offset LayerNorm** — encoder norms are affine-free ``LayerNorm`` scaled + by ``gamma + 1``. +* **Context adapter** — the decoder adds a learned absolute position table + (``pos_emb``) to the encoder output, then optionally projects it to the decoder + width, before cross-attention. + +The decoder itself (partial interleaved RoPE, cached causal self-attention, +cross-attention, fused gate/up SiLU MLP, bias-free LayerNorms) is identical to +offline Moonshine and is reused directly. +""" + +from __future__ import annotations + +from typing import overload + +import onnx_ir as ir +from onnxscript import OpBuilder, nn + +from mobius._configs import MoonshineStreamingConfig +from mobius.components import Conv1d, Embedding, Linear, SiLU, get_activation +from mobius.models.moonshine import ( + MoonshineAttention, + MoonshineDecoderModel, + MoonshineForConditionalGeneration, +) + + +class MoonshineStreamingLayerNorm(nn.Module): + """Affine-free LayerNorm scaled by ``gamma + 1``. + + Mirrors HF ``MoonshineStreamingLayerNorm``: an ``nn.LayerNorm`` with + ``elementwise_affine=False`` followed by a multiplication with + ``gamma + unit_offset``. The checkpoint therefore stores ``gamma`` (not + ``weight``), initialised to zero-mean around the unit scale. + """ + + def __init__(self, hidden_size: int, eps: float = 1e-5): + super().__init__() + self.gamma = nn.Parameter([hidden_size]) + self._eps = eps + + def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: + # LayerNorm without affine followed by ``* (gamma + 1)`` is exactly + # LayerNormalization with scale = gamma + 1 and no bias. + return op.LayerNormalization( + hidden_states, + op.Add(self.gamma, 1.0), + None, + epsilon=self._eps, + axis=-1, + ) + + +class MoonshineStreamingAsinhCompression(nn.Module): + """Learned log-domain gain followed by ``asinh`` dynamic-range compression.""" + + def __init__(self): + super().__init__() + # Scalar parameter; kept in float32 because it feeds the float32 + # normalisation stage of the front end. + self.log_k = nn.Parameter([]) + self.log_k._keep_float32 = True + + def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: + return op.Asinh(op.Mul(op.Exp(self.log_k), hidden_states)) + + +class MoonshineStreamingCausalConv1d(Conv1d): + """Left-padded strided 1-D convolution that also carries the frame mask. + + HF ``MoonshineStreamingCausalConv1d`` pads the input by + ``(kernel_size - 1) * dilation`` on the left only, so an output frame never + depends on a future input frame. The padding mask is pushed through the same + receptive field with an "any valid" reduction and used to zero invalid output + frames. + """ + + def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int): + super().__init__( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=(kernel_size - 1, 0), + bias=True, + ) + self._left_pad = kernel_size - 1 + self._window = kernel_size + self._stride = stride + + @overload + def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: ... + + @overload + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + mask: ir.Value, + dtype: ir.DataType, + ) -> tuple[ir.Value, ir.Value]: ... + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + mask: ir.Value | None = None, + dtype: ir.DataType = ir.DataType.FLOAT, + ) -> ir.Value | tuple[ir.Value, ir.Value]: + """Run the convolution and downsample the mask. + + Args: + hidden_states: ``(B, C_in, T)`` channel-first frames. + mask: ``(B, T)`` bool frame-validity mask. + dtype: Compute dtype used to zero masked output frames. + + Returns: + ``((B, C_out, T_out), (B, T_out))`` hidden states and mask. + """ + hidden_states = super().forward(op, hidden_states) # (B, C_out, T_out) + if mask is None: + return hidden_states + + # An output frame is valid when *any* input frame in its (left-padded) + # receptive field is valid — max-pooling over the same window with the + # same stride is exactly HF's ``conv1d(mask, ones) > 0``. + mask_values = op.Cast(op.Unsqueeze(mask, [1]), to=ir.DataType.FLOAT) # (B, 1, T) + mask_values = op.Pad(mask_values, [0, 0, self._left_pad, 0, 0, 0]) + mask_values = op.MaxPool( + mask_values, kernel_shape=[self._window], strides=[self._stride] + ) + mask_4d = op.Greater(mask_values, 0.0) # (B, 1, T_out) bool + hidden_states = op.Mul(hidden_states, op.Cast(mask_4d, to=dtype)) + return hidden_states, op.Squeeze(mask_4d, [1]) + + +class MoonshineStreamingEncoderEmbedder(nn.Module): + """Raw-waveform framing front end: CMVN, asinh compression, causal convs.""" + + def __init__(self, config: MoonshineStreamingConfig): + super().__init__() + hidden_size = config.encoder_hidden_size + self.comp = MoonshineStreamingAsinhCompression() + self.conv1 = MoonshineStreamingCausalConv1d( + hidden_size, 2 * hidden_size, kernel_size=5, stride=2 + ) + self.conv2 = MoonshineStreamingCausalConv1d( + 2 * hidden_size, hidden_size, kernel_size=5, stride=2 + ) + self.linear = Linear(config.frame_length, hidden_size, bias=False) + self.activation = SiLU() + self._frame_length = config.frame_length + self._cmvn_eps = config.encoder_cmvn_eps + self._dtype = config.dtype + + def forward( + self, + op: OpBuilder, + input_values: ir.Value, + attention_mask: ir.Value, + ) -> tuple[ir.Value, ir.Value]: + # Raw audio (B, L) -> (B, T, frame_len). ``L`` must be a multiple of + # ``frame_len``; the processor enforces this with pad_to_multiple_of. + frames = op.Reshape(input_values, [0, -1, self._frame_length]) + + # Per-frame CMVN and the asinh gain run in float32: squared waveform + # amplitudes are ~1e-6 and would collapse into float16 subnormals. + frames = op.Cast(frames, to=ir.DataType.FLOAT) + mean = op.ReduceMean(frames, [-1], keepdims=1) # (B, T, 1) + centered = op.Sub(frames, mean) + variance = op.ReduceMean(op.Mul(centered, centered), [-1], keepdims=1) + normalized = op.Div(centered, op.Sqrt(op.Add(variance, self._cmvn_eps))) + hidden_states = self.comp(op, normalized) # (B, T, frame_len) + if self._dtype != ir.DataType.FLOAT: + hidden_states = op.Cast(hidden_states, to=self._dtype) + hidden_states = self.activation(op, self.linear(op, hidden_states)) # (B, T, D) + + # Only frames fully covered by real samples are valid; HF uses integer + # division of the sample-level mask by the frame length. + valid_frames = op.Div( + op.ReduceSum(attention_mask, [-1], keepdims=1), self._frame_length + ) # (B, 1) + frame_count = op.Shape(hidden_states, start=1, end=2) + positions = op.Unsqueeze(op.Range(0, op.Squeeze(frame_count, [0]), 1), [0]) + frame_mask = op.Less(positions, valid_frames) # (B, T) bool + hidden_states = op.Mul( + hidden_states, op.Cast(op.Unsqueeze(frame_mask, [2]), to=self._dtype) + ) + + # Channel-first for the convolutions; each halves the frame rate. + hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1]) # (B, D, T) + hidden_states, frame_mask = self.conv1(op, hidden_states, frame_mask, self._dtype) + hidden_states = self.activation(op, hidden_states) + hidden_states, frame_mask = self.conv2(op, hidden_states, frame_mask, self._dtype) + hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1]) # (B, T2, D) + return hidden_states, frame_mask + + +class MoonshineStreamingEncoderMLP(nn.Module): + """Feed-forward network of an encoder layer (``fc1`` -> act -> ``fc2``).""" + + def __init__(self, config: MoonshineStreamingConfig): + super().__init__() + self.fc1 = Linear(config.encoder_hidden_size, config.encoder_intermediate_size) + self.fc2 = Linear(config.encoder_intermediate_size, config.encoder_hidden_size) + self._activation = get_activation(config.encoder_hidden_act) + + def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: + return self.fc2(op, self._activation(op, self.fc1(op, hidden_states))) + + +class MoonshineStreamingEncoderLayer(nn.Module): + """Pre-norm encoder layer with unit-offset norms and no positional rotation.""" + + def __init__(self, config: MoonshineStreamingConfig): + super().__init__() + self.input_layernorm = MoonshineStreamingLayerNorm( + config.encoder_hidden_size, eps=config.layer_norm_eps + ) + self.self_attn = MoonshineAttention( + config, + num_heads=config.encoder_num_attention_heads, + num_key_value_heads=config.encoder_num_key_value_heads, + hidden_size=config.encoder_hidden_size, + head_dim=config.encoder_head_dim, + # Upstream gates every encoder projection — including o_proj — on + # the encoder sub-config's attention_bias. + qkv_bias=config.encoder_attention_bias, + o_bias=config.encoder_attention_bias, + ) + self.post_attention_layernorm = MoonshineStreamingLayerNorm( + config.encoder_hidden_size, eps=config.layer_norm_eps + ) + self.mlp = MoonshineStreamingEncoderMLP(config) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + attention_bias: ir.Value, + ) -> ir.Value: + residual = hidden_states + hidden_states = self.input_layernorm(op, hidden_states) + # Streaming encoder attention carries no RoPE — the sliding window and + # the causal convolution stem supply all positional structure. + hidden_states, _ = self.self_attn(op, hidden_states, attention_bias=attention_bias) + hidden_states = op.Add(residual, hidden_states) + + residual = hidden_states + hidden_states = self.post_attention_layernorm(op, hidden_states) + hidden_states = self.mlp(op, hidden_states) + return op.Add(residual, hidden_states) + + +def _sliding_window_attention_bias( + op: OpBuilder, + frame_mask: ir.Value, + frame_count: ir.Value, + window: tuple[int, int], + dtype: ir.DataType, +) -> ir.Value: + """Additive ``[B, 1, T, T]`` bias for one asymmetric bidirectional window. + + HF combines the padding mask with + ``(0 <= q - kv < left) or (0 < kv - q < right)``. Expressed as bounds on + ``dist = q - kv`` that is ``-(right - 1) <= dist <= left - 1``, with the + lower bound clamped to ``0`` when ``right == 0`` (fully causal layer). + + Args: + frame_mask: ``(B, T)`` bool mask of valid encoder frames. + frame_count: ``(1,)`` int64 encoder sequence length. + window: ``(left, right)`` window of the layer. + dtype: Compute dtype of the emitted bias. + """ + left, right = window + lower = -(right - 1) if right >= 1 else 0 + upper = left - 1 + + positions = op.Range(0, op.Squeeze(frame_count, [0]), 1) # (T,) + distance = op.Sub(op.Unsqueeze(positions, [1]), op.Unsqueeze(positions, [0])) # (T, T) + in_window = op.And( + op.GreaterOrEqual(distance, lower), + op.LessOrEqual(distance, upper), + ) + # (1, 1, T, T) window AND (B, 1, 1, T) key validity -> (B, 1, T, T). + allowed = op.And(op.Unsqueeze(in_window, [0, 1]), op.Unsqueeze(frame_mask, [1, 2])) + return op.Cast(op.Where(allowed, 0.0, float(dtype.min)), to=dtype) + + +class MoonshineStreamingEncoderModel(nn.Module): + """Streaming audio encoder: framing front end plus windowed transformer.""" + + def __init__(self, config: MoonshineStreamingConfig): + super().__init__() + self.embedder = MoonshineStreamingEncoderEmbedder(config) + self.layers = nn.ModuleList( + [ + MoonshineStreamingEncoderLayer(config) + for _ in range(config.encoder_num_hidden_layers) + ] + ) + self.final_norm = MoonshineStreamingLayerNorm( + config.encoder_hidden_size, eps=config.layer_norm_eps + ) + self._windows = config.encoder_sliding_windows + self._dtype = config.dtype + + def forward( + self, + op: OpBuilder, + input_values: ir.Value, + attention_mask: ir.Value, + ) -> tuple[ir.Value, ir.Value]: + hidden_states, frame_mask = self.embedder(op, input_values, attention_mask) + frame_count = op.Shape(hidden_states, start=1, end=2) + + # Layers share only a handful of distinct windows, so build each bias + # once and reuse it across the layers that declare it. + biases: dict[tuple[int, int], ir.Value] = {} + for window in self._windows: + if window not in biases: + biases[window] = _sliding_window_attention_bias( + op, frame_mask, frame_count, window, self._dtype + ) + + for layer, window in zip(self.layers, self._windows): + hidden_states = layer(op, hidden_states, biases[window]) + + hidden_states = self.final_norm(op, hidden_states) + encoder_attention_mask = op.Cast(frame_mask, to=ir.DataType.INT64) + return hidden_states, encoder_attention_mask + + +class MoonshineStreamingDecoderModel(MoonshineDecoderModel): + """Moonshine decoder with the streaming context adapter on encoder states.""" + + def __init__(self, config: MoonshineStreamingConfig): + super().__init__(config) + # Absolute learned positions added to the encoder output. The table is + # sized by the decoder's ``max_position_embeddings`` but indexed by the + # encoder frame position, exactly as upstream does. + self.pos_emb = Embedding(config.max_position_embeddings, config.encoder_hidden_size) + if config.encoder_hidden_size != config.hidden_size: + self.proj = Linear(config.encoder_hidden_size, config.hidden_size, bias=False) + else: + # HF uses ``nn.Identity`` here, which contributes no checkpoint entry. + self.proj = None + + def forward( + self, + op: OpBuilder, + decoder_input_ids: ir.Value, + encoder_hidden_states: ir.Value, + encoder_attention_mask: ir.Value, + position_ids: ir.Value, + past_key_values: list[tuple[ir.Value, ir.Value]] | None = None, + ) -> tuple[ir.Value, list[tuple[ir.Value, ir.Value]]]: + # Context adapter: encoder_hidden_states (B, E, De) + pos_emb[0:E] + encoder_length = op.Shape(encoder_hidden_states, start=1, end=2) + encoder_positions = op.Range(0, op.Squeeze(encoder_length, [0]), 1) # (E,) + encoder_hidden_states = op.Add( + encoder_hidden_states, self.pos_emb(op, encoder_positions) + ) + if self.proj is not None: + encoder_hidden_states = self.proj(op, encoder_hidden_states) # (B, E, D) + + return super().forward( + op, + decoder_input_ids, + encoder_hidden_states, + encoder_attention_mask, + position_ids, + past_key_values, + ) + + +class _MoonshineStreamingModel(nn.Module): + def __init__(self, config: MoonshineStreamingConfig): + super().__init__() + self.encoder = MoonshineStreamingEncoderModel(config) + self.decoder = MoonshineStreamingDecoderModel(config) + + +class MoonshineStreamingForConditionalGeneration(MoonshineForConditionalGeneration): + """Moonshine Streaming encoder-decoder model for low-latency ASR. + + Replicates Hugging Face ``MoonshineStreamingForConditionalGeneration``: a + causal framing/convolution front end feeding a windowed (bounded-lookahead) + transformer encoder, and a cached Moonshine decoder that cross-attends to the + position-adapted encoder context. + """ + + default_task: str = "speech-to-text" + category: str = "Speech-to-Text" + config_class: type = MoonshineStreamingConfig + + def __init__(self, config: MoonshineStreamingConfig): + nn.Module.__init__(self) + self.config = config + self.model = _MoonshineStreamingModel(config) + self.proj_out = self.model.decoder.proj_out + + +__all__ = ["MoonshineStreamingForConditionalGeneration"] diff --git a/src/mobius/tasks/_speech_to_text.py b/src/mobius/tasks/_speech_to_text.py index 243da219f..59fa10957 100644 --- a/src/mobius/tasks/_speech_to_text.py +++ b/src/mobius/tasks/_speech_to_text.py @@ -133,7 +133,7 @@ def _build_decoder( encoder_hidden_states = builder.input( "encoder_hidden_states", dtype=config.dtype, - shape=[batch, encoder_seq_len, config.hidden_size], + shape=[batch, encoder_seq_len, config.encoder_output_size], ) position_ids = builder.input( "position_ids", diff --git a/testdata/cases/audio/moonshine-streaming-tiny.yaml b/testdata/cases/audio/moonshine-streaming-tiny.yaml new file mode 100644 index 000000000..ffa486c46 --- /dev/null +++ b/testdata/cases/audio/moonshine-streaming-tiny.yaml @@ -0,0 +1,23 @@ +model_id: "moonshine-ai/moonshine-streaming-tiny" +model_type: "moonshine_streaming" +revision: "f8e9dfd8c562c257c151a907b7b7f2fe8ff8511a" +task_type: "speech-to-text" +dtype: "float32" + +inputs: + audio: + - "652-129742-0006.flac" + +generation: + max_new_tokens: 50 + do_sample: false + eos_token_id: 2 + exact_match: true + +level: "L4+L5" + +notes: >- + Moonshine Streaming tiny: causal framing front end (5 ms frames, CMVN, asinh + compression, two left-padded stride-2 convolutions) feeding a windowed encoder + with per-layer asymmetric (left, right) lookahead, and a cached Moonshine + decoder that cross-attends to the position-adapted encoder context. diff --git a/testdata/golden/audio/moonshine-streaming-tiny.json b/testdata/golden/audio/moonshine-streaming-tiny.json new file mode 100644 index 000000000..9aceb5cfe --- /dev/null +++ b/testdata/golden/audio/moonshine-streaming-tiny.json @@ -0,0 +1,37 @@ +{ + "top1_id": 9243, + "top2_id": 5777, + "top10_ids": [ + 9243, + 5777, + 8251, + 1530, + 315, + 9159, + 678, + 476, + 319, + 3037 + ], + "top10_logits": [ + "0x1.67f3c80000000p+3", + "0x1.4ca65e0000000p+3", + "0x1.4464480000000p+3", + "0x1.436b680000000p+3", + "0x1.3cdb2e0000000p+3", + "0x1.0a4b140000000p+3", + "0x1.049ab00000000p+3", + "0x1.01bf440000000p+3", + "0x1.f0596e0000000p+2", + "0x1.cd77c60000000p+2" + ], + "logits_summary": [ + "0x1.67f3c80000000p+3", + "-0x1.887c780000000p+4", + "-0x1.4de05fb4c5380p+3", + "0x1.057eaedb36406p+2" + ], + "input_ids": [ + 1 + ] +} diff --git a/testdata/golden/audio/moonshine-streaming-tiny_generation.json b/testdata/golden/audio/moonshine-streaming-tiny_generation.json new file mode 100644 index 000000000..4a2af5e93 --- /dev/null +++ b/testdata/golden/audio/moonshine-streaming-tiny_generation.json @@ -0,0 +1,40 @@ +{ + "model_id": "moonshine-ai/moonshine-streaming-tiny", + "prompt": "652-129742-0006.flac", + "generated_tokens": [ + 9243, + 352, + 361, + 13609, + 1122, + 11586, + 895, + 29892, + 2125, + 11220, + 1045, + 2356, + 5777, + 352, + 361, + 13609, + 29892, + 2867, + 964, + 14202, + 29892, + 4417, + 15795, + 29892, + 1236, + 2496, + 322, + 13848, + 387, + 279, + 304, + 4259, + 29889 + ], + "generated_text": "Cauliflower mayonnaise, take cold boiled cauliflower, break into branches, adding salt, pepper and vinegar to season." +} diff --git a/tests/_test_configs.py b/tests/_test_configs.py index 23264a0b0..ac45414f7 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -48,6 +48,7 @@ MiniMaxConfig, MllamaConfig, MoonshineConfig, + MoonshineStreamingConfig, MuseGlimmerConfig, NanoChatConfig, NemotronHConfig, @@ -3554,6 +3555,25 @@ def vl_overrides(model_type: str) -> dict: }, True, ), + # --- Moonshine Streaming (causal framing front end + windowed encoder) --- + ( + "moonshine_streaming", + { + "_config_cls": MoonshineStreamingConfig, + "num_key_value_heads": TINY_HEADS, + "partial_rotary_factor": 0.8, + "rope_type": "default", + "rope_interleave": True, + "mlp_bias": True, + "tie_word_embeddings": False, + "encoder_num_hidden_layers": TINY_LAYERS, + "encoder_num_attention_heads": TINY_HEADS, + "encoder_num_key_value_heads": TINY_HEADS, + # Asymmetric lookahead on layer 0, fully causal on layer 1. + "encoder_sliding_windows": ((16, 4), (16, 0)), + }, + True, + ), # --- Whisper (speech-to-text, encoder-decoder) --- ( "whisper", diff --git a/tests/build_graph/speech_test.py b/tests/build_graph/speech_test.py index 9c8839aac..74ad82ec4 100644 --- a/tests/build_graph/speech_test.py +++ b/tests/build_graph/speech_test.py @@ -284,6 +284,120 @@ def test_moonshine_architecture_initializers(self): assert "Sigmoid" not in decoder_ops +class TestBuildGraphMoonshineStreaming: + """Verify Moonshine Streaming's causal framing encoder and cached decoder.""" + + def _config(self, **overrides): + from mobius._configs import MoonshineStreamingConfig + + options = dict( + _config_cls=MoonshineStreamingConfig, + num_key_value_heads=TINY_HEADS, + partial_rotary_factor=0.8, + rope_type="default", + rope_interleave=True, + mlp_bias=True, + tie_word_embeddings=False, + encoder_num_hidden_layers=TINY_LAYERS, + encoder_num_attention_heads=TINY_HEADS, + encoder_num_key_value_heads=TINY_HEADS, + encoder_sliding_windows=((16, 4), (16, 0)), + decoder_start_token_id=1, + ) + options.update(overrides) + return _base_config(**options) + + def _build(self, config): + from mobius.models import MoonshineStreamingForConditionalGeneration + from mobius.tasks import SpeechToTextTask + + return build_from_module( + MoonshineStreamingForConditionalGeneration(config), + config, + task=SpeechToTextTask(), + ) + + def test_package_and_io(self): + package = self._build(self._config()) + + assert set(package) == {"encoder", "decoder"} + encoder_inputs = {value.name for value in package["encoder"].graph.inputs} + encoder_outputs = {value.name for value in package["encoder"].graph.outputs} + decoder_inputs = {value.name for value in package["decoder"].graph.inputs} + assert encoder_inputs == {"input_values", "attention_mask"} + assert encoder_outputs == {"encoder_hidden_states", "encoder_attention_mask"} + assert "encoder_attention_mask" in decoder_inputs + assert "position_ids" in decoder_inputs + for layer_idx in range(TINY_LAYERS): + assert f"past_key_values.{layer_idx}.key" in decoder_inputs + + def test_architecture_initializers_match_huggingface_names(self): + package = self._build(self._config()) + encoder_initializers = set(package["encoder"].graph.initializers) + decoder_initializers = set(package["decoder"].graph.initializers) + + assert "encoder.embedder.linear.weight" in encoder_initializers + assert "encoder.embedder.linear.bias" not in encoder_initializers + assert "encoder.embedder.comp.log_k" in encoder_initializers + for conv in ("conv1", "conv2"): + assert f"encoder.embedder.{conv}.weight" in encoder_initializers + assert f"encoder.embedder.{conv}.bias" in encoder_initializers + + assert "encoder.layers.0.input_layernorm.gamma" in encoder_initializers + assert "encoder.layers.0.input_layernorm.weight" not in encoder_initializers + assert "encoder.final_norm.gamma" in encoder_initializers + assert "encoder.layers.0.self_attn.q_proj.bias" not in encoder_initializers + assert "encoder.layers.0.mlp.fc1.bias" in encoder_initializers + assert not any("rotary" in name for name in encoder_initializers) + + assert "decoder.pos_emb.weight" in decoder_initializers + assert "decoder.embed_tokens.weight" in decoder_initializers + assert "decoder.proj_out.weight" in decoder_initializers + assert "decoder.layers.0.encoder_attn.q_proj.weight" in decoder_initializers + assert "decoder.proj.weight" not in decoder_initializers + + def test_encoder_conv_padding_is_causal(self): + package = self._build(self._config()) + conv_pads = [ + tuple(node.attributes["pads"].as_ints()) + for node in package["encoder"].graph + if node.op_type == "Conv" + ] + assert conv_pads, "encoder should contain convolutions" + assert all(pads == (4, 0) for pads in conv_pads), conv_pads + + def test_narrower_encoder_adds_projection(self): + config = self._config(encoder_hidden_size=32, encoder_head_dim=8) + package = self._build(config) + assert config.encoder_output_size == 32 + decoder_initializers = set(package["decoder"].graph.initializers) + assert "decoder.proj.weight" in decoder_initializers + assert "decoder.pos_emb.weight" in decoder_initializers + encoder_input = next( + value + for value in package["decoder"].graph.inputs + if value.name == "encoder_hidden_states" + ) + assert encoder_input.shape[2] == 32 + + def test_attention_bias_follows_upstream_gating(self): + package = self._build(self._config(attn_qkv_bias=True, encoder_attention_bias=True)) + encoder_initializers = set(package["encoder"].graph.initializers) + decoder_initializers = set(package["decoder"].graph.initializers) + + for projection in ("q_proj", "k_proj", "v_proj", "o_proj"): + assert f"encoder.layers.0.self_attn.{projection}.bias" in encoder_initializers + for projection in ("q_proj", "k_proj", "v_proj"): + assert f"decoder.layers.0.self_attn.{projection}.bias" in decoder_initializers + assert f"decoder.layers.0.encoder_attn.{projection}.bias" in decoder_initializers + assert "decoder.layers.0.self_attn.o_proj.bias" not in decoder_initializers + assert "decoder.layers.0.encoder_attn.o_proj.bias" not in decoder_initializers + + def test_sliding_window_length_is_validated(self): + with pytest.raises(ValueError, match="encoder_sliding_windows"): + self._config(encoder_sliding_windows=((16, 4),)) + + class TestBuildGraphGlmAsr: """Verify GLM-ASR's audio encoder, projector, embedding, and decoder split.""" diff --git a/tests/moonshine_streaming_integration_test.py b/tests/moonshine_streaming_integration_test.py new file mode 100644 index 000000000..7ee7be995 --- /dev/null +++ b/tests/moonshine_streaming_integration_test.py @@ -0,0 +1,891 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Numerical parity tests for Moonshine Streaming speech recognition. + +Covers the pieces that make ``moonshine_streaming`` different from offline +Moonshine: the raw-waveform framing front end, causal strided convolutions with +mask propagation, per-layer asymmetric ``(left, right)`` sliding windows, the +unit-offset encoder LayerNorm, and the decoder's absolute position context +adapter. Real-weight cases run against a pinned checkpoint revision and real +nonzero speech, and chain the ONNX decoder off the ONNX encoder output rather +than a HuggingFace intermediate. +""" + +from __future__ import annotations + +from pathlib import Path + +import librosa +import numpy as np +import pytest +import torch +import transformers + +from mobius import build, build_from_module +from mobius._configs import MoonshineStreamingConfig +from mobius._testing.comparison import assert_logits_close +from mobius._testing.golden import ( + discover_test_cases, + generation_json_path_for_case, + load_generation_golden, +) +from mobius._testing.ort_inference import OnnxModelSession +from mobius.models import MoonshineStreamingForConditionalGeneration +from mobius.tasks import SpeechToTextTask + +_MODEL_ID = "moonshine-ai/moonshine-streaming-tiny" +_REVISION = "f8e9dfd8c562c257c151a907b7b7f2fe8ff8511a" +_AUDIO_PATH = Path(__file__).parent.parent / "testdata" / "652-129742-0006.flac" +_EOS_TOKEN_ID = 2 + +pytestmark = pytest.mark.skipif( + not hasattr(transformers, "MoonshineStreamingConfig"), + reason="transformers build has no moonshine_streaming support", +) + + +def _tiny_hf_config( + sliding_windows: list[list[int]] | None = None, attention_bias: bool = False +): + """A production-shaped but tiny Moonshine Streaming config. + + Keeps the real frame length (80 samples), the real partial rotary factor and + a mixed lookahead schedule so the window logic is genuinely exercised. + """ + encoder_config_cls = transformers.models.moonshine_streaming.configuration_moonshine_streaming.MoonshineStreamingEncoderConfig + encoder_config = encoder_config_cls( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=16, + hidden_act="gelu", + attention_bias=attention_bias, + sliding_windows=sliding_windows or [[16, 4], [16, 0]], + sample_rate=16_000, + frame_ms=5.0, + ) + return transformers.MoonshineStreamingConfig( + encoder_config=encoder_config, + vocab_size=256, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=16, + max_position_embeddings=128, + attention_bias=attention_bias, + rope_parameters={ + "rope_type": "default", + "rope_theta": 10_000.0, + "partial_rotary_factor": 0.8, + }, + tie_word_embeddings=False, + decoder_start_token_id=1, + ) + + +def _weighted_package(hf_model, hf_config): + config = MoonshineStreamingConfig.from_transformers(hf_config) + module = MoonshineStreamingForConditionalGeneration(config) + package = build_from_module(module, config, task=SpeechToTextTask()) + state_dict = { + name: value.detach().clone() for name, value in hf_model.state_dict().items() + } + package.apply_weights(module.preprocess_weights(state_dict)) + return package, config + + +def _run_encoder(package, input_values, attention_mask): + session = OnnxModelSession(package["encoder"]) + try: + return session.run( + { + "input_values": input_values.astype(np.float32), + "attention_mask": attention_mask.astype(np.int64), + } + ) + finally: + session.close() + + +def _empty_cache_feeds( + config: MoonshineStreamingConfig, + batch_size: int = 1, + dtype: np.dtype | type = np.float32, +): + feeds = {} + for layer_idx in range(config.num_hidden_layers): + for kind in ("key", "value"): + feeds[f"past_key_values.{layer_idx}.{kind}"] = np.zeros( + (batch_size, config.num_key_value_heads, 0, config.head_dim), + dtype=dtype, + ) + return feeds + + +def _decoder_feeds( + config, + decoder_input_ids, + encoder_hidden_states, + encoder_attention_mask, + past_key_values=None, + position_offset=0, + dtype: np.dtype | type = np.float32, +): + batch_size, sequence_length = decoder_input_ids.shape + position_ids = np.arange( + position_offset, position_offset + sequence_length, dtype=np.int64 + ) + return { + "decoder_input_ids": decoder_input_ids.astype(np.int64), + "encoder_hidden_states": encoder_hidden_states.astype(dtype), + "encoder_attention_mask": encoder_attention_mask.astype(np.int64), + "position_ids": np.broadcast_to( + position_ids[None, :], (batch_size, sequence_length) + ).copy(), + **(past_key_values or _empty_cache_feeds(config, batch_size, dtype=dtype)), + } + + +def _run_decoder(package, config, *args, **kwargs): + session = OnnxModelSession(package["decoder"]) + try: + return session.run(_decoder_feeds(config, *args, **kwargs)) + finally: + session.close() + + +def _onnx_greedy_generate(package, config, encoder_outputs, max_new_tokens, start_id): + """Greedy decode entirely through the exported ONNX decoder.""" + session = OnnxModelSession(package["decoder"]) + try: + dtype = session.get_input_dtype("past_key_values.0.key") or np.float32 + cache = _empty_cache_feeds(config, dtype=dtype) + current = np.array([[start_id]], dtype=np.int64) + generated: list[int] = [] + for step in range(max_new_tokens): + outputs = session.run( + _decoder_feeds( + config, + current, + encoder_outputs["encoder_hidden_states"], + encoder_outputs["encoder_attention_mask"], + past_key_values=cache, + position_offset=step, + dtype=dtype, + ) + ) + next_token = int(np.asarray(outputs["logits"], dtype=np.float32)[0, -1].argmax()) + generated.append(next_token) + if next_token == _EOS_TOKEN_ID: + break + current = np.array([[next_token]], dtype=np.int64) + cache = { + f"past_key_values.{layer_idx}.{kind}": outputs[f"present.{layer_idx}.{kind}"] + for layer_idx in range(config.num_hidden_layers) + for kind in ("key", "value") + } + return generated + finally: + session.close() + + +@pytest.fixture(scope="module") +def synthetic_models(): + torch.manual_seed(42) + hf_config = _tiny_hf_config() + hf_model = transformers.MoonshineStreamingForConditionalGeneration(hf_config).eval() + package, config = _weighted_package(hf_model, hf_config) + return hf_model, package, config + + +def _nonmutating_hf_greedy(hf_model, encoder_hidden_states, encoder_attention_mask, steps): + """Greedy decode with a pristine encoder context handed to every step. + + HuggingFace's ``MoonshineStreamingDecoder`` adds its absolute position table + to ``encoder_hidden_states`` with ``+=``, mutating the caller's tensor. With + cached cross-attention the mutated tensor is never read again, so the result + is unaffected — but the reference this test compares against must not depend + on that, so each step gets a fresh clone. This is the same semantics the + exported ONNX decoder implements: encoder context is added exactly once. + + Returns ``(tokens, per_step_logits)``. + """ + from transformers.cache_utils import DynamicCache, EncoderDecoderCache + + cache = EncoderDecoderCache( + DynamicCache(config=hf_model.config), DynamicCache(config=hf_model.config) + ) + current = torch.tensor([[hf_model.config.decoder_start_token_id]], dtype=torch.long) + tokens: list[int] = [] + per_step: list[np.ndarray] = [] + for step in range(steps): + with torch.no_grad(): + output = hf_model.model.decoder( + input_ids=current, + encoder_hidden_states=encoder_hidden_states.clone(), + encoder_attention_mask=encoder_attention_mask, + past_key_values=cache, + position_ids=torch.tensor([[step]], dtype=torch.long), + use_cache=True, + ) + logits = hf_model.proj_out(output.last_hidden_state) + cache = output.past_key_values + per_step.append(logits[0, -1].float().numpy()) + token = int(logits[0, -1].argmax()) + tokens.append(token) + if token == _EOS_TOKEN_ID: + break + current = torch.tensor([[token]], dtype=torch.long) + return tokens, per_step + + +@pytest.fixture(scope="module") +def real_models(): + package = build(_MODEL_ID, dtype="f32", load_weights=True, revision=_REVISION) + hf_model = transformers.AutoModelForSpeechSeq2Seq.from_pretrained( + _MODEL_ID, revision=_REVISION + ).eval() + processor = transformers.AutoProcessor.from_pretrained(_MODEL_ID, revision=_REVISION) + config = package.config + assert isinstance(config, MoonshineStreamingConfig) + return hf_model, processor, package, config + + +@pytest.fixture(scope="module") +def real_audio_inputs(real_models): + _hf_model, processor, _package, _config = real_models + audio, _sample_rate = librosa.load(str(_AUDIO_PATH), sr=16_000) + assert np.abs(audio).max() > 0.1, "fixture must be real nonzero speech" + return processor(audio, sampling_rate=16_000, return_tensors="np") + + +class TestMoonshineStreamingConfigExtraction: + """Config extraction preserves the streaming-specific architecture fields.""" + + def test_extracts_encoder_and_window_schedule(self): + config = MoonshineStreamingConfig.from_transformers(_tiny_hf_config()) + assert config.model_type == "moonshine_streaming" + assert config.encoder_sliding_windows == ((16, 4), (16, 0)) + assert config.encoder_hidden_act == "gelu" + assert config.decoder_hidden_act == "silu" + assert config.partial_rotary_factor == pytest.approx(0.8) + assert config.rope_interleave is True + assert config.tie_word_embeddings is False + # 16 kHz * 5 ms = 80 raw samples per encoder frame. + assert config.frame_length == 80 + assert config.encoder_output_size == config.encoder_hidden_size == 64 + + def test_extracts_from_raw_pinned_json(self): + """The pinned checkpoint's raw JSON (nested dict sub-config) resolves.""" + raw = transformers.AutoConfig.from_pretrained(_MODEL_ID, revision=_REVISION) + config = MoonshineStreamingConfig.from_transformers(raw) + assert config.encoder_num_hidden_layers == 6 + assert config.num_hidden_layers == 6 + assert config.encoder_sliding_windows == ( + (16, 4), + (16, 4), + (16, 0), + (16, 0), + (16, 4), + (16, 4), + ) + assert config.head_dim == 40 + assert config.encoder_head_dim == 40 + assert config.vocab_size == 32_768 + assert config.decoder_start_token_id == 1 + + def test_rejects_other_model_types(self): + config = _tiny_hf_config() + config.model_type = "moonshine" + with pytest.raises(ValueError, match="moonshine_streaming"): + MoonshineStreamingConfig.from_transformers(config) + + +def test_moonshine_streaming_l5_generation_reference_is_declared_and_valid(): + cases = [case for case in discover_test_cases(level="L5") if case.model_id == _MODEL_ID] + assert len(cases) == 1 + case = cases[0] + assert case.revision == _REVISION + generation_path = generation_json_path_for_case(case) + assert generation_path.exists() + generated_tokens = load_generation_golden(case) + assert len(generated_tokens) >= 20 + assert all(isinstance(token, int) for token in generated_tokens) + + +class TestMoonshineStreamingSyntheticParity: + """L3 random-weight parity against a reduced HuggingFace streaming model.""" + + def test_encoder_hidden_states_and_mask(self, synthetic_models): + hf_model, package, _config = synthetic_models + rng = np.random.default_rng(42) + # 200 frames of 80 samples; only the first 60 frames are real audio, so + # the causal-conv mask propagation is genuinely exercised. + input_values = (rng.standard_normal((1, 80 * 200)) * 0.05).astype(np.float32) + attention_mask = np.ones_like(input_values, dtype=np.int64) + attention_mask[:, 80 * 60 :] = 0 + + with torch.no_grad(): + expected = hf_model.get_encoder()( + input_values=torch.from_numpy(input_values), + attention_mask=torch.from_numpy(attention_mask), + ) + actual = _run_encoder(package, input_values, attention_mask) + + expected_mask = expected.attention_mask.numpy() + np.testing.assert_array_equal( + actual["encoder_attention_mask"].astype(bool), expected_mask + ) + # Padded frames are excluded everywhere downstream, and HuggingFace and + # ONNX fill fully-masked attention rows differently, so parity is + # asserted on the valid frames the model actually consumes. + valid = expected_mask[0].astype(bool) + assert valid.sum() < expected_mask.shape[1], "padding should shrink the mask" + assert_logits_close( + actual["encoder_hidden_states"][:, valid], + expected.last_hidden_state.numpy()[:, valid], + rtol=1e-3, + atol=1e-3, + ) + + @pytest.mark.parametrize( + "sliding_windows", + [[[2, 1], [3, 0]], [[16, 4], [16, 0]], [[64, 64], [64, 64]]], + ids=["tight", "default", "global"], + ) + def test_encoder_matches_each_window_schedule(self, sliding_windows): + """Asymmetric left/right windows and lookahead reproduce upstream.""" + torch.manual_seed(7) + hf_config = _tiny_hf_config(sliding_windows) + hf_model = transformers.MoonshineStreamingForConditionalGeneration(hf_config).eval() + package, _config = _weighted_package(hf_model, hf_config) + + rng = np.random.default_rng(3) + input_values = (rng.standard_normal((1, 80 * 40)) * 0.05).astype(np.float32) + attention_mask = np.ones_like(input_values, dtype=np.int64) + with torch.no_grad(): + expected = hf_model.get_encoder()( + input_values=torch.from_numpy(input_values), + attention_mask=torch.from_numpy(attention_mask), + ) + actual = _run_encoder(package, input_values, attention_mask) + assert_logits_close( + actual["encoder_hidden_states"], + expected.last_hidden_state.numpy(), + rtol=1e-3, + atol=1e-3, + ) + + @pytest.mark.parametrize("attention_bias", [False, True], ids=["nobias", "bias"]) + def test_encoder_matches_with_and_without_attention_bias(self, attention_bias): + """Upstream gates encoder q/k/v/o (and decoder q/k/v) on attention_bias.""" + torch.manual_seed(11) + hf_config = _tiny_hf_config(attention_bias=attention_bias) + hf_model = transformers.MoonshineStreamingForConditionalGeneration(hf_config).eval() + package, config = _weighted_package(hf_model, hf_config) + assert config.encoder_attention_bias is attention_bias + assert config.attn_qkv_bias is attention_bias + + rng = np.random.default_rng(5) + input_values = (rng.standard_normal((1, 80 * 40)) * 0.05).astype(np.float32) + attention_mask = np.ones_like(input_values, dtype=np.int64) + with torch.no_grad(): + expected = hf_model.get_encoder()( + input_values=torch.from_numpy(input_values), + attention_mask=torch.from_numpy(attention_mask), + ) + actual = _run_encoder(package, input_values, attention_mask) + assert_logits_close( + actual["encoder_hidden_states"], + expected.last_hidden_state.numpy(), + rtol=1e-3, + atol=1e-3, + ) + + def test_batched_ragged_padding_matches_huggingface(self, synthetic_models): + """Two rows of different real length: mask, encoder and decoder all agree.""" + hf_model, package, config = synthetic_models + rng = np.random.default_rng(21) + input_values = (rng.standard_normal((3, 80 * 150)) * 0.05).astype(np.float32) + attention_mask = np.ones_like(input_values, dtype=np.int64) + # Ragged: 150, 90 and 40 valid frames. + attention_mask[1, 80 * 90 :] = 0 + attention_mask[2, 80 * 40 :] = 0 + + with torch.no_grad(): + expected = hf_model.get_encoder()( + input_values=torch.from_numpy(input_values), + attention_mask=torch.from_numpy(attention_mask), + ) + actual = _run_encoder(package, input_values, attention_mask) + + expected_mask = expected.attention_mask.numpy() + np.testing.assert_array_equal( + actual["encoder_attention_mask"].astype(bool), expected_mask + ) + assert len(set(expected_mask.sum(axis=-1).tolist())) == 3, ( + "each row should keep a different number of valid frames" + ) + expected_hidden = expected.last_hidden_state.numpy() + for row in range(input_values.shape[0]): + valid = expected_mask[row].astype(bool) + assert_logits_close( + actual["encoder_hidden_states"][row][valid], + expected_hidden[row][valid], + rtol=1e-3, + atol=1e-3, + ) + + decoder_input_ids = np.array([[1, 5], [1, 9], [1, 13]], dtype=np.int64) + with torch.no_grad(): + prefill = hf_model.model.decoder( + input_ids=torch.from_numpy(decoder_input_ids), + encoder_hidden_states=expected.last_hidden_state.clone(), + encoder_attention_mask=expected.attention_mask, + use_cache=True, + ) + expected_logits = hf_model.proj_out(prefill.last_hidden_state).numpy() + + session = OnnxModelSession(package["decoder"]) + try: + actual_logits = session.run( + _decoder_feeds( + config, + decoder_input_ids, + expected_hidden, + expected_mask.astype(np.int64), + past_key_values=_empty_cache_feeds(config, batch_size=3), + ) + )["logits"] + finally: + session.close() + assert_logits_close(actual_logits, expected_logits, rtol=1e-3, atol=1e-3) + + def test_decoder_prefill_and_cached_decode(self, synthetic_models): + hf_model, package, config = synthetic_models + rng = np.random.default_rng(7) + input_values = (rng.standard_normal((1, 80 * 120)) * 0.05).astype(np.float32) + attention_mask = np.ones_like(input_values, dtype=np.int64) + decoder_input_ids = np.array([[1, 17, 29]], dtype=np.int64) + + with torch.no_grad(): + encoder_output = hf_model.get_encoder()( + input_values=torch.from_numpy(input_values), + attention_mask=torch.from_numpy(attention_mask), + ) + # The decoder adds pos_emb to encoder_hidden_states in place, so + # hand it a clone and keep the pristine tensor for the ONNX feed. + expected_prefill = hf_model.model.decoder( + input_ids=torch.from_numpy(decoder_input_ids), + encoder_hidden_states=encoder_output.last_hidden_state.clone(), + encoder_attention_mask=encoder_output.attention_mask, + use_cache=True, + ) + expected_prefill_logits = hf_model.proj_out( + expected_prefill.last_hidden_state + ).numpy() + + encoder_hidden_states = encoder_output.last_hidden_state.numpy() + encoder_attention_mask = encoder_output.attention_mask.numpy() + actual_prefill = _run_decoder( + package, + config, + decoder_input_ids, + encoder_hidden_states, + encoder_attention_mask, + ) + assert_logits_close( + actual_prefill["logits"], expected_prefill_logits, rtol=1e-3, atol=1e-3 + ) + + next_token = expected_prefill_logits[:, -1:].argmax(axis=-1).astype(np.int64) + with torch.no_grad(): + expected_decode = hf_model.model.decoder( + input_ids=torch.from_numpy(next_token), + encoder_hidden_states=encoder_output.last_hidden_state.clone(), + encoder_attention_mask=encoder_output.attention_mask, + past_key_values=expected_prefill.past_key_values, + use_cache=True, + ) + expected_decode_logits = hf_model.proj_out( + expected_decode.last_hidden_state + ).numpy() + + cache = { + f"past_key_values.{layer_idx}.{kind}": actual_prefill[ + f"present.{layer_idx}.{kind}" + ] + for layer_idx in range(config.num_hidden_layers) + for kind in ("key", "value") + } + actual_decode = _run_decoder( + package, + config, + next_token, + encoder_hidden_states, + encoder_attention_mask, + past_key_values=cache, + position_offset=decoder_input_ids.shape[1], + ) + assert_logits_close( + actual_decode["logits"], expected_decode_logits, rtol=1e-3, atol=1e-3 + ) + + +@pytest.mark.integration +@pytest.mark.integration_fast +class TestMoonshineStreamingRealWeightParity: + """L3 checkpoint parity on real speech, pinned to one revision.""" + + def test_processor_contract(self, real_models, real_audio_inputs): + """The audio processor emits frame-aligned raw samples plus a mask.""" + _hf_model, processor, _package, config = real_models + assert set(real_audio_inputs) == {"input_values", "attention_mask"} + input_values = real_audio_inputs["input_values"] + assert input_values.ndim == 2 + assert input_values.dtype == np.float32 + # Framing reshape requires a multiple of the frame length. + assert input_values.shape[1] % config.frame_length == 0 + assert real_audio_inputs["attention_mask"].shape == input_values.shape + + # The waveform is consumed raw: no normalisation, 16 kHz, and the frame + # alignment the encoder needs is guaranteed by pad_to_multiple_of. + extractor = processor.feature_extractor + assert extractor.sampling_rate == config.encoder_sample_rate + assert extractor.do_normalize is False + assert extractor.return_attention_mask is True + assert extractor.pad_to_multiple_of == config.frame_length + + def test_ort_genai_runtime_is_rejected(self, real_models, tmp_path): + """ORT GenAI cannot host the variable-length raw-waveform encoder.""" + from mobius.integrations.ort_genai import write_ort_genai_config + + _hf_model, _processor, package, _config = real_models + with pytest.raises(NotImplementedError, match="raw-waveform encoder"): + write_ort_genai_config(package, str(tmp_path)) + + def test_real_encoder_hidden_states(self, real_models, real_audio_inputs): + hf_model, _processor, package, _config = real_models + with torch.no_grad(): + expected = hf_model.get_encoder()( + input_values=torch.from_numpy(real_audio_inputs["input_values"]), + attention_mask=torch.from_numpy( + real_audio_inputs["attention_mask"].astype(np.int64) + ), + ) + actual = _run_encoder( + package, + real_audio_inputs["input_values"], + real_audio_inputs["attention_mask"], + ) + + np.testing.assert_array_equal( + actual["encoder_attention_mask"].astype(bool), expected.attention_mask.numpy() + ) + assert_logits_close( + actual["encoder_hidden_states"], + expected.last_hidden_state.numpy(), + rtol=1e-3, + atol=1e-3, + ) + + def test_real_decoder_prefill_from_onnx_encoder(self, real_models, real_audio_inputs): + """Chain ONNX encoder -> ONNX decoder and compare with the full HF model.""" + hf_model, _processor, package, config = real_models + start_id = config.decoder_start_token_id + decoder_input_ids = np.array([[start_id]], dtype=np.int64) + + with torch.no_grad(): + expected = hf_model( + input_values=torch.from_numpy(real_audio_inputs["input_values"]), + attention_mask=torch.from_numpy( + real_audio_inputs["attention_mask"].astype(np.int64) + ), + decoder_input_ids=torch.from_numpy(decoder_input_ids), + ) + + encoder_outputs = _run_encoder( + package, + real_audio_inputs["input_values"], + real_audio_inputs["attention_mask"], + ) + actual = _run_decoder( + package, + config, + decoder_input_ids, + encoder_outputs["encoder_hidden_states"], + encoder_outputs["encoder_attention_mask"], + ) + assert_logits_close(actual["logits"], expected.logits.numpy(), rtol=1e-3, atol=1e-3) + + def test_golden_reference_guard_blocks_encoder_mutation(self, real_models): + """The golden generator's guard really stops the in-place ``+=``. + + Without it the reference decoder rewrites the caller's encoder context + on every step, so this asserts the unguarded mutation exists (otherwise + the guard would be silently vacuous) and that the guard removes it. + """ + import sys + + sys.path.insert(0, str(Path(__file__).parent.parent / "scripts")) + try: + from generate_golden import non_mutating_encoder_context + finally: + sys.path.pop(0) + + hf_model, _processor, _package, _config = real_models + with torch.no_grad(): + encoder = hf_model.get_encoder()( + input_values=torch.from_numpy(np.zeros((1, 80 * 40), dtype=np.float32) + 0.01), + attention_mask=torch.ones((1, 80 * 40), dtype=torch.long), + ) + pristine = encoder.last_hidden_state.clone() + + def decode_once(states): + with torch.no_grad(): + hf_model.model.decoder( + input_ids=torch.tensor([[1]], dtype=torch.long), + encoder_hidden_states=states, + encoder_attention_mask=encoder.attention_mask, + use_cache=True, + ) + + unguarded = pristine.clone() + decode_once(unguarded) + assert not torch.equal(unguarded, pristine), ( + "upstream decoder is expected to mutate encoder_hidden_states in place" + ) + + guarded = pristine.clone() + with non_mutating_encoder_context(hf_model): + decode_once(guarded) + assert torch.equal(guarded, pristine) + + def test_decoder_adds_encoder_context_exactly_once(self, real_models, real_audio_inputs): + """The ONNX decoder must not accumulate encoder context across steps. + + HuggingFace's decoder adds ``pos_emb`` to ``encoder_hidden_states`` in + place, so a decoder that re-used a mutated buffer would drift. The + exported graph takes the encoder output as a read-only input and adds + the position table on every call, so replaying step 0 after a full + decode must reproduce the original step-0 logits bit for bit, and the + input buffer must come back untouched. + """ + _hf_model, _processor, package, config = real_models + encoder_outputs = _run_encoder( + package, + real_audio_inputs["input_values"], + real_audio_inputs["attention_mask"], + ) + encoder_hidden_states = encoder_outputs["encoder_hidden_states"] + pristine = encoder_hidden_states.copy() + + def step_zero_logits(): + session = OnnxModelSession(package["decoder"]) + try: + return session.run( + _decoder_feeds( + config, + np.array([[config.decoder_start_token_id]], dtype=np.int64), + encoder_hidden_states, + encoder_outputs["encoder_attention_mask"], + ) + )["logits"] + finally: + session.close() + + first = step_zero_logits() + _onnx_greedy_generate( + package, + config, + encoder_outputs, + max_new_tokens=20, + start_id=config.decoder_start_token_id, + ) + replayed = step_zero_logits() + + np.testing.assert_array_equal(encoder_hidden_states, pristine) + np.testing.assert_array_equal(replayed, first) + + def test_stepwise_logit_parity_with_nonmutating_reference( + self, real_models, real_audio_inputs + ): + """Every decode step matches a non-mutating HuggingFace reference. + + Token-level agreement can hide logit drift, so this compares the full + vocabulary distribution at every step against a reference that is + handed a pristine encoder context each call. + """ + hf_model, processor, package, config = real_models + encoder_outputs = _run_encoder( + package, + real_audio_inputs["input_values"], + real_audio_inputs["attention_mask"], + ) + + with torch.no_grad(): + expected_encoder = hf_model.get_encoder()( + input_values=torch.from_numpy(real_audio_inputs["input_values"]), + attention_mask=torch.from_numpy( + real_audio_inputs["attention_mask"].astype(np.int64) + ), + ) + expected_tokens, expected_logits = _nonmutating_hf_greedy( + hf_model, + expected_encoder.last_hidden_state, + expected_encoder.attention_mask, + steps=50, + ) + + session = OnnxModelSession(package["decoder"]) + try: + cache = _empty_cache_feeds(config) + current = np.array([[config.decoder_start_token_id]], dtype=np.int64) + actual_tokens: list[int] = [] + for step in range(len(expected_logits)): + outputs = session.run( + _decoder_feeds( + config, + current, + encoder_outputs["encoder_hidden_states"], + encoder_outputs["encoder_attention_mask"], + past_key_values=cache, + position_offset=step, + ) + ) + assert_logits_close( + outputs["logits"][0, -1], + expected_logits[step], + rtol=1e-3, + atol=1e-3, + ) + token = int(outputs["logits"][0, -1].argmax()) + actual_tokens.append(token) + if token == _EOS_TOKEN_ID: + break + current = np.array([[token]], dtype=np.int64) + cache = { + f"past_key_values.{layer_idx}.{kind}": outputs[ + f"present.{layer_idx}.{kind}" + ] + for layer_idx in range(config.num_hidden_layers) + for kind in ("key", "value") + } + finally: + session.close() + + assert len(actual_tokens) == len(expected_tokens) + assert actual_tokens == expected_tokens + + # The committed L5 golden must describe this same non-mutating decode. + content = ( + expected_tokens[:-1] + if expected_tokens and expected_tokens[-1] == _EOS_TOKEN_ID + else expected_tokens + ) + golden = load_generation_golden( + next( + case for case in discover_test_cases(level="L5") if case.model_id == _MODEL_ID + ) + ) + assert len(content) == len(golden) + assert content == golden + assert processor.decode(content, skip_special_tokens=True) + + def test_float16_cpu_parity_and_transcript(self, real_audio_inputs, real_models): + """fp16 keeps full-logit parity and an identical transcript on ORT CPU. + + fp16 on CUDA is semantically exact too, but ORT's CUDA fp16 fused + attention kernel zeroes the first encoder frame — the sparsest masked + query row, which sees only 4 of 456 keys under the ``(16, 4)`` window. + That is an execution-provider defect, not a graph defect: the identical + graph is accurate here on CPU, and fp32/bf16 are accurate on CUDA. + """ + hf_model, processor, _package, _config = real_models + package = build(_MODEL_ID, dtype="f16", load_weights=True, revision=_REVISION) + config = package.config + + session = OnnxModelSession(package["encoder"]) + try: + encoder_outputs = session.run( + { + "input_values": real_audio_inputs["input_values"].astype(np.float16), + "attention_mask": real_audio_inputs["attention_mask"].astype(np.int64), + } + ) + finally: + session.close() + + with torch.no_grad(): + expected = hf_model.get_encoder()( + input_values=torch.from_numpy(real_audio_inputs["input_values"]), + attention_mask=torch.from_numpy( + real_audio_inputs["attention_mask"].astype(np.int64) + ), + ) + actual_hidden = np.asarray(encoder_outputs["encoder_hidden_states"], dtype=np.float32) + np.testing.assert_array_equal( + encoder_outputs["encoder_attention_mask"].astype(bool), + expected.attention_mask.numpy(), + ) + assert not np.isnan(actual_hidden).any() + np.testing.assert_allclose( + actual_hidden, expected.last_hidden_state.numpy(), rtol=1e-2, atol=1e-1 + ) + + tokens = _onnx_greedy_generate( + package, + config, + encoder_outputs, + max_new_tokens=50, + start_id=config.decoder_start_token_id, + ) + content = tokens[:-1] if tokens and tokens[-1] == _EOS_TOKEN_ID else tokens + golden = load_generation_golden( + next( + case for case in discover_test_cases(level="L5") if case.model_id == _MODEL_ID + ) + ) + assert len(content) == len(golden) + assert content == golden + assert processor.decode(content, skip_special_tokens=True) + + def test_real_generation_matches_huggingface(self, real_models, real_audio_inputs): + """Full ONNX pipeline transcribes identically to HuggingFace generate().""" + hf_model, processor, package, config = real_models + encoder_outputs = _run_encoder( + package, + real_audio_inputs["input_values"], + real_audio_inputs["attention_mask"], + ) + actual_tokens = _onnx_greedy_generate( + package, + config, + encoder_outputs, + max_new_tokens=50, + start_id=config.decoder_start_token_id, + ) + + with torch.no_grad(): + generated = hf_model.generate( + input_values=torch.from_numpy(real_audio_inputs["input_values"]), + attention_mask=torch.from_numpy( + real_audio_inputs["attention_mask"].astype(np.int64) + ), + max_new_tokens=50, + do_sample=False, + ) + expected_tokens = generated[0].tolist() + if expected_tokens and expected_tokens[0] == config.decoder_start_token_id: + expected_tokens = expected_tokens[1:] + + assert len(actual_tokens) == len(expected_tokens) + assert actual_tokens == expected_tokens + transcript = processor.decode(actual_tokens, skip_special_tokens=True) + assert transcript == processor.decode(expected_tokens, skip_special_tokens=True) + assert len(transcript.split()) > 5, transcript