diff --git a/optimum/commands/export/openvino.py b/optimum/commands/export/openvino.py index 578fd65216..e38108372e 100644 --- a/optimum/commands/export/openvino.py +++ b/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, ) diff --git a/optimum/exporters/openvino/__main__.py b/optimum/exporters/openvino/__main__.py index ca1c6dd2e6..241f7bdcb0 100644 --- a/optimum/exporters/openvino/__main__.py +++ b/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( diff --git a/tests/openvino/test_exporters_cli.py b/tests/openvino/test_exporters_cli.py index 7817fedda9..b74dba04c8 100644 --- a/tests/openvino/test_exporters_cli.py +++ b/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, + ) + 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)