apply review comments

huggingface · Jun 6, 2024 · 0c0cedb · 0c0cedb
1 parent d2a25ff
commit 0c0cedb
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 11 deletions.
diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py
@@ -155,18 +155,13 @@ def __exit__(self, exc_type, exc_value, traceback):
 
 class ArcticModelPatcher(MixtralModelPatcher):
     def __enter__(self):
-        super().__enter__()
         # model initialize some weights for matrix multiplication in bfloat16, that lead to inconsistency of dtype
         try:
             self._model.to(torch.float32)
         except Exception:
             pass
 
-        for layer in self._model.model.layers:
-            layer.block_sparse_moe._unpatched_forward = layer.block_sparse_moe.forward
-            layer.block_sparse_moe.forward = types.MethodType(
-                _mixtral_sparse_moe_block_forward, layer.block_sparse_moe
-            )
+        super().__enter__()
 
 
 def _chatglm_transformer_forward(

diff --git a/tests/openvino/test_modeling.py b/tests/openvino/test_modeling.py
@@ -565,7 +565,7 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase):
         "dbrx",
         "qwen2-moe",
         "jais",
-        "snowflake",
+        "arctic",
     )
     GENERATION_LENGTH = 100
     REMOTE_CODE_MODELS = (
@@ -583,7 +583,7 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase):
         "xverse",
         "internlm",
         "codegen2",
-        "snowflake",
+        "arctic",
     )
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES)
@@ -625,7 +625,7 @@ def test_compare_to_transformers(self, model_arch):
 
         set_seed(SEED)
         transformers_model = AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs)
-        if model_arch in ["qwen", "snowflake"]:
+        if model_arch in ["qwen", "arctic"]:
             transformers_model.to(torch.float32)
 
         with torch.no_grad():
@@ -873,7 +873,7 @@ def test_beam_search(self, model_arch):
         )
         transformers_model = AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs)
 
-        if model_arch == "snowflake":
+        if model_arch == "arctic":
             transformers_model.to(torch.float32)
         tokenizer.pad_token_id = tokenizer.eos_token_id
         tokens = tokenizer(["Today is a nice day and I am longer", "This is me"], return_tensors="pt", padding=True)

diff --git a/tests/openvino/utils_tests.py b/tests/openvino/utils_tests.py
@@ -110,7 +110,7 @@
     "latent-consistency": "echarlaix/tiny-random-latent-consistency",
     "sew": "hf-internal-testing/tiny-random-SEWModel",
     "sew_d": "asapp/sew-d-tiny-100k-ft-ls100h",
-    "snowflake": "katuni4ka/tiny-random-snowflake",
+    "arctic": "katuni4ka/tiny-random-snowflake",
     "swin": "hf-internal-testing/tiny-random-SwinModel",
     "t5": "hf-internal-testing/tiny-random-t5",
     "trocr": "microsoft/trocr-small-handwritten",