test

app.py

@@ -11,7 +11,8 @@ from config_store import (
 )
 
 import gradio as gr
-from huggingface_hub import whoami, login, logout
+from huggingface_hub import whoami
+from huggingface_hub.errors import GatedRepoError
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
 from optimum_benchmark.launchers.device_isolation_utils import *  # noqa
 from optimum_benchmark.backends.openvino.utils import TASKS_TO_OVMODEL
@@ -37,7 +38,7 @@ BENCHMARKS_REPO_ID = "optimum-benchmark/OpenVINO-Benchmarks"
 TASKS = set(TASKS_TO_OVMODEL.keys()) & set(TASKS_TO_MODEL_LOADERS.keys())
 
 
-def parse_configs(inputs, token: str):
+def parse_configs(inputs):
     configs = {"process": {}, "inference": {}, "pytorch": {}, "openvino": {}}
 
     for key, value in inputs.items():
@@ -63,14 +64,12 @@ def parse_configs(inputs, token: str):
         model=model,
         device=DEVICE,
         **configs["pytorch"],
-        model_kwargs={"token": token},
     )
     configs["openvino"] = OVConfig(
         task=task,
         model=model,
         device=DEVICE,
         **configs["openvino"],
-        model_kwargs={"token": token},
     )
 
     return configs
@@ -87,7 +86,7 @@ def run_benchmark(inputs, oauth_token: Optional[gr.OAuthToken]):
     gr.Info(f"📩 Benchmark will be saved under {BENCHMARKS_REPO_ID}/{folder}")
 
     outputs = {backend: "Running..." for backend in BACKENDS}
-    configs = parse_configs(inputs, token=oauth_token.token)
+    configs = parse_configs(inputs)
     yield tuple(outputs[b] for b in BACKENDS)
 
     for backend in BACKENDS:
@@ -101,9 +100,6 @@ def run_benchmark(inputs, oauth_token: Optional[gr.OAuthToken]):
             )
             benchmark_report = Benchmark.launch(benchmark_config)
 
-            # make sure the token is never pushed to the hub
-            del benchmark_config.backend.model_kwargs["token"]
-
             benchmark_config.push_to_hub(
                 repo_id=BENCHMARKS_REPO_ID,
                 subfolder=benchmark_name,
@@ -115,6 +111,11 @@ def run_benchmark(inputs, oauth_token: Optional[gr.OAuthToken]):
                 token=BENCHMARKS_HF_TOKEN,
             )
 
+        except GatedRepoError:
+            outputs[backend] = f"🔒 Model {configs[backend].model} is gated."
+            yield tuple(outputs[b] for b in BACKENDS)
+            gr.Info("🔒 Gated Repo Error while trying to access the model.")
+
         except Exception:
             outputs[backend] = f"\n```python-traceback\n{traceback.format_exc()}```\n"
             yield tuple(outputs[b] for b in BACKENDS)
@@ -129,6 +130,12 @@ def run_benchmark(inputs, oauth_token: Optional[gr.OAuthToken]):
 def update_task(model_id):
     try:
         inferred_task = infer_task_from_model_name_or_path(model_id)
+
+    except GatedRepoError:
+        raise gr.Error(
+            f"Model {model_id} is gated, please use optimum-benchmark locally to benchmark it."
+        )
+
     except Exception:
         raise gr.Error(
             f"Error while inferring task for {model_id}, please select a task manually."