[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,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)."""

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"],