changes

app.py

@@ -1,7 +1,6 @@
 import os
 import time
 import traceback
-import logging
 from typing import Optional
 
 from config_store import (
@@ -36,10 +35,8 @@ BACKENDS = ["pytorch", "openvino"]
 TASKS = set(TASKS_TO_OVMODEL.keys()) & set(TASKS_TO_MODEL_LOADERS.keys())
 
 
-def run_benchmark(kwargs, oauth_token: Optional[gr.OAuthToken] = None):
+def run_benchmark(inputs, oauth_token: Optional[gr.OAuthToken] = None):
     if oauth_token.token is None or oauth_token.token == "":
-        yield tuple(None for _ in BACKENDS)
-
         raise gr.Error("Please login to be able to run the benchmark.")
 
     timestamp = time.strftime("%Y-%m-%d-%H-%M-%S")
@@ -47,17 +44,15 @@ def run_benchmark(kwargs, oauth_token: Optional[gr.OAuthToken] = None):
     repo_id = f"{name}/benchmarks"
     token = oauth_token.token
 
-    create_repo(repo_id, token=token, repo_type="dataset", exist_ok=True)
-    gr.Info(f'Created repository "{repo_id}" where results will be pushed.')
+    try:
+        create_repo(repo_id, token=token, repo_type="dataset", exist_ok=True)
+        gr.Info(f"📂 Created dataset repository {repo_id} on the Hub.")
+    except Exception:
+        raise gr.Error(f"Error while creating dataset repository {repo_id} on the Hub.")
 
-    configs = {
-        "process": {},
-        "inference": {},
-        "pytorch": {},
-        "openvino": {},
-    }
+    configs = {"process": {}, "inference": {}, "pytorch": {}, "openvino": {}}
 
-    for key, value in kwargs.items():
+    for key, value in inputs.items():
         if key.label == "model":
             model = value
         elif key.label == "task":
@@ -76,16 +71,10 @@ def run_benchmark(kwargs, oauth_token: Optional[gr.OAuthToken] = None):
     configs["process"] = ProcessConfig(**configs.pop("process"))
     configs["inference"] = InferenceConfig(**configs.pop("inference"))
     configs["pytorch"] = PyTorchConfig(
-        task=task,
-        model=model,
-        device=DEVICE,
-        **configs["pytorch"],
+        task=task, model=model, device=DEVICE, **configs["pytorch"]
     )
     configs["openvino"] = OVConfig(
-        task=task,
-        model=model,
-        device=DEVICE,
-        **configs["openvino"],
+        task=task, model=model, device=DEVICE, **configs["openvino"]
     )
 
     outputs = {
@@ -95,8 +84,6 @@ def run_benchmark(kwargs, oauth_token: Optional[gr.OAuthToken] = None):
 
     yield tuple(outputs[b] for b in BACKENDS)
 
-    setup_logging(level="INFO", prefix="MAIN-PROCESS")
-
     for backend in BACKENDS:
         try:
             benchmark_name = f"{timestamp}/{backend}"
@@ -118,20 +105,19 @@ def run_benchmark(kwargs, oauth_token: Optional[gr.OAuthToken] = None):
                 repo_id=repo_id, subfolder=benchmark_name, token=oauth_token.token
             )
 
-            outputs[backend] = f"\n{benchmark_report.to_markdown_text()}"
+        except Exception:
+            outputs[backend] = f"\n```python-traceback\n{traceback.format_exc()}```\n"
 
             yield tuple(outputs[b] for b in BACKENDS)
 
-            gr.Info(f"Pushed benchmark to {repo_id}/{benchmark_name}")
+            gr.Info(f"❌ Error while running benchmark for {backend} backend.")
 
-        except Exception:
-            outputs[backend] = f"\n```python\n{traceback.format_exc()}```"
+        else:
+            outputs[backend] = f"\n{benchmark_report.to_markdown_text()}\n"
 
             yield tuple(outputs[b] for b in BACKENDS)
 
-            raise gr.Error(f"Error while running benchmark for {backend}")
-
-    logging.getLogger().setLevel(logging.NOTSET)
+            gr.Info(f"✅ Benchmark for {backend} backend ran successfully.")
 
 
 def update_task(model_id):
@@ -195,11 +181,11 @@ with gr.Blocks() as demo:
     with gr.Row():
         button = gr.Button(value="Run Benchmark", variant="primary")
 
-    with gr.Row() as markdown_outputs:
-        with gr.Accordion(label="PyTorch Output", open=True, visible=True):
-            pytorch_output = gr.Markdown()
-        with gr.Accordion(label="OpenVINO Output", open=True, visible=True):
-            openvino_output = gr.Markdown()
+    with gr.Row():
+        with gr.Accordion(label="PyTorch Report", open=True, visible=True):
+            pytorch_report = gr.Markdown()
+        with gr.Accordion(label="OpenVINO Report", open=True, visible=True):
+            openvino_report = gr.Markdown()
 
     model.submit(inputs=model, outputs=task, fn=update_task)
 
@@ -215,8 +201,8 @@ with gr.Blocks() as demo:
             *openvino_config.values(),
         },
         outputs={
-            pytorch_output,
-            openvino_output,
+            pytorch_report,
+            openvino_report,
         },
         concurrency_limit=1,
     )
@@ -225,5 +211,6 @@ with gr.Blocks() as demo:
 if __name__ == "__main__":
     os.environ["LOG_TO_FILE"] = "0"
     os.environ["LOG_LEVEL"] = "INFO"
+    setup_logging(level="INFO", prefix="MAIN-PROCESS")
 
     demo.queue(max_size=10).launch()