diff --git a/docs/source/openvino/models.mdx b/docs/source/openvino/models.mdx index 902dd2ad75..db24f15496 100644 --- a/docs/source/openvino/models.mdx +++ b/docs/source/openvino/models.mdx @@ -38,6 +38,7 @@ Here is the list of the supported architectures : - CodeGen2 - Cohere - Cohere2 +- Cohere ASR - ConvBERT - ConvNeXt - DBRX @@ -216,4 +217,4 @@ Here is the list of the supported architectures : - Qwen3-*-DFlash - Qwen3.5-*-DFlash - Qwen3.6-*-DFlash -- Qwen3-Coder-30B-A3B-DFlash \ No newline at end of file +- Qwen3-Coder-30B-A3B-DFlash diff --git a/optimum/exporters/openvino/input_generators.py b/optimum/exporters/openvino/input_generators.py index 34f75923ba..e78792ba43 100644 --- a/optimum/exporters/openvino/input_generators.py +++ b/optimum/exporters/openvino/input_generators.py @@ -1396,6 +1396,92 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int ) +class CohereAsrDummyAudioInputGenerator(DummyAudioInputGenerator): + """Dummy input generator for the Cohere ASR Conformer encoder. + + The feature extractor pads to the longest sample of the batch instead of a fixed 30s window, + so the encoder takes a companion `length` input holding the real frame count per sample.""" + + SUPPORTED_INPUT_NAMES = ("input_features", "input_values", "length") + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name != "length": + return super().generate(input_name, framework=framework, int_dtype=int_dtype, float_dtype=float_dtype) + + # Dummy batches are built without padding, so every sample covers all nb_max_frames frames + return self.random_int_tensor( + shape=[self.batch_size], + min_value=self.nb_max_frames, + max_value=self.nb_max_frames + 1, + framework=framework, + dtype=int_dtype, + ) + + +class CohereAsrDummySeq2SeqDecoderTextInputGenerator(DummySeq2SeqDecoderTextInputGenerator): + """Dummy input generator for the Cohere ASR decoder. + + Encoder states reach the decoder before `encoder_decoder_proj` is applied, so `encoder_outputs` + has to be sized with the Conformer `d_model` and not with the decoder hidden size.""" + + def __init__(self, task, normalized_config, **kwargs): + super().__init__(task, normalized_config, **kwargs) + self.hidden_size = normalized_config.encoder_hidden_size + + +class CohereAsrNativeDummyAudioInputGenerator(DummyAudioInputGenerator): + """Dummy input generator for the transformers native Cohere ASR encoder. + + Features are time major here rather than laid out as in whisper, and padded batches are + described by a frame level `attention_mask` instead of by an explicit frame count.""" + + SUPPORTED_INPUT_NAMES = ("input_features", "attention_mask") + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "attention_mask": + return self.random_int_tensor( + shape=[self.batch_size, self.nb_max_frames], + min_value=1, + max_value=2, + framework=framework, + dtype=int_dtype, + ) + + return self.random_float_tensor( + shape=[self.batch_size, self.nb_max_frames, self.feature_size], + min_value=-1, + max_value=1, + framework=framework, + dtype=float_dtype, + ) + + +class CohereAsrNativeDummySeq2SeqDecoderTextInputGenerator(DummySeq2SeqDecoderTextInputGenerator): + """Dummy input generator for the transformers native Cohere ASR decoder. + + Encoder states are projected only once they are already inside the decoder, so `encoder_outputs` + has to be sized with the Parakeet width, and cross attention runs against the subsampled states + so the mask that reaches the decoder is shorter than the frame level one.""" + + SUPPORTED_INPUT_NAMES = DummySeq2SeqDecoderTextInputGenerator.SUPPORTED_INPUT_NAMES + ("attention_mask",) + + def __init__(self, task, normalized_config, **kwargs): + super().__init__(task, normalized_config, **kwargs) + self.hidden_size = normalized_config.encoder_hidden_size + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "attention_mask": + return self.random_int_tensor( + shape=[self.batch_size, self.sequence_length], + min_value=1, + max_value=2, + framework=framework, + dtype=int_dtype, + ) + + return super().generate(input_name, framework=framework, int_dtype=int_dtype, float_dtype=float_dtype) + + class Qwen3ASRDummySeq2SeqPastKeyValuesGenerator(DummySeq2SeqPastKeyValuesGenerator): """Custom KV cache generator for Qwen3-ASR with GQA (num_key_value_heads != num_attention_heads). Qwen3-ASR has no cross-attention, so only self-attention KV cache is generated (2 per layer).""" diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index 1bb72ff0e4..aac12b0ae3 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -39,6 +39,10 @@ from optimum.exporters.openvino.input_generators import ( AquilaDummyPastKeyValuesGenerator, ChatGLM2DummyPastKeyValuesGenerator, + CohereAsrDummyAudioInputGenerator, + CohereAsrDummySeq2SeqDecoderTextInputGenerator, + CohereAsrNativeDummyAudioInputGenerator, + CohereAsrNativeDummySeq2SeqDecoderTextInputGenerator, DeciDummyPastKeyValuesGenerator, DummyAudioPhi4MMInputGenerator, DummyFluxTextInputGenerator, @@ -104,6 +108,8 @@ BloomModelPatcher, ChatGLMModelPatcher, CodeGenModelPatcher, + CohereAsrModelPatcher, + CohereAsrNativeModelPatcher, CommonImageEmbeddingsModelPatcher, DBRXModelPatcher, DeciLMModelPatcher, @@ -351,6 +357,15 @@ def init_model_configs(): "Qwen3OmniMoeForConditionalGeneration", ) + TasksManager._CUSTOM_CLASSES[("pt", "cohere_asr", "automatic-speech-recognition")] = ( + "transformers", + "AutoModelForSpeechSeq2Seq", + ) + TasksManager._CUSTOM_CLASSES[("pt", "cohere_asr", "automatic-speech-recognition-with-past")] = ( + "transformers", + "AutoModelForSpeechSeq2Seq", + ) + if is_diffusers_available() and "fill" not in TasksManager._DIFFUSERS_TASKS_TO_MODEL_LOADERS: TasksManager._DIFFUSERS_TASKS_TO_MODEL_LOADERS["fill"] = "FluxFillPipeline" TasksManager._DIFFUSERS_TASKS_TO_MODEL_MAPPINGS["fill"] = {"flux": "FluxFillPipeline"} @@ -4358,6 +4373,216 @@ def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], dire inputs_or_outputs[f"{name}.{i}.decoder.value"] = {0: "batch_size", 2: decoder_sequence_name} +@register_in_tasks_manager( + "cohere_asr", + *[ + "automatic-speech-recognition", + "automatic-speech-recognition-with-past", + ], + library_name="transformers", +) +class CohereAsrOpenVINOConfig(AudioToTextOpenVINOConfig): + """Picks the export config matching the Cohere ASR variant that was loaded. + + The transformers native model and the Hub modeling file share the `cohere_asr` model type but + have nothing else in common, and only the native one carries a nested `encoder_config`.""" + + def __new__(cls, config: "PretrainedConfig", *args, **kwargs): + if cls is CohereAsrOpenVINOConfig: + variant = ( + CohereAsrNativeOpenVINOConfig + if getattr(config, "encoder_config", None) is not None + else CohereAsrRemoteOpenVINOConfig + ) + return super().__new__(variant) + return super().__new__(cls) + + +class CohereAsrNativeOpenVINOConfig(CohereAsrOpenVINOConfig): + """Export config for the transformers native `CohereAsrForConditionalGeneration`. + + The Parakeet encoder takes time major features next to a frame level mask, and returns that mask + subsampled because cross attention runs against the shortened states.""" + + DUMMY_INPUT_GENERATOR_CLASSES = ( + CohereAsrNativeDummyAudioInputGenerator, + CohereAsrNativeDummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + ) + + _MODEL_PATCHER = CohereAsrNativeModelPatcher + + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( + encoder_num_layers="encoder_layers", + decoder_num_layers="num_hidden_layers", + hidden_size="hidden_size", + num_attention_heads="num_attention_heads", + decoder_num_attention_heads="num_key_value_heads", + feature_size="num_mel_bins", + # Parakeet width, kept apart from the decoder hidden size above because the encoder states + # are projected only once they are already inside the decoder + encoder_hidden_size="encoder_hidden_size", + allow_new=True, + ) + + def __init__( + self, + config: "PretrainedConfig", + task: str = "automatic-speech-recognition", + int_dtype: str = "int64", + float_dtype: str = "fp32", + preprocessors: Optional[List[Any]] = None, + **kwargs, + ): + # The decoder settings sit on the top level config while the encoder keeps its own section, + # and NormalizedConfig resolves every field on the top level + encoder_config = config.encoder_config + config.encoder_layers = encoder_config.num_hidden_layers + config.num_mel_bins = encoder_config.num_mel_bins + config.encoder_hidden_size = encoder_config.hidden_size + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + preprocessors=preprocessors, + **kwargs, + ) + + def _create_dummy_input_generator_classes(self, **kwargs) -> List["DummyInputGenerator"]: + generators = super()._create_dummy_input_generator_classes(**kwargs) + if self._behavior is ConfigBehavior.DECODER: + # Both generators can produce an `attention_mask`, and past the encoder it has to follow + # the subsampled length, so the decoder side one has to be picked first + generators.sort(key=lambda generator: not isinstance(generator, DummySeq2SeqDecoderTextInputGenerator)) + return generators + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + common_inputs = {} + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + # Time major, unlike the whisper style layout the shared audio config describes + common_inputs["input_features"] = {0: "batch_size", 1: "encoder_sequence_length", 2: "feature_size"} + else: + common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} + + common_inputs["attention_mask"] = {0: "batch_size", 1: "encoder_sequence_length"} + + if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: + common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} + if self.use_past_in_inputs: + self.add_past_key_values(common_inputs, direction="inputs") + + return common_inputs + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + common_outputs = super().outputs + if self._behavior is ConfigBehavior.ENCODER: + # Subsampled by the convolutional front end, so the decoder cannot reuse the frame mask. + # Named apart from the encoder input of the same meaning to keep the export from + # disambiguating it into `attention_mask_1` + common_outputs["encoder_attention_mask"] = {0: "batch_size", 1: "encoder_sequence_length"} + return common_outputs + + @property + def torch_to_ov_output_map(self) -> Dict[str, str]: + return {"attention_mask": "encoder_attention_mask"} + + +class CohereAsrRemoteOpenVINOConfig(CohereAsrOpenVINOConfig): + """Export config for the Hub modeling file (Conformer encoder, transformer decoder). + + The encoder consumes variable length features plus a `length` input, so this derives from the + generic audio config rather than from Whisper, whose overrides pin a static 3000 frame shape.""" + + DUMMY_INPUT_GENERATOR_CLASSES = ( + CohereAsrDummyAudioInputGenerator, + CohereAsrDummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + ) + + _MODEL_PATCHER = CohereAsrModelPatcher + + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( + encoder_num_layers="encoder_layers", + decoder_num_layers="decoder_layers", + hidden_size="hidden_size", + num_attention_heads="num_attention_heads", + decoder_num_attention_heads="num_key_value_heads", + feature_size="num_mel_bins", + # Conformer d_model, kept apart from the decoder hidden size above because the decoder + # dummy inputs are sized before encoder_decoder_proj runs + encoder_hidden_size="d_model", + allow_new=True, + ) + + @staticmethod + def _read_field(source, field, default=None): + """Read `field` from a mapping or from an object, falling back to `default`""" + if source is None: + return default + if isinstance(source, dict): + return source.get(field, default) + return getattr(source, field, default) + + @classmethod + def _flatten_nested_config(cls, config: "PretrainedConfig"): + # The checkpoint keeps encoder and decoder settings in nested NeMo style sections, while + # NormalizedConfig resolves every field on the top level config + read_field = cls._read_field + + encoder_config = getattr(config, "encoder", None) + if encoder_config is not None: + config.encoder_layers = read_field(encoder_config, "n_layers") + config.num_mel_bins = read_field(encoder_config, "feat_in") + config.d_model = read_field(encoder_config, "d_model") + + decoder_config = read_field(getattr(config, "transf_decoder", None), "config_dict") + if decoder_config is not None: + config.decoder_layers = read_field(decoder_config, "num_layers") + config.hidden_size = read_field(decoder_config, "hidden_size") + config.num_attention_heads = read_field(decoder_config, "num_attention_heads") + config.num_key_value_heads = read_field( + decoder_config, "num_key_value_heads", read_field(decoder_config, "num_attention_heads") + ) + config.vocab_size = getattr(config, "vocab_size", None) or read_field(decoder_config, "vocab_size") + + preprocessor_config = getattr(config, "preprocessor", None) + if getattr(config, "num_mel_bins", None) is None and preprocessor_config is not None: + config.num_mel_bins = read_field(preprocessor_config, "features") + + if getattr(config, "decoder_start_token_id", None) is None: + config.decoder_start_token_id = 0 + + def __init__( + self, + config: "PretrainedConfig", + task: str = "automatic-speech-recognition", + int_dtype: str = "int64", + float_dtype: str = "fp32", + preprocessors: Optional[List[Any]] = None, + **kwargs, + ): + self._flatten_nested_config(config) + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + preprocessors=preprocessors, + **kwargs, + ) + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + common_inputs = super().inputs + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + # Real frame count per sample, needed by ConvSubsampling to mask the padded tail + common_inputs["length"] = {0: "batch_size"} + return common_inputs + + @register_in_tasks_manager( "t5", *["feature-extraction", "feature-extraction-with-past", "text2text-generation", "text2text-generation-with-past"], diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index a93cf65a7d..d712af87eb 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -16,6 +16,7 @@ import inspect import logging import math +import sys import types from dataclasses import dataclass from types import SimpleNamespace @@ -10244,6 +10245,206 @@ def __exit__(self, exc_type, exc_value, traceback): del attn._orig_forward +class CohereAsrModelPatcher(OVSeq2SeqModelPatcher): + """Makes the Cohere ASR remote code traceable. + + Two spots break `torch.jit.trace`: the relative positional encoding caches its table behind + `torch._dynamo.disable`, and the cache length helper casts the kv length to a Python int.""" + + def __enter__(self): + super().__enter__() + self._patched_positional_encodings = [] + self._patch_positional_encoding() + self._patch_cache_seq_length() + self._patch_decoder_inputs() + + def _patch_positional_encoding(self): + # Matched by class name so the patch also covers the encoder only submodule that the + # encoder export traces + for positional_encoding in self._model.modules(): + if type(positional_encoding).__name__ != "RelPositionalEncoding": + continue + positional_encoding._orig_materialize_pe = positional_encoding._materialize_pe + positional_encoding._materialize_pe = self._build_materialize_pe(positional_encoding) + self._patched_positional_encodings.append(positional_encoding) + + @staticmethod + def _build_materialize_pe(positional_encoding): + # Rebuilding the table on every call, instead of reusing a cached one, keeps it a pure + # function of the arguments so trace sees identical constants on its verification pass + def materialize_pe(length, device, dtype): + effective_length = max(length, getattr(positional_encoding, "max_len", length)) + positions = torch.arange( + effective_length - 1, -effective_length, -1, dtype=torch.float32, device=device + ).unsqueeze(1) + positional_embeddings = positional_encoding._create_pe(positions=positions, dtype=dtype) + if hasattr(positional_encoding, "pe"): + positional_encoding.pe = positional_embeddings + else: + positional_encoding.register_buffer("pe", positional_embeddings, persistent=False) + + return materialize_pe + + @staticmethod + def _cache_seq_length(past_key_values): + if past_key_values is None: + return 0 + if hasattr(past_key_values, "get_seq_length"): + return past_key_values.get_seq_length() + # The exporter feeds the legacy layout as a list, which the remote helper only accepts as a + # tuple, so without this the length collapses to a constant zero in the traced graph + if isinstance(past_key_values, (list, tuple)) and past_key_values: + return past_key_values[0][0].shape[-2] + return 0 + + def _patch_cache_seq_length(self): + # The remote helper casts the cache length to int, which freezes the kv length as a graph + # constant and leaves the exported decoder valid only for the traced step count + remote_module = sys.modules.get(type(self._model).__module__) + if remote_module is None or not hasattr(remote_module, "_get_cache_seq_length"): + return + + self._cache_seq_length_module = remote_module + self._orig_get_cache_seq_length = remote_module._get_cache_seq_length + remote_module._get_cache_seq_length = self._cache_seq_length + + def _patch_decoder_inputs(self): + # `positions` is not part of the exported decoder signature, and the remote code defaults it + # to arange(tgt_len), which restarts at zero on every cached step and makes the decoder loop + self._forward_without_positions = self._model.forward + + @functools.wraps(self._forward_without_positions) + def forward_with_running_positions(*args, **kwargs): + decoder_ids = kwargs.get("input_ids") + if decoder_ids is None: + decoder_ids = kwargs.get("decoder_input_ids") + past_key_values = kwargs.get("past_key_values") + if decoder_ids is not None and past_key_values is None and self.real_config._behavior == "decoder": + # The remote code only fills a cache it is handed, so the cacheless decoder would + # export no `present` outputs and leave the second submodel with an empty cache + past_key_values = EncoderDecoderCache(DynamicCache(), DynamicCache()) + kwargs["past_key_values"] = past_key_values + if kwargs.get("positions") is None and decoder_ids is not None and past_key_values is not None: + # Reuses the cache length the remote code already resolves dynamically, so the + # offset stays a graph value instead of the traced step index + past_length = self._cache_seq_length(past_key_values) + positions = torch.arange(past_length, past_length + decoder_ids.shape[1], device=decoder_ids.device) + # Kept at batch one so the embedding add broadcasts, otherwise the traced batch size + # would be frozen and beam search could not widen it + kwargs["positions"] = positions.unsqueeze(0) + return self._forward_without_positions(*args, **kwargs) + + self._model.forward = forward_with_running_positions + + def __exit__(self, exc_type, exc_value, traceback): + super().__exit__(exc_type, exc_value, traceback) + self._forward_without_positions = None + + for positional_encoding in getattr(self, "_patched_positional_encodings", []): + if hasattr(positional_encoding, "_orig_materialize_pe"): + positional_encoding._materialize_pe = positional_encoding._orig_materialize_pe + del positional_encoding._orig_materialize_pe + + remote_module = getattr(self, "_cache_seq_length_module", None) + if remote_module is not None and hasattr(self, "_orig_get_cache_seq_length"): + remote_module._get_cache_seq_length = self._orig_get_cache_seq_length + del self._orig_get_cache_seq_length + self._cache_seq_length_module = None + + +class CohereAsrNativeModelPatcher(OVSeq2SeqModelPatcher): + """Reconnects the encoder mask for the transformers native Cohere ASR decoder. + + The model reads the subsampled mask off the encoder output, but the decoder submodel is traced + with the encoder states handed in as a bare tensor, so the mask has to be reattached.""" + + def __enter__(self): + super().__enter__() + self._strip_forward_annotations() + self._patch_inner_forward() + self._patch_encoder_attention() + + def _patch_encoder_attention(self): + # A padded frame masks its own attention row out completely, and the softmax over that row + # comes back as NaN here where the fused torch kernel returns zeros instead + inner_model = getattr(self._model, "model", None) + # The encoder submodel is traced on its own, so it is already the model being patched + encoder = getattr(inner_model, "encoder", None) if inner_model is not None else self._model + self._encoder_attentions = [layer.self_attn for layer in getattr(encoder, "layers", [])] + + for attention in self._encoder_attentions: + attention._forward_with_masked_rows = attention.forward + attention.forward = functools.partial(self._attention_without_masked_rows, attention) + + @staticmethod + def _attention_without_masked_rows(attention, hidden_states, *args, attention_mask=None, **kwargs): + attn_output, attn_weights = attention._forward_with_masked_rows( + hidden_states, *args, attention_mask=attention_mask, **kwargs + ) + if attention_mask is not None: + # Zeroed rather than left as they are so the depthwise convolution downstream, the only + # step that mixes across time, keeps reading padded frames as empty + fully_masked_rows = attention_mask.any(dim=-1).logical_not().transpose(1, 2) + attn_output = attn_output.masked_fill(fully_masked_rows, 0.0) + return attn_output, attn_weights + + def _strip_forward_annotations(self): + # The tracing wrapper is built by executing the forward signature as source, so an + # annotation naming a type it does not import drops the export to positional arguments and + # the shared seq2seq patcher can then no longer find the cache it is meant to convert + signature = inspect.signature(self._model.forward) + self._model.forward.__signature__ = signature.replace( + parameters=[ + parameter.replace(annotation=inspect.Parameter.empty) for parameter in signature.parameters.values() + ] + ) + + def _patch_inner_forward(self): + # Patched on the inner model because that is the first place the arguments carry names, and + # also where the mask is read off the encoder output the decoder submodel never produces + self._inner_model = getattr(self._model, "model", None) + if self._inner_model is None: + return + + self._inner_forward = self._inner_model.forward + + @functools.wraps(self._inner_forward) + def patched_inner_forward(*args, **kwargs): + encoder_outputs = kwargs.get("encoder_outputs") + if encoder_outputs is not None: + # Bypassing the encoder leaves `attention_mask` carrying the subsampled mask rather + # than the frame level one, which is exactly what cross attention needs + kwargs["encoder_outputs"] = self._as_encoder_output(encoder_outputs, kwargs.get("attention_mask")) + if kwargs.get("use_cache") is None: + # Never filled in from the config here, and left unset the cache is dropped from the + # outputs, so the exported decoder would have nothing to hand to the next step + kwargs["use_cache"] = True + return self._inner_forward(*args, **kwargs) + + self._inner_model.forward = patched_inner_forward + + @staticmethod + def _as_encoder_output(encoder_outputs, encoder_mask): + # The dummy generator hands the states over as a tuple, which the model cannot read the + # mask off, so they are rewrapped instead of being passed straight through + hidden_states = getattr(encoder_outputs, "last_hidden_state", None) + if hidden_states is None: + hidden_states = encoder_outputs[0] if isinstance(encoder_outputs, (list, tuple)) else encoder_outputs + + rewrapped = BaseModelOutput(last_hidden_state=hidden_states) + rewrapped.attention_mask = encoder_mask + return rewrapped + + def __exit__(self, exc_type, exc_value, traceback): + super().__exit__(exc_type, exc_value, traceback) + for attention in getattr(self, "_encoder_attentions", []): + attention.forward = attention._forward_with_masked_rows + self._encoder_attentions = [] + if getattr(self, "_inner_model", None) is not None: + self._inner_model.forward = self._inner_forward + self._inner_model = None + + class KokoroModelPatcher(ModelPatcher): """ Patches the Kokoro TTS model for OpenVINO export by redirecting forward diff --git a/optimum/intel/openvino/modeling_seq2seq.py b/optimum/intel/openvino/modeling_seq2seq.py index 642100260e..6bd57cb4fa 100644 --- a/optimum/intel/openvino/modeling_seq2seq.py +++ b/optimum/intel/openvino/modeling_seq2seq.py @@ -685,6 +685,12 @@ def forward( # get decoder inputs from shifting lm labels to the right decoder_input_ids = self._shift_right(labels) + # An encoder that changes the length of the sequence reports the mask that matches its own + # output, and the incoming one no longer lines up with the states cross attention sees + encoder_attention_mask = getattr(encoder_outputs, "attention_mask", None) + if encoder_attention_mask is None: + encoder_attention_mask = attention_mask + # Decode if past_key_values is None or self.decoder_with_past is None: decoder_outputs = self.decoder( @@ -693,7 +699,7 @@ def forward( ), attention_mask=decoder_attention_mask, encoder_hidden_states=encoder_outputs.last_hidden_state, - encoder_attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, past_key_values=past_key_values, cache_position=cache_position, ) @@ -702,7 +708,7 @@ def forward( input_ids=decoder_input_ids[:, -1:], # Cut decoder_input_ids if past is used attention_mask=decoder_attention_mask, encoder_hidden_states=encoder_outputs.last_hidden_state, - encoder_attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, past_key_values=past_key_values, cache_position=cache_position, ) @@ -750,7 +756,9 @@ def _reshape(self, model: openvino.Model, batch_size: int, sequence_length: int, elif is_decoder and not inputs.get_any_name().startswith("encoder"): if not inputs.get_any_name().startswith("beam_idx"): shapes[inputs][1] = -1 - else: + # Rank 1 encoder inputs such as the cohere_asr `length` only carry a batch dimension, + # which the assignment above already covers + elif len(shapes[inputs]) > 1: shapes[inputs][1] = sequence_length model.reshape(shapes) return model @@ -904,6 +912,9 @@ def forward( self, input_ids: torch.LongTensor = None, attention_mask: torch.LongTensor = None, + # Listed explicitly rather than picked up from kwargs because generate() inspects this + # signature when it decides which encoder arguments to forward + length: Optional[torch.LongTensor] = None, **kwargs, ) -> BaseModelOutput: self.compile() @@ -917,6 +928,20 @@ def forward( attention_mask = torch.ones_like(inputs[self.main_input_name]) inputs["attention_mask"] = attention_mask + # Frame count per sample, used by the Conformer encoder to mask the padded tail + if "length" in self.input_names: + if length is None: + # Same default as the eager ConformerEncoder, so callers that bypass generate() + # get the full time dimension for every sample. The traced graph always wants it + encoder_features = inputs[self.main_input_name] + length = torch.full( + (encoder_features.shape[0],), + encoder_features.shape[-1], + dtype=torch.int64, + device=getattr(encoder_features, "device", None), + ) + inputs["length"] = length + # Qwen3-ASR requires input_features chunking before passing to encoder for processing of long audios. if getattr(self.config, "model_type", None) == "qwen3_asr": input_features = inputs["input_features"] @@ -928,11 +953,18 @@ def forward( return BaseModelOutput(last_hidden_state=audio_features) # Run inference - last_hidden_state = torch.from_numpy( - self.request(inputs, share_inputs=True, share_outputs=True)["last_hidden_state"] - ).to(self.device) + request_outputs = self.request(inputs, share_inputs=True, share_outputs=True) + last_hidden_state = torch.from_numpy(request_outputs["last_hidden_state"]).to(self.device) + encoder_outputs = BaseModelOutput(last_hidden_state=last_hidden_state) + + # Encoders that subsample the time axis hand back the mask that goes with the shortened + # states, which is the one cross attention needs rather than the one that came in + if "encoder_attention_mask" in self.output_names: + encoder_outputs.attention_mask = torch.from_numpy(request_outputs["encoder_attention_mask"]).to( + self.device + ) - return BaseModelOutput(last_hidden_state=last_hidden_state) + return encoder_outputs class OVDecoder(OVModelPart): @@ -1613,3 +1645,21 @@ def _get_logits_processor(self, generation_config: GenerationConfig, *args, **kw logits_processor = super()._get_logits_processor(generation_config, *args, **kwargs) generation_config.forced_decoder_ids = forced_decoder_ids return logits_processor + + +def _register_ov_speech_seq2seq_for_pipeline_autodetection(): + """Let transformers pipelines recognise `OVModelForSpeechSeq2Seq` as a seq2seq ASR model. + + `AutomaticSpeechRecognitionPipeline` classifies models by class name and treats anything it + does not know as CTC, which calls the model directly instead of going through generate().""" + try: + from transformers.models.auto.modeling_auto import MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES + except ImportError: + return + + # Whisper is matched on model_type instead, so only the other architectures need this entry + if OVModelForSpeechSeq2Seq.__name__ not in MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES.values(): + MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES["_optimum_intel_ov_speech_seq2seq"] = OVModelForSpeechSeq2Seq.__name__ + + +_register_ov_speech_seq2seq_for_pipeline_autodetection() diff --git a/tests/openvino/test_export.py b/tests/openvino/test_export.py index d9928449a3..6d90aa01f7 100644 --- a/tests/openvino/test_export.py +++ b/tests/openvino/test_export.py @@ -112,6 +112,7 @@ class ExportModelTest(unittest.TestCase): "lfm2_moe": OVModelForCausalLM, "qwen3_asr": OVModelForSpeechSeq2Seq, "fun_asr": OVModelForSpeechSeq2Seq, + "cohere_asr": OVModelForSpeechSeq2Seq, "mamba": OVModelForCausalLM, "falcon_mamba": OVModelForCausalLM, "gemma4": OVModelForVisualCausalLM, diff --git a/tests/openvino/test_exporters_cli.py b/tests/openvino/test_exporters_cli.py index c57789a4c5..96c0913ce4 100644 --- a/tests/openvino/test_exporters_cli.py +++ b/tests/openvino/test_exporters_cli.py @@ -136,6 +136,7 @@ class OVCLIExportTestCase(unittest.TestCase): ("text-generation-with-past", "mamba"), ("text-generation-with-past", "falcon_mamba"), ("text-to-image", "flux.2-klein"), + ("automatic-speech-recognition", "cohere_asr"), ] # filter architectures depending on min/max transformers supported versions SUPPORTED_ARCHITECTURES = [ @@ -173,6 +174,7 @@ class OVCLIExportTestCase(unittest.TestCase): "speecht5": 2, "kokoro": 0, # uses g2p, no tokenizer "clip": 2, + "cohere_asr": 2, "mamba": 2, "falcon_mamba": 2, "qwen3": 2, @@ -1110,8 +1112,10 @@ def test_exporters_cli_int8(self, task: str, model_type: str): model = self._load_exported_ov_model(model_type, task, tmpdir, model_kwargs) expected_int8 = _ARCHITECTURES_TO_EXPECTED_INT8[model_type] expected_int8 = {k: {"int8": v} for k, v in expected_int8.items()} - if task.startswith("text2text-generation") and (not task.endswith("with-past") or model.decoder.stateful): - del expected_int8["decoder_with_past"] + if task.startswith(("text2text-generation", "automatic-speech-recognition")) and ( + not task.endswith("with-past") or model.decoder.stateful + ): + expected_int8.pop("decoder_with_past", None) check_compression_state_per_model(self, model.ov_models, expected_int8) @parameterized.expand(SUPPORTED_SD_HYBRID_ARCHITECTURES) diff --git a/tests/openvino/test_seq2seq.py b/tests/openvino/test_seq2seq.py index 371b729b5e..82d635dc56 100644 --- a/tests/openvino/test_seq2seq.py +++ b/tests/openvino/test_seq2seq.py @@ -349,6 +349,8 @@ def test_compare_with_and_without_past_key_values(self): class OVModelForSpeechSeq2SeqIntegrationTest(OVSeq2SeqTestMixin): SUPPORTED_ARCHITECTURES = ("whisper",) + if "cohere_asr" in CONFIG_MAPPING_NAMES: + SUPPORTED_ARCHITECTURES += ("cohere_asr",) OVMODEL_CLASS = OVModelForSpeechSeq2Seq AUTOMODEL_CLASS = AutoModelForSpeechSeq2Seq TASK = "automatic-speech-recognition" @@ -360,6 +362,18 @@ def _generate_random_audio_data(self): audio_data = 0.5 * np.sin(2 * np.pi * 220 * t) return audio_data + def test_pipeline_autodetection_registers_non_whisper_architectures(self): + # Without the registration the ASR pipeline falls back to its CTC branch for every non + # Whisper architecture and calls the model instead of generate() + from transformers.models.auto.modeling_auto import MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES + + self.assertIn( + OVModelForSpeechSeq2Seq.__name__, + MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES.values(), + "OVModelForSpeechSeq2Seq must be registered so that transformers.pipeline routes " + "non-Whisper OpenVINO ASR architectures through generate()", + ) + @parameterized.expand(SUPPORTED_ARCHITECTURES) def test_compare_to_transformers(self, model_arch): set_seed(SEED) @@ -369,7 +383,11 @@ def test_compare_to_transformers(self, model_arch): model_id, export=True, ov_config=F32_CONFIG, device=OPENVINO_DEVICE ) ov_model_stateless = self.OVMODEL_CLASS.from_pretrained( - model_id, export=True, ov_config=F32_CONFIG, stateful=False, device=OPENVINO_DEVICE + model_id, + export=True, + ov_config=F32_CONFIG, + stateful=False, + device=OPENVINO_DEVICE, ) self._check_openvino_model_attributes(ov_model, use_cache=True, stateful=True) self._check_openvino_model_attributes(ov_model_stateless, use_cache=True, stateful=False) @@ -377,7 +395,10 @@ def test_compare_to_transformers(self, model_arch): processor = AutoProcessor.from_pretrained(model_id) data = self._generate_random_audio_data() pt_features = processor.feature_extractor(data, return_tensors="pt") - decoder_start_token_id = transformers_model.config.decoder_start_token_id + # Only describes how the batch was split for long audio, and the models take neither it nor + # any other argument the feature extractor is free to add + pt_features.pop("audio_chunk_index", None) + decoder_start_token_id = getattr(transformers_model.config, "decoder_start_token_id", None) or 0 decoder_inputs = {"decoder_input_ids": torch.ones((1, 1), dtype=torch.long) * decoder_start_token_id} with torch.no_grad(): @@ -385,6 +406,7 @@ def test_compare_to_transformers(self, model_arch): for input_type in ["pt", "np"]: features = processor.feature_extractor(data, return_tensors=input_type) + features.pop("audio_chunk_index", None) if input_type == "np": decoder_inputs = {"decoder_input_ids": np.ones((1, 1), dtype=np.int64) * decoder_start_token_id} @@ -455,6 +477,129 @@ def test_pipeline(self, model_arch): del model gc.collect() + @pytest.mark.run_slow + @slow + def test_cohere_asr_generate_non_30s_multiple_audio(self): + # The encoder used to inherit the Whisper dummy generator, which pins input_features to + # 3000 frames, so any audio that was not a multiple of 30s failed at inference time + if "cohere_asr" not in self.SUPPORTED_ARCHITECTURES: + self.skipTest("cohere_asr is not available in this transformers version") + + model_id = MODEL_NAMES["cohere_asr"] + model = self.OVMODEL_CLASS.from_pretrained(model_id, export=True, device=OPENVINO_DEVICE) + processor = AutoProcessor.from_pretrained(model_id) + + encoder_shapes = { + encoder_input.get_any_name(): encoder_input.get_partial_shape() + for encoder_input in model.encoder.model.inputs + } + self.assertIn("attention_mask", encoder_shapes, "encoder must expose an `attention_mask` input") + input_features_shape = encoder_shapes["input_features"] + self.assertTrue( + input_features_shape[1].is_dynamic, + f"encoder `input_features` time dim must be dynamic, got {input_features_shape}", + ) + encoder_output_names = {encoder_output.get_any_name() for encoder_output in model.encoder.model.outputs} + self.assertIn( + "encoder_attention_mask", + encoder_output_names, + "encoder must return the subsampled mask that cross attention runs against", + ) + + np.random.seed(SEED) + for duration_in_seconds in (3, 7, 11): + audio = (np.random.randn(16000 * duration_in_seconds).astype(np.float32)) * 0.01 + inputs = processor(audio, language="en", sampling_rate=16000, return_tensors="pt") + inputs.pop("audio_chunk_index", None) + generated_tokens = model.generate(**inputs, max_new_tokens=8) + self.assertEqual(generated_tokens.shape[0], 1) + self.assertGreater(generated_tokens.shape[1], 1) + + del model + gc.collect() + + @pytest.mark.run_slow + @slow + def test_cohere_asr_padded_batch_matches_single(self): + # Clips of different lengths are padded to a common size, and the frame level mask has to + # survive the eightfold subsampling for cross attention to skip the padded tail + if "cohere_asr" not in self.SUPPORTED_ARCHITECTURES: + self.skipTest("cohere_asr is not available in this transformers version") + + model_id = MODEL_NAMES["cohere_asr"] + model = self.OVMODEL_CLASS.from_pretrained(model_id, export=True, device=OPENVINO_DEVICE) + processor = AutoProcessor.from_pretrained(model_id) + + np.random.seed(SEED) + long_audio = (np.random.randn(16000 * 9).astype(np.float32)) * 0.01 + short_audio = long_audio[: 16000 * 4] + + generate_kwargs = {"max_new_tokens": 8, "do_sample": False, "num_beams": 1} + batched_inputs = processor([long_audio, short_audio], language="en", sampling_rate=16000, return_tensors="pt") + batched_inputs.pop("audio_chunk_index", None) + batched_tokens = model.generate(**batched_inputs, **generate_kwargs) + + for row, audio in enumerate([long_audio, short_audio]): + single_inputs = processor(audio, language="en", sampling_rate=16000, return_tensors="pt") + single_inputs.pop("audio_chunk_index", None) + single_tokens = model.generate(**single_inputs, **generate_kwargs) + self.assertTrue(torch.equal(batched_tokens[row : row + 1], single_tokens)) + + del model + gc.collect() + + @pytest.mark.run_slow + @slow + def test_cohere_asr_with_past_decoder_is_stateful(self): + # The with-past export used to produce a plain decoder without beam_idx or KV cache state, + # which stateful consumers such as openvino_genai.WhisperPipeline require + if "cohere_asr" not in self.SUPPORTED_ARCHITECTURES: + self.skipTest("cohere_asr is not available in this transformers version") + + model_id = MODEL_NAMES["cohere_asr"] + model = self.OVMODEL_CLASS.from_pretrained(model_id, export=True, device=OPENVINO_DEVICE, stateful=True) + self.assertTrue(model_has_state(model.decoder.model)) + decoder_input_names = {decoder_input.get_any_name() for decoder_input in model.decoder.model.inputs} + self.assertIn("beam_idx", decoder_input_names) + self.assertFalse( + any(name.startswith("past_key_values") for name in decoder_input_names), + "the cache has to live in the graph rather than be handed in as inputs", + ) + self.assertGreater(len(model.decoder.model.get_sinks()), 0) + + # More than one decode step also has to run, since the cache length used to be baked into + # the graph as a constant while tracing + processor = AutoProcessor.from_pretrained(model_id) + np.random.seed(SEED) + audio = (np.random.randn(16000 * 5).astype(np.float32)) * 0.01 + inputs = processor(audio, language="en", sampling_rate=16000, return_tensors="pt") + inputs.pop("audio_chunk_index", None) + generated_tokens = model.generate(**inputs, max_new_tokens=12) + self.assertGreater(generated_tokens.shape[1], 1) + + del model + gc.collect() + + @pytest.mark.run_slow + @slow + def test_cohere_asr_exported_processor_is_self_contained(self): + # Reloading the processor straight from an export directory has to work without copying + # tokenizer files over by hand + if "cohere_asr" not in self.SUPPORTED_ARCHITECTURES: + self.skipTest("cohere_asr is not available in this transformers version") + + model_id = MODEL_NAMES["cohere_asr"] + with TemporaryDirectory() as tmp_dir: + model = self.OVMODEL_CLASS.from_pretrained(model_id, export=True, device=OPENVINO_DEVICE) + model.save_pretrained(tmp_dir) + processor = AutoProcessor.from_pretrained(model_id) + processor.save_pretrained(tmp_dir) + + reloaded_processor = AutoProcessor.from_pretrained(tmp_dir) + self.assertIsNotNone(reloaded_processor.tokenizer) + del model + gc.collect() + class OVModelForImageTextToTextIntegrationTest(OVSeq2SeqTestMixin): SUPPORTED_ARCHITECTURES = ["vision-encoder-decoder", "trocr"] diff --git a/tests/openvino/utils_tests.py b/tests/openvino/utils_tests.py index 74418694cc..b6fca87445 100644 --- a/tests/openvino/utils_tests.py +++ b/tests/openvino/utils_tests.py @@ -181,6 +181,7 @@ def _create_tiny_kokoro_model(): "chatglm4": "optimum-intel-internal-testing/tiny-random-chatglm4", "codegen": "optimum-intel-internal-testing/tiny-random-CodeGenForCausalLM", "codegen2": "optimum-intel-internal-testing/tiny-random-codegen2", + "cohere_asr": "optimum-intel-internal-testing/tiny-random-cohere-asr", "data2vec-text": "optimum-intel-internal-testing/tiny-random-Data2VecTextModel", "data2vec-vision": "optimum-intel-internal-testing/tiny-random-Data2VecVisionModel", "data2vec-audio": "optimum-intel-internal-testing/tiny-random-Data2VecAudioModel", @@ -636,6 +637,11 @@ def _resolve_cached_model_paths(model_names: dict) -> dict: "decoder": 30, "decoder_with_past": 30, }, + "cohere_asr": { + "encoder": 60, + "decoder": 48, + "decoder_with_past": 48, + }, } TEST_IMAGE_URL = "http://images.cocodataset.org/val2017/000000039769.jpg"