Add ability to pass model_dir to .ci/scripts/test_huggingface_optimum_model.py

Author: metascroy

.ci/scripts/test_huggingface_optimum_model.py

@@ -270,6 +270,7 @@ def test_vit(model_id, model_dir, recipe, *, quantize=False, run_only=False):
     parser.add_argument("--model", type=str, required=True)
     parser.add_argument("--recipe", type=str, required=True)
     parser.add_argument("--quantize", action="store_true", help="Enable quantization")
+    parser.add_argument("--model_dir", type=str, required=False)
     args = parser.parse_args()
 
     model_to_model_id_and_test_function = {
@@ -294,11 +295,20 @@ def test_vit(model_id, model_dir, recipe, *, quantize=False, run_only=False):
             f"Unknown model name: {args.model}. Available models: {model_to_model_id_and_test_function.keys()}"
         )
 
-    with tempfile.TemporaryDirectory() as tmp_dir:
-        model_id, test_fn = model_to_model_id_and_test_function[args.model]
+    model_id, test_fn = model_to_model_id_and_test_function[args.model]
+    if args.model_dir is None:
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            test_fn(
+                model_id=model_id,
+                model_dir=tmp_dir,
+                recipe=args.recipe,
+                quantize=args.quantize,
+            )
+    else:
         test_fn(
             model_id=model_id,
-            model_dir=tmp_dir,
+            model_dir=args.model_dir,
             recipe=args.recipe,
             quantize=args.quantize,
+            run_only=False,
         )