[OpenVINO] support CohereLabs/tiny-aya-base (#1623)

* [OpenVINO] support cohere/tiny-aya-base

* add docs

* add tests

* add tests

* fix tests

* fix tests

* fix tests

* Update Cohere2 minimum transformers version to 4.48.0

Author: Mohamed-Ashraf273

docs/source/openvino/models.mdx

@@ -37,6 +37,7 @@ Here is the list of the supported architectures :
 - CodeGen
 - CodeGen2
 - Cohere
+- Cohere2
 - ConvBert
 - ConvNext
 - DBRX

optimum/exporters/openvino/model_configs.py

@@ -888,6 +888,21 @@ class ArceeOpenVINOConfig(LlamaOpenVINOConfig):
     MIN_TRANSFORMERS_VERSION = "4.53.0"
 
 
+@register_in_tasks_manager(
+    "cohere2",
+    *[
+        "feature-extraction",
+        "feature-extraction-with-past",
+        "text-generation",
+        "text-generation-with-past",
+        "text-classification",
+    ],
+    library_name="transformers",
+)
+class Cohere2OpenVINOConfig(LlamaOpenVINOConfig):
+    MIN_TRANSFORMERS_VERSION = "4.48.0"
+
+
 class QwenDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator):
     def __init__(
         self,

tests/openvino/test_decoder.py

@@ -110,6 +110,9 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase):
 
     SUPPORTED_ARCHITECTURES += SUPPORTED_SSM_ARCHITECTURES
 
+    if is_transformers_version(">=", "4.48.0"):
+        SUPPORTED_ARCHITECTURES += ("cohere2",)
+
     if is_transformers_version(">=", "4.46.0"):
         SUPPORTED_ARCHITECTURES += ("glm", "mistral-nemo", "phimoe")
 
@@ -208,6 +211,7 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase):
         "starcoder2": 2,
         "dbrx": 2,
         "cohere": 2,
+        "cohere2": 2,
         "qwen2": 2,
         "qwen2_moe": 4,
         "arctic": 4,

tests/openvino/test_export.py

@@ -95,6 +95,9 @@ class ExportModelTest(unittest.TestCase):
         "ltx-video": OVLTXPipeline,
     }
 
+    if is_transformers_version(">=", "4.48.0"):
+        SUPPORTED_ARCHITECTURES.update({"cohere2": OVModelForCausalLM})
+
     if is_transformers_version(">=", "4.49"):
         SUPPORTED_ARCHITECTURES.update({"zamba2": OVModelForCausalLM})
 

tests/openvino/test_exporters_cli.py

@@ -116,6 +116,13 @@ class OVCLIExportTestCase(unittest.TestCase):
         ("zero-shot-image-classification", "clip"),
     ]
 
+    if is_transformers_version(">=", "4.48.0"):
+        SUPPORTED_ARCHITECTURES.extend(
+            [
+                ("text-generation-with-past", "cohere2"),
+            ]
+        )
+
     if is_transformers_version(">=", "4.54.0"):
         SUPPORTED_ARCHITECTURES.extend(
             [

tests/openvino/test_genai.py

@@ -78,6 +78,8 @@ class LLMPipelineTestCase(unittest.TestCase):
         "granitemoe",
     )
 
+    if is_transformers_version(">=", "4.48.0"):
+        SUPPORTED_ARCHITECTURES += ("cohere2",)
     if is_transformers_version(">=", "4.46.0"):
         SUPPORTED_ARCHITECTURES += ("glm", "mistral-nemo", "phimoe", "opt")
         if is_transformers_version("<", "4.54.0"):

tests/openvino/test_quantization.py

@@ -1074,6 +1074,9 @@ class OVWeightCompressionTest(unittest.TestCase):
     if is_transformers_version("<", "4.52.0"):
         SUPPORTED_ARCHITECTURES_WITH_AUTO_COMPRESSION.append((OVModelForVisualCausalLM, "minicpmo", True))
 
+    if is_transformers_version(">=", "4.48.0"):
+        SUPPORTED_ARCHITECTURES_WITH_AUTO_COMPRESSION.append((OVModelForCausalLM, "cohere2", False))
+
     if is_transformers_version(">=", "4.54.0"):
         SUPPORTED_ARCHITECTURES_WITH_AUTO_COMPRESSION.append((OVModelForCausalLM, "exaone4", True))
 

tests/openvino/utils_tests.py

@@ -57,6 +57,7 @@
     "clip": "optimum-intel-internal-testing/tiny-random-CLIPModel",
     "convbert": "optimum-intel-internal-testing/tiny-random-ConvBertForSequenceClassification",
     "cohere": "optimum-intel-internal-testing/tiny-random-CohereForCausalLM",
+    "cohere2": "optimum-intel-internal-testing/tiny-random-aya-base",
     "chatglm": "optimum-intel-internal-testing/tiny-random-chatglm",
     "chatglm4": "optimum-intel-internal-testing/tiny-random-chatglm4",
     "codegen": "optimum-intel-internal-testing/tiny-random-CodeGenForCausalLM",
@@ -237,6 +238,7 @@
     "albert": {"model": 84},
     "vit": {"model": 64},
     "blenderbot": {"model": 70},
+    "cohere2": {"model": 30},
     "gpt2": {"model": 44},
     "granitemoehybrid": {"model": 118},
     "wav2vec2": {"model": 34},