use search

app.py

@@ -13,6 +13,7 @@ from config_store import (
 
 import gradio as gr
 from huggingface_hub import create_repo, whoami
+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
 from optimum_benchmark.backends.transformers_utils import TASKS_TO_MODEL_LOADERS
@@ -25,24 +26,13 @@ from optimum_benchmark import (
     OVConfig,
 )
 from optimum_benchmark.logging_utils import setup_logging
+from optimum_benchmark.task_utils import infer_task_from_model_name_or_path
 
 
 DEVICE = "cpu"
 LAUNCHER = "process"
 SCENARIO = "inference"
 BACKENDS = ["pytorch", "openvino"]
-MODELS = [
-    "openai-community/gpt2",
-    "google-bert/bert-base-uncased",
-    "hf-internal-testing/tiny-random-LlamaForCausalLM",
-    "hf-internal-testing/tiny-random-BertForSequenceClassification",
-]
-MODELS_TO_TASKS = {
-    "openai-community/gpt2": "text-generation",
-    "google-bert/bert-base-uncased": "text-classification",
-    "hf-internal-testing/tiny-random-LlamaForCausalLM": "text-generation",
-    "hf-internal-testing/tiny-random-BertForSequenceClassification": "text-classification",
-}
 TASKS = set(TASKS_TO_OVMODEL.keys()) & set(TASKS_TO_MODEL_LOADERS.keys())
 
 
@@ -159,19 +149,19 @@ def build_demo():
             "</p>"
         )
 
-        model = gr.Dropdown(
+        model = HuggingfaceHubSearch(
             label="model",
-            choices=MODELS,
-            value=MODELS[0],
-            info="Model to run the benchmark on.",
+            search_type="model",
+            value="openai-community/gpt2",
+            placeholder="Search for a model",
+            sumbit_on_select=True,
         )
         task = gr.Dropdown(
             label="task",
             choices=TASKS,
-            value=MODELS_TO_TASKS[MODELS[0]],
+            value="text-generation",
             info="Task to run the benchmark on.",
         )
-
         backends = gr.CheckboxGroup(
             interactive=True,
             label="backends",
@@ -202,9 +192,7 @@ def build_demo():
             with gr.Accordion(label="OpenVINO Output", open=True, visible=True):
                 openvino_output = gr.Markdown()
 
-        model.change(
-            inputs=model, outputs=task, fn=lambda value: MODELS_TO_TASKS[value]
-        )
+        model.submit(inputs=model, outputs=task, fn=infer_task_from_model_name_or_path)
 
         backends.change(
             inputs=backends,

requirements.txt

@@ -1 +1,2 @@
+gradio_huggingfacehub_search
 optimum-benchmark[openvino]@git+https://github.com/huggingface/optimum-benchmark.git@markdown-report