import gradio as gr
import altair as alt
from datasets import load_dataset
import pandas as pd

model_id = "ybelkada/model_cards_correct_tag"
dataset = load_dataset(model_id, split="train").to_pandas()

# Convert dataset to a pandas DataFrame and sort by commit_dates
df = pd.DataFrame(dataset)
df["commit_dates"] = pd.to_datetime(df["commit_dates"])  # Convert commit_dates to datetime format
df = df.sort_values(by="commit_dates")

def plot_fn():
    line_chart = alt.Chart(df).mark_line().encode(
        x=alt.X('commit_dates:T', axis=alt.Axis(title='Date')),
        y=alt.Y('total_transformers_model:Q', axis=alt.Axis(title='Count'), scale=alt.Scale(zero=False)),
        color=alt.value('blue'),
        tooltip=['commit_dates:T', 'total_transformers_model:Q'],
    ).properties(width=600, height=400)

    line_chart_missing_library = alt.Chart(df).mark_line().encode(
        x=alt.X('commit_dates:T', axis=alt.Axis(title='Date')),
        y=alt.Y('missing_library_name:Q', axis=alt.Axis(title='Count'), scale=alt.Scale(zero=False)),
        color=alt.value('orange'),
        tooltip=['commit_dates:T', 'missing_library_name:Q'],
    ).properties(width=600, height=400)
    
    chart = (line_chart + line_chart_missing_library).properties(width=600, height=400)

    return chart

with gr.Blocks() as demo:
    plot = gr.Plot(label="Plot")
    demo.load(plot_fn, inputs=[], outputs=[plot])

if __name__ == "__main__":
    demo.launch()