Fix ONNXRT example with upgraded optimum 1.14.0

examples/onnxrt/nlp/huggingface_model/text_generation/llama/quantization/ptq_static/prepare_model.py

@@ -1,6 +1,9 @@
 import argparse
 import os
 import subprocess
+import optimum.version
+from packaging.version import Version
+OPTIMUM114_VERSION = Version("1.14.0")
 
 
 def parse_arguments():
@@ -12,20 +15,37 @@ def parse_arguments():
 
 def prepare_model(input_model, output_model):
     print("\nexport model...")
-    subprocess.run(
-        [
-            "optimum-cli",
-            "export",
-            "onnx",
-            "--model",
-            f"{input_model}",
-            "--task",
-            "text-generation-with-past",
-            f"{output_model}",
-        ],
-        stdout=subprocess.PIPE,
-        text=True,
-    )
+    if Version(optimum.version.__version__) >= OPTIMUM114_VERSION:
+        subprocess.run(
+            [
+                "optimum-cli",
+                "export",
+                "onnx",
+                "--model",
+                f"{input_model}",
+                "--task",
+                "text-generation-with-past",
+                "--legacy",
+                f"{output_model}",
+            ],
+            stdout=subprocess.PIPE,
+            text=True,
+        )
+    else:
+        subprocess.run(
+            [
+                "optimum-cli",
+                "export",
+                "onnx",
+                "--model",
+                f"{input_model}",
+                "--task",
+                "text-generation-with-past",
+                f"{output_model}",
+            ],
+            stdout=subprocess.PIPE,
+            text=True,
+        )
 
     assert os.path.exists(output_model), f"{output_model} doesn't exist!"
 

examples/onnxrt/nlp/huggingface_model/text_generation/llama/quantization/weight_only/README.md

@@ -17,7 +17,7 @@ pip install -r requirements.txt
 ## 2. Prepare Model
 
 ```bash
-optimum-cli export onnx --model decapoda-research/llama-7b-hf --task text-generation-with-past ./llama_7b
+python prepare_model.py  --input_model="decapoda-research/llama-7b-hf" --output_model="./llama_7b"
 ```
 
 # Run

examples/onnxrt/nlp/huggingface_model/text_generation/llama/quantization/weight_only/prepare_model.py

@@ -0,0 +1,55 @@
+import argparse
+import os
+import subprocess
+import optimum.version
+from packaging.version import Version
+OPTIMUM114_VERSION = Version("1.14.0")
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input_model", type=str, required=False, default="")
+    parser.add_argument("--output_model", type=str, required=True)
+    return parser.parse_args()
+
+
+def prepare_model(input_model, output_model):
+    print("\nexport model...")
+    if Version(optimum.version.__version__) >= OPTIMUM114_VERSION:
+        subprocess.run(
+            [
+                "optimum-cli",
+                "export",
+                "onnx",
+                "--model",
+                f"{input_model}",
+                "--task",
+                "text-generation-with-past",
+                "--legacy",
+                f"{output_model}",
+            ],
+            stdout=subprocess.PIPE,
+            text=True,
+        )
+    else:
+        subprocess.run(
+            [
+                "optimum-cli",
+                "export",
+                "onnx",
+                "--model",
+                f"{input_model}",
+                "--task",
+                "text-generation-with-past",
+                f"{output_model}",
+            ],
+            stdout=subprocess.PIPE,
+            text=True,
+        )
+
+    assert os.path.exists(output_model), f"{output_model} doesn't exist!"
+
+
+if __name__ == "__main__":
+    args = parse_arguments()
+    prepare_model(args.input_model, args.output_model)