Add --revision to optimum-cli export openvino

Author: helena-intel

optimum/commands/export/openvino.py

@@ -263,6 +263,11 @@ def parse_args_openvino(parser: "ArgumentParser"):
             "reduces quantization error. Valid only when activations quantization is enabled."
         ),
     )
+    optional_group.add_argument(
+        "--revision",
+        type=str,
+        help=("Revision of the model to load."),
+    )
     optional_group.add_argument(
         "--model-kwargs",
         type=json.loads,
@@ -332,7 +337,7 @@ def run(self):
         from ...intel.openvino.configuration import _DEFAULT_4BIT_WQ_CONFIG, OVConfig, get_default_quantization_config
 
         if self.args.library is None:
-            # TODO: add revision, subfolder and token to args
+            # TODO: add subfolder and token to args
             library_name = _infer_library_from_model_name_or_path(
                 model_name_or_path=self.args.model, cache_dir=self.args.cache_dir
             )
@@ -427,6 +432,7 @@ def run(self):
                 self.args.model,
                 cache_dir=self.args.cache_dir,
                 trust_remote_code=self.args.trust_remote_code,
+                revision=self.args.revision,
             )
             if getattr(config, "model_type", "").replace("_", "-") in MULTI_MODAL_TEXT_GENERATION_MODELS:
                 task = "image-text-to-text"
@@ -473,7 +479,9 @@ def run(self):
             else:
                 raise NotImplementedError(f"Quantization isn't supported for class {class_name}.")
 
-            model = model_cls.from_pretrained(self.args.model, export=True, quantization_config=quantization_config)
+            model = model_cls.from_pretrained(
+                self.args.model, export=True, quantization_config=quantization_config, revision=self.args.revision
+            )
             model.save_pretrained(self.args.output)
             if not self.args.disable_convert_tokenizer:
                 maybe_convert_tokenizers(library_name, self.args.output, model, task=task)
@@ -529,6 +537,7 @@ def run(self):
                 trust_remote_code=self.args.trust_remote_code,
                 variant=self.args.variant,
                 cache_dir=self.args.cache_dir,
+                revision=self.args.revision,
             )
             model.save_pretrained(self.args.output)
 
@@ -551,6 +560,7 @@ def run(self):
                 convert_tokenizer=not self.args.disable_convert_tokenizer,
                 library_name=library_name,
                 variant=self.args.variant,
+                revision=self.args.revision,
                 model_kwargs=self.args.model_kwargs,
                 # **input_shapes,
             )

optimum/exporters/openvino/__main__.py

@@ -383,7 +383,9 @@ class StoreAttr(object):
 
     try:
         if library_name == "open_clip":
-            model = _OpenClipForZeroShotImageClassification.from_pretrained(model_name_or_path, cache_dir=cache_dir)
+            model = _OpenClipForZeroShotImageClassification.from_pretrained(
+                model_name_or_path, cache_dir=cache_dir, revision=revision
+            )
         else:
             model = TasksManager.get_model_from_task(
                 task,
@@ -407,7 +409,7 @@ class StoreAttr(object):
             if pad_token_id is not None:
                 model.config.pad_token_id = pad_token_id
             else:
-                tok = AutoTokenizer.from_pretrained(model_name_or_path)
+                tok = AutoTokenizer.from_pretrained(model_name_or_path, revision=revision)
                 pad_token_id = getattr(tok, "pad_token_id", None)
                 if pad_token_id is None:
                     raise ValueError(

tests/openvino/test_exporters_cli.py

@@ -201,9 +201,11 @@ class OVCLIExportTestCase(unittest.TestCase):
             "whisper",
             "int8",
             "--dataset librispeech --num-samples 1 --smooth-quant-alpha 0.9 --trust-remote-code",
-            {"encoder": 10, "decoder": 12, "decoder_with_past": 11}
-            if is_transformers_version("<=", "4.36.0")
-            else {"encoder": 8, "decoder": 12, "decoder_with_past": 25},
+            (
+                {"encoder": 10, "decoder": 12, "decoder_with_past": 11}
+                if is_transformers_version("<=", "4.36.0")
+                else {"encoder": 8, "decoder": 12, "decoder_with_past": 25}
+            ),
             (
                 {"encoder": {"int8": 8}, "decoder": {"int8": 11}, "decoder_with_past": {"int8": 9}}
                 if is_transformers_version("<=", "4.36.0")
@@ -215,9 +217,11 @@ class OVCLIExportTestCase(unittest.TestCase):
             "whisper",
             "f8e4m3",
             "--dataset librispeech --num-samples 1 --smooth-quant-alpha 0.9 --trust-remote-code",
-            {"encoder": 10, "decoder": 12, "decoder_with_past": 11}
-            if is_transformers_version("<=", "4.36.0")
-            else {"encoder": 8, "decoder": 12, "decoder_with_past": 25},
+            (
+                {"encoder": 10, "decoder": 12, "decoder_with_past": 11}
+                if is_transformers_version("<=", "4.36.0")
+                else {"encoder": 8, "decoder": 12, "decoder_with_past": 25}
+            ),
             (
                 {"encoder": {"f8e4m3": 8}, "decoder": {"f8e4m3": 11}, "decoder_with_past": {"f8e4m3": 9}}
                 if is_transformers_version("<=", "4.36.0")
@@ -1142,3 +1146,24 @@ def test_export_openvino_with_custom_variant(self):
             model = eval(_HEAD_TO_AUTOMODELS["stable-diffusion"]).from_pretrained(tmpdir, compile=False)
             for component in ["text_encoder", "tokenizer", "unet", "vae_encoder", "vae_decoder"]:
                 self.assertIsNotNone(getattr(model, component))
+
+    def test_export_openvino_with_revision(self):
+        with TemporaryDirectory() as tmpdir:
+            subprocess.run(
+                f"optimum-cli export openvino --model hf-internal-testing/tiny-random-MistralForCausalLM --revision 7158fab {tmpdir}",
+                shell=True,
+                check=True,
+            )
+            model = eval(_HEAD_TO_AUTOMODELS["text-generation"]).from_pretrained(tmpdir, compile=False)
+
+        with TemporaryDirectory() as tmpdir:
+            result = subprocess.run(
+                f"optimum-cli export openvino --model hf-internal-testing/tiny-random-MistralForCausalLM --revision 7158fac {tmpdir}",
+                shell=True,
+                check=False,
+                text=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            self.assertNotEqual(result.returncode, 0)
+            self.assertIn("not a valid git identifier", result.stderr)