replicate to other generation types

app.py

@@ -153,7 +153,7 @@ def get_plot(model_name, plot_eager, generate_type):
         ci="sd", palette="dark", alpha=.6, height=6
     )
     g.despine(left=True)
-    g.set_axis_labels("GPU", "Generation time (ms)")
+    g.set_axis_labels("GPU", "Generation time (ms) -- LOWER IS BETTER")
     g.legend.set_title("Framework")
     return plt.gcf()
 
@@ -164,7 +164,7 @@ with demo:
         """
         # TensorFlow XLA Text Generation Benchmark
         Pick a tab for the type of generation (or other information), and then select a model from the dropdown menu.
-        You can also ommit results from TensorFlow Eager Execution, if you wish to better compare the performance of
+        You can also omit results from TensorFlow Eager Execution, if you wish to better compare the performance of
         PyTorch to TensorFlow with XLA.
         """
     )
@@ -195,9 +195,58 @@ with demo:
             model_selector.change(fn=plot_fn, inputs=[model_selector, eager_enabler], outputs=plot)
             eager_enabler.change(fn=plot_fn, inputs=[model_selector, eager_enabler], outputs=plot)
         with gr.TabItem("Sample"):
-            gr.Button("New Tiger")
+            gr.Markdown(
+                """
+                ### Sample benchmark parameters
+                - `max_new_tokens = 128`;
+                - `temperature = 2.0`;
+                - `top_k = 50`;
+                - `pad_to_multiple_of = 64` for Tensorflow XLA models. Others do not pad (input prompts between 2 and 33 tokens).
+                """
+            )
+            with gr.Row():
+                model_selector = gr.Dropdown(
+                    choices=["DistilGPT2", "GPT2", "OPT-1.3B", "GPTJ-6B", "T5 Small", "T5 Base", "T5 Large", "T5 3B"],
+                    value="T5 Small",
+                    label="Model",
+                    interactive=True,
+                )
+                eager_enabler = gr.Radio(
+                    ["Yes", "No"],
+                    value="Yes",
+                    label="Plot TF Eager Execution?",
+                    interactive=True
+                )
+            plot_fn = functools.partial(get_plot, generate_type="Sample")
+            plot = gr.Plot(value=plot_fn("T5 Small", "Yes"))  # Show plot when the gradio app is initialized
+            model_selector.change(fn=plot_fn, inputs=[model_selector, eager_enabler], outputs=plot)
+            eager_enabler.change(fn=plot_fn, inputs=[model_selector, eager_enabler], outputs=plot)
         with gr.TabItem("Beam Search"):
-            gr.Button("New Tiger")
+            gr.Markdown(
+                """
+                ### Beam Search benchmark parameters
+                - `max_new_tokens = 256`;
+                - `num_beams = 16`;
+                - `pad_to_multiple_of = 64` for Tensorflow XLA models. Others do not pad (input prompts between 2 and 33 tokens).
+                """
+            )
+            with gr.Row():
+                model_selector = gr.Dropdown(
+                    choices=["DistilGPT2", "GPT2", "OPT-1.3B", "GPTJ-6B", "T5 Small", "T5 Base", "T5 Large", "T5 3B"],
+                    value="T5 Small",
+                    label="Model",
+                    interactive=True,
+                )
+                eager_enabler = gr.Radio(
+                    ["Yes", "No"],
+                    value="Yes",
+                    label="Plot TF Eager Execution?",
+                    interactive=True
+                )
+            plot_fn = functools.partial(get_plot, generate_type="Beam Search")
+            plot = gr.Plot(value=plot_fn("T5 Small", "Yes"))  # Show plot when the gradio app is initialized
+            model_selector.change(fn=plot_fn, inputs=[model_selector, eager_enabler], outputs=plot)
+            eager_enabler.change(fn=plot_fn, inputs=[model_selector, eager_enabler], outputs=plot)
         with gr.TabItem("Benchmark Information"):
             gr.Dataframe(
                 headers=["Parameter", "Value"],