IlyasMoutawwakil HF staff commited on
Commit
f43ba09
1 Parent(s): 04d4e6b

filter cuda

Browse files
Files changed (1) hide show
  1. app.py +6 -3
app.py CHANGED
@@ -1,3 +1,4 @@
 
1
  import random
2
  import gradio as gr
3
  from optimum_benchmark.task_utils import (
@@ -14,9 +15,11 @@ from config_store import (
14
  get_pytorch_config,
15
  )
16
 
 
 
17
  BACKENDS = ["pytorch", "onnxruntime", "openvino", "neural-compressor"]
 
18
  BENCHMARKS = ["inference", "training"]
19
- DEVICES = ["cpu", "cuda"]
20
 
21
 
22
  with gr.Blocks() as demo:
@@ -32,7 +35,7 @@ with gr.Blocks() as demo:
32
  model = gr.Textbox(
33
  label="model",
34
  value="optimum/distilbert-base-uncased-finetuned-sst-2-english",
35
- info="Model to run the benchmark on. In the particular case of this space, only models that are hosted on huggingface.co/models can be benchmarked.",
36
  )
37
  task = gr.Dropdown(
38
  label="task",
@@ -48,7 +51,7 @@ with gr.Blocks() as demo:
48
  )
49
  experiment = gr.Textbox(
50
  label="experiment_name",
51
- value=f"experiment_{random.getrandbits(16)}",
52
  info="Name of the experiment. Will be used to create a folder where results are stored.",
53
  )
54
  model.submit(fn=infer_task_from_model_name_or_path, inputs=model, outputs=task)
 
1
+ import torch
2
  import random
3
  import gradio as gr
4
  from optimum_benchmark.task_utils import (
 
15
  get_pytorch_config,
16
  )
17
 
18
+ cuda_available = torch.cuda.is_available()
19
+
20
  BACKENDS = ["pytorch", "onnxruntime", "openvino", "neural-compressor"]
21
+ DEVICES = ["cpu", "cuda"] if cuda_available else ["cpu"]
22
  BENCHMARKS = ["inference", "training"]
 
23
 
24
 
25
  with gr.Blocks() as demo:
 
35
  model = gr.Textbox(
36
  label="model",
37
  value="optimum/distilbert-base-uncased-finetuned-sst-2-english",
38
+ info="Model to run the benchmark on. Press enter to infer the task automatically.",
39
  )
40
  task = gr.Dropdown(
41
  label="task",
 
51
  )
52
  experiment = gr.Textbox(
53
  label="experiment_name",
54
+ value=f"awesome-experiment-{random.randint(0, 1000)}",
55
  info="Name of the experiment. Will be used to create a folder where results are stored.",
56
  )
57
  model.submit(fn=infer_task_from_model_name_or_path, inputs=model, outputs=task)