[OpenVINO][Models] CohereASR conversion fix

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
+- Qwen3-Coder-30B-A3B-DFlash

optimum/exporters/openvino/input_generators.py

@@ -1396,6 +1396,39 @@ 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 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)."""

optimum/exporters/openvino/model_configs.py

@@ -39,6 +39,8 @@
 from optimum.exporters.openvino.input_generators import (
     AquilaDummyPastKeyValuesGenerator,
     ChatGLM2DummyPastKeyValuesGenerator,
+    CohereAsrDummyAudioInputGenerator,
+    CohereAsrDummySeq2SeqDecoderTextInputGenerator,
     DeciDummyPastKeyValuesGenerator,
     DummyAudioPhi4MMInputGenerator,
     DummyFluxTextInputGenerator,
@@ -104,6 +106,7 @@
     BloomModelPatcher,
     ChatGLMModelPatcher,
     CodeGenModelPatcher,
+    CohereAsrModelPatcher,
     CommonImageEmbeddingsModelPatcher,
     DBRXModelPatcher,
     DeciLMModelPatcher,
@@ -351,6 +354,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 +4370,107 @@ 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):
+    """Export config for CohereAsrForConditionalGeneration (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"],

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,113 @@ 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 KokoroModelPatcher(ModelPatcher):
     """
     Patches the Kokoro TTS model for OpenVINO export by redirecting forward

optimum/intel/openvino/modeling_seq2seq.py

@@ -750,7 +750,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 +906,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 +922,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"]
@@ -1613,3 +1632,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()

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,