Fix sentence transformers CLI export (#674)

* fix sentence transformers export after latest changes merge

* add test case

* Update optimum/commands/export/openvino.py

Co-authored-by: Ella Charlaix <80481427+echarlaix@users.noreply.github.com>

---------

Co-authored-by: Ella Charlaix <80481427+echarlaix@users.noreply.github.com>

Author: eaidova

optimum/commands/export/openvino.py

@@ -214,7 +214,13 @@ def run(self):
             quantization_config["dataset"] = self.args.dataset
             ov_config = OVConfig(quantization_config=quantization_config)
 
-        library_name = TasksManager.infer_library_from_model(self.args.model)
+        library_name = TasksManager.infer_library_from_model(self.args.model, library_name=self.args.library)
+        if library_name == "sentence_transformers" and self.args.library is None:
+            logger.warning(
+                "Library name is not specified. There are multiple possible variants: `sentence_transformers`, `transformers`."
+                "`transformers` will be selected. If you want to load your model with the `sentence-transformers` library instead, please set --library sentence_transformers"
+            )
+            library_name = "transformers"
 
         if (
             library_name == "diffusers"

setup.py

@@ -53,6 +53,7 @@
     "transformers_stream_generator",
     "einops",
     "tiktoken",
+    "sentence_transformers",
 ]
 
 QUALITY_REQUIRE = ["black~=23.1", "ruff>=0.0.241"]

tests/openvino/test_exporters_cli.py

@@ -218,3 +218,31 @@ def test_exporters_cli_help(self):
             shell=True,
             check=True,
         )
+
+    def test_exporters_cli_sentence_transformers(self):
+        model_id = MODEL_NAMES["bge"]
+        with TemporaryDirectory() as tmpdir:
+            # default export creates transformers model
+            subprocess.run(
+                f"optimum-cli export openvino --model {model_id} --task feature-extraction {tmpdir}",
+                shell=True,
+                check=True,
+            )
+            model = eval(_HEAD_TO_AUTOMODELS["feature-extraction"]).from_pretrained(tmpdir, compile=False)
+            self.assertTrue("last_hidden_state" in model.output_names)
+            # export with transformers lib creates transformers model
+            subprocess.run(
+                f"optimum-cli export openvino --model {model_id} --task feature-extraction --library transformers {tmpdir}",
+                shell=True,
+                check=True,
+            )
+            model = eval(_HEAD_TO_AUTOMODELS["feature-extraction"]).from_pretrained(tmpdir, compile=False)
+            self.assertTrue("last_hidden_state" in model.output_names)
+            # export with sentence_transformers lib creates sentence_transformers model
+            subprocess.run(
+                f"optimum-cli export openvino --model {model_id} --task feature-extraction --library sentence_transformers {tmpdir}",
+                shell=True,
+                check=True,
+            )
+            model = eval(_HEAD_TO_AUTOMODELS["feature-extraction"]).from_pretrained(tmpdir, compile=False)
+            self.assertFalse("last_hidden_state" in model.output_names)

tests/openvino/utils_tests.py

@@ -19,6 +19,7 @@
 MODEL_NAMES = {
     "albert": "hf-internal-testing/tiny-random-albert",
     "audio_spectrogram_transformer": "Ericwang/tiny-random-ast",
+    "bge": "BAAI/bge-small-en-v1.5",
     "beit": "hf-internal-testing/tiny-random-BeitForImageClassification",
     "bert": "hf-internal-testing/tiny-random-bert",
     "bart": "hf-internal-testing/tiny-random-bart",