Applied comments.

Author: popovaan

optimum/exporters/openvino/model_configs.py

@@ -14,9 +14,11 @@
 
 import enum
 import logging
+import math
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
+import torch
 from transformers import AutoConfig, PretrainedConfig, PreTrainedModel
 
 from optimum.exporters.onnx.config import OnnxConfig, TextDecoderOnnxConfig, TextDecoderWithPositionIdsOnnxConfig
@@ -4370,10 +4372,6 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int
                 dtype=float_dtype,
             )
         if input_name == "image_position_ids":
-            import math
-
-            import torch
-
             # Create position ids as a grid. The patch count = h_patches * w_patches
             # where both are divisible by pooling_kernel_size for correct pooling.
             k = self.pooling_kernel_size

optimum/intel/openvino/modeling_decoder.py

@@ -33,18 +33,19 @@
 from transformers.modeling_outputs import CausalLMOutputWithPast, ModelOutput
 from transformers.utils.hub import PushToHubMixin
 
+from ..utils.import_utils import compare_versions, is_transformers_version
 
-try:
+
+if is_transformers_version("<", "5.5"):
     from transformers.models.mamba.modeling_mamba import MambaCache
-except ImportError:
+else:
     MambaCache = object
 
 from optimum.utils.normalized_config import NormalizedConfigManager
 
 from ...exporters.openvino import ensure_stateful_is_available, main_export, patch_stateful
 from ...exporters.openvino.stateful import model_has_state
 from ...exporters.openvino.utils import SSM_MODELS
-from ..utils.import_utils import compare_versions
 from ..utils.modeling_utils import MULTI_QUERY_ATTN_MODELS
 from .configuration import (
     OVConfig,

tests/openvino/test_seq2seq.py

@@ -787,7 +787,7 @@ def compare_outputs(inputs, ov_model, transformers_model, generation_config):
             with torch.no_grad():
                 transformers_outputs = transformers_model(**transformers_inputs)
             self.assertTrue(
-                torch.allclose(ov_outputs.logits, transformers_outputs.logits.to(torch.float32), atol=4e-3),
+                torch.allclose(ov_outputs.logits, transformers_outputs.logits, atol=4e-3),
                 f"Max abs diff {(torch.abs(ov_outputs.logits - transformers_outputs.logits).max())}",
             )
 

tests/openvino/utils_tests.py

@@ -383,7 +383,7 @@
         "text_embeddings_per_layer_model": 1,
     },
     "gemma4_moe": {
-        "lm_model": 44,
+        "lm_model": 48,
         "text_embeddings_model": 1,
         "vision_embeddings_model": 10,
         "text_embeddings_per_layer_model": 0,