fix style

huggingface · Feb 29, 2024 · 41805d8 · 41805d8
1 parent 0aa3f29
commit 41805d8
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 15 deletions.
diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py
@@ -32,8 +32,12 @@
 
 
 def init_model_configs():
-
-    supported_model_types = ["_SUPPORTED_MODEL_TYPE", "_DIFFUSERS_SUPPORTED_MODEL_TYPE", "_TIMM_SUPPORTED_MODEL_TYPE", "_SENTENCE_TRANSFORMERS_SUPPORTED_MODEL_TYPE"]
+    supported_model_types = [
+        "_SUPPORTED_MODEL_TYPE",
+        "_DIFFUSERS_SUPPORTED_MODEL_TYPE",
+        "_TIMM_SUPPORTED_MODEL_TYPE",
+        "_SENTENCE_TRANSFORMERS_SUPPORTED_MODEL_TYPE",
+    ]
 
     for supported_models_config in supported_model_types:
         supported_models = getattr(TasksManager, supported_models_config)
@@ -42,9 +46,8 @@ def init_model_configs():
                 continue
             onnx_config = export_configs["onnx"]
             supported_models[model]["openvino"] = deepcopy(onnx_config)
-
-        setattr(TasksManager, supported_models_config, supported_models)
 
+        setattr(TasksManager, supported_models_config, supported_models)
 
 
 init_model_configs()

diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py
@@ -71,7 +71,7 @@ def patch_model_with_bettertransformer(model):
     return model
 
 
-def mixtral_sparse_moe_block_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+def _mixtral_sparse_moe_block_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
     """ """
     batch_size, sequence_length, hidden_dim = hidden_states.shape
     hidden_states = hidden_states.view(-1, hidden_dim)
@@ -113,7 +113,9 @@ def __enter__(self):
         super().__enter__()
         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)
+            layer.block_sparse_moe.forward = types.MethodType(
+                _mixtral_sparse_moe_block_forward, layer.block_sparse_moe
+            )
 
     def __exit__(self, exc_type, exc_value, traceback):
         super().__exit__(exc_type, exc_value, traceback)

diff --git a/optimum/intel/openvino/modeling_decoder.py b/optimum/intel/openvino/modeling_decoder.py
@@ -311,7 +311,9 @@ def _reshape(
             shapes[inputs][0] = -1
             input_name = inputs.get_any_name()
             if input_name.startswith("past_key_values"):
-                if (len(inputs.partial_shape) == 3 and input_name.endswith("value")) or self.config.model_type == "chatglm":
+                if (
+                    len(inputs.partial_shape) == 3 and input_name.endswith("value")
+                ) or self.config.model_type == "chatglm":
                     shapes[inputs][1] = -1
                 else:
                     shapes[inputs][2] = -1

diff --git a/setup.py b/setup.py
@@ -48,7 +48,6 @@
         "openvino>=2023.3",
         "onnx",
         "onnxruntime",
-        "transformers>=4.36.0",
         "optimum @ git+https://github.com/huggingface/optimum.git#egg=optimum",
     ],
     "openvino-tokenizers": ["openvino-tokenizers[transformers]"],

diff --git a/tests/openvino/test_modeling.py b/tests/openvino/test_modeling.py
@@ -497,7 +497,7 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase):
     )
     GENERATION_LENGTH = 100
     IS_SUPPORT_STATEFUL = is_openvino_version(">=", "2023.3")
-    REMOTE_CODE_MODELS = ("chatglm", )
+    REMOTE_CODE_MODELS = ("chatglm",)
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES)
     def test_compare_to_transformers(self, model_arch):
@@ -513,13 +513,16 @@ def test_compare_to_transformers(self, model_arch):
 
         model_kwargs = {}
         if model_arch in self.REMOTE_CODE_MODELS:
-            model_kwargs = {"config": AutoConfig.from_pretrained(model_id, trust_remote_code=True), "trust_remote_code": True}
+            model_kwargs = {
+                "config": AutoConfig.from_pretrained(model_id, trust_remote_code=True),
+                "trust_remote_code": True,
+            }
         ov_model = OVModelForCausalLM.from_pretrained(model_id, export=True, ov_config=F32_CONFIG, **model_kwargs)
         self.assertIsInstance(ov_model.config, PretrainedConfig)
         self.assertTrue(ov_model.use_cache)
-        self.assertEqual(ov_model.stateful, self.IS_SUPPORT_STATEFUL and model_arch not in not_stateful)
+        self.assertEqual(ov_model.stateful, self.IS_SUPPORT_STATEFUL and ov_model.config.model_type not in not_stateful)
         transformers_model = AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs)
-        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=model_arch in remote_code)
+        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=model_arch in self.REMOTE_CODE_MODELS)
         tokens = tokenizer(
             "This is a sample", return_tensors="pt", return_token_type_ids=False if model_arch == "llama" else None
         )
@@ -552,9 +555,14 @@ def test_pipeline(self, model_arch):
         model_kwargs = {}
         model_id = MODEL_NAMES[model_arch]
         if model_arch in self.REMOTE_CODE_MODELS:
-            model_kwargs = {"config": AutoConfig.from_pretrained(model_id, trust_remote_code=True), "trust_remote_code": True}
+            model_kwargs = {
+                "config": AutoConfig.from_pretrained(model_id, trust_remote_code=True),
+                "trust_remote_code": True,
+            }
         tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=model_arch in self.REMOTE_CODE_MODELS)
-        model = OVModelForCausalLM.from_pretrained(model_id, export=True, use_cache=False, compile=False, **model_kwargs)
+        model = OVModelForCausalLM.from_pretrained(
+            model_id, export=True, use_cache=False, compile=False, **model_kwargs
+        )
         model.config.encoder_no_repeat_ngram_size = 0
         model.to("cpu")
         model.half()
@@ -573,7 +581,10 @@ def test_multiple_inputs(self, model_arch):
         set_seed(SEED)
         model_kwargs = {}
         if model_arch in self.REMOTE_CODE_MODELS:
-            model_kwargs = {"config": AutoConfig.from_pretrained(model_id, trust_remote_code=True), "trust_remote_code": True}
+            model_kwargs = {
+                "config": AutoConfig.from_pretrained(model_id, trust_remote_code=True),
+                "trust_remote_code": True,
+            }
         model = OVModelForCausalLM.from_pretrained(model_id, export=True, compile=False, **model_kwargs)
         tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=model_arch in self.REMOTE_CODE_MODELS)
         tokenizer.pad_token = tokenizer.eos_token