app.py

import functools
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

import gradio as gr


BENCHMARK_DATA = {
    "Greedy Search": {
        "DistilGPT2": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "GPT2": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "OPT-1.3B": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "GPTJ-6B": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "T5 Small": {
            "T4": [1, 2, 3, 4],
            "3090": [],
            "A100": [],
        },
        "T5 Base": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "T5 Large": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "T5 3B": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
    },
    "Sample": {
        "DistilGPT2": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "GPT2": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "OPT-1.3B": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "GPTJ-6B": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "T5 Small": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "T5 Base": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "T5 Large": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "T5 3B": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
    },
    "Beam Search": {
        "DistilGPT2": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "GPT2": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "OPT-1.3B": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "GPTJ-6B": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "T5 Small": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "T5 Base": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "T5 Large": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
        "T5 3B": {
            "T4": [],
            "3090": [],
            "A100": [],
        },
    },
}


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()

demo = gr.Blocks()

with demo:
    with gr.Tabs():
        with gr.TabItem("Greedy Search"):
            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,
            )
            plot_fn = functools.partial(get_plot, generate_type="Greedy Search")
            plot = gr.Plot(value=plot_fn("T5 Small"))  # Show plot when the gradio app is initialized
            model_selector.change(fn=get_plot, inputs=[model_selector], outputs=plot)
        with gr.TabItem("Sample"):
            gr.Button("New Tiger")
        with gr.TabItem("Beam Search"):
            gr.Button("New Tiger")
        with gr.TabItem("Benchmark Information"):
            gr.Dataframe(
                headers=["Parameter", "Value"],
                value=[
                    ["Transformers Version", "4.22.dev0"],
                    ["TensorFlow Version", "2.9.1"],
                    ["Pytorch Version", "1.11.0"],
                    ["OS", "22.04 LTS (3090) / Debian 10 (other GPUs)"],
                    ["CUDA", "11.6 (3090) / 11.3 (others GPUs)"],
                    ["Number of runs", "100 (the first run was discarded to ignore compilation time)"],
                    ["Is there code to reproduce?", "Yes -- https://gist.github.com/gante/f0017e3f13ac11b0c02e4e4db351f52f"],
                ],
            )

demo.launch()