playing around with gradio

app.py

@@ -1,11 +1,11 @@
 import functools
-import matplotlib
-matplotlib.use('Agg')
-import matplotlib.pyplot as plt
 
 import gradio as gr
+import seaborn as sns
+import pandas as pd
 
 
+# benchmark order: pytorch, tf eager, tf xla; units = ms
 BENCHMARK_DATA = {
     "Greedy Search": {
         "DistilGPT2": {
@@ -29,9 +29,9 @@ BENCHMARK_DATA = {
             "A100": [],
         },
         "T5 Small": {
-            "T4": [1, 2, 3, 4],
-            "3090": [],
-            "A100": [],
+            "T4": [99.88, 1527.73, 18.78],
+            "3090": [55.09, 665.70, 9.25],
+            "A100": [124.91, 1642.07, 13.72],
         },
         "T5 Base": {
             "T4": [],
@@ -137,10 +137,19 @@ BENCHMARK_DATA = {
 
 
 def get_plot(model_name, generate_type):
-    data = BENCHMARK_DATA[generate_type][model_name]["T4"]
-    plt.plot(data)
-    plt.title(model_name)
-    return plt.gcf()
+    df = pd.DataFrame(BENCHMARK_DATA[generate_type][model_name])
+    df["framework"] = ["PyTorch", "TF (Eager Execition)", "TF (XLA)"]
+    df = pd.melt(df, id_vars=["framework"], value_vars=["T4", "3090", "A100"])
+
+    g = sns.catplot(
+        data=df, kind="bar",
+        x="variable", y="value", hue="framework",
+        ci="sd", palette="dark", alpha=.6, height=6
+    )
+    g.despine(left=True)
+    # g.set_axis_labels("", "Body mass (g)")
+    # g.legend.set_title("")
+    return g.gcf()
 
 demo = gr.Blocks()