models.py

import streamlit as st
import pandas as pd
from datasets import load_dataset
from ast import literal_eval
import altair as alt
import plotly.graph_objs as go
import matplotlib.pyplot as plt

nlp_tasks = ["text-classification", "text-generation", "text2text-generation", "token-classification", "fill-mask", "question-answering",
             "translation", "conversational", "sentence-similarity", "summarization", "multiple-choice", "zero-shot-classification", "table-question-answering"
]
audio_tasks = ["automatic-speech-recognition", "audio-classification", "text-to-speech", "audio-to-audio", "voice-activity-detection"]
cv_tasks = ["image-classification", "image-segmentation", "zero-shot-image-classification", "image-to-image", "unconditional-image-generation", "object-detection"]
multimodal = ["feature-extraction", "text-to-image", "visual-question-answering", "image-to-text", "document-question-answering"]
tabular = ["tabular-classification", "tabular-regression"]

modalities = {
    "nlp": nlp_tasks,
    "audio": audio_tasks,
    "cv": cv_tasks,
    "multimodal": multimodal,
    "tabular": tabular,
    "rl": ["reinforcement-learning"]
}

def modality(row):
  pipeline = row["pipeline"]
  for modality, tasks in modalities.items():
    if pipeline in tasks:
      return modality
  if type(pipeline) == "str":
    return "unk_modality"
  return None

supported_revisions = ["27_09_22"]

def process_dataset(version):
    # Load dataset at specified revision
    dataset = load_dataset("open-source-metrics/model-repos-stats", revision=version)

    # Convert to pandas dataframe
    data = dataset["train"].to_pandas()

    # Add modality column
    data["modality"] = data.apply(modality, axis=1)

    # Bin the model card length into some bins
    data["length_bins"] = pd.cut(data["text_length"], [0, 200, 1000, 2000, 3000, 4000, 5000, 7500, 10000, 20000, 50000])

    return data

base = st.selectbox(
    'What revision do you want to use',
    supported_revisions)
data = process_dataset(base)

def eval_tags(row):
    tags = row["tags"]
    if tags == "none" or tags == [] or tags == "{}":
        return []
    if tags[0] != "[":
        tags = str([tags])
    val = literal_eval(tags)
    if isinstance(val, dict):
        return []
    return val

data["tags"] = data.apply(eval_tags, axis=1)

total_samples = data.shape[0]
st.metric(label="Total models", value=total_samples)

tab1, tab2, tab3, tab4, tab5, tab6, tab7, tab8 = st.tabs(["Language", "License", "Pipeline", "Discussion Features", "Libraries", "Model Cards", "Super users", "Raw Data"])

with tab1:
    st.header("Languages info")

    data.loc[data.languages == "False", 'languages'] = None
    data.loc[data.languages == {}, 'languages'] = None

    no_lang_count = data["languages"].isna().sum()
    data["languages"] = data["languages"].fillna('none')

    def make_list(row):
        languages = row["languages"]
        if languages == "none":
            return []
        return literal_eval(languages)

    def language_count(row):
        languages = row["languages"]
        leng = len(languages)
        return leng

    data["languages"] = data.apply(make_list, axis=1)
    data["language_count"] = data.apply(language_count, axis=1)

    models_with_langs = data[data["language_count"] > 0]
    langs = models_with_langs["languages"].explode()
    langs = langs[langs != {}]
    total_langs = len(langs.unique())

    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric(label="Language Specified", value=total_samples-no_lang_count)
    with col2:
        st.metric(label="No Language Specified", value=no_lang_count)
    with col3:
        st.metric(label="Total Unique Languages", value=total_langs)

    st.subheader("Count of languages per model repo")
    st.text("Some repos are for multiple languages, so the count is greater than 1")
    linguality = st.selectbox(
        'All or just Multilingual',
        ["All", "Just Multilingual", "Three or more languages"])

    filter = 0
    if linguality == "Just Multilingual":
        filter = 1
    elif linguality == "Three or more languages":
        filter = 2

    models_with_langs = data[data["language_count"] > filter]
    df1 = models_with_langs['language_count'].value_counts()
    st.bar_chart(df1)

    st.subheader("Most frequent languages")
    linguality_2 = st.selectbox(
        'All or filtered',
        ["All", "No English", "Remove top 10"])

    filter = 0
    if linguality_2 == "All":
        filter = 0
    elif linguality_2 == "No English":
        filter = 1
    else:
        filter = 2

    models_with_langs = data[data["language_count"] > 0]
    langs = models_with_langs["languages"].explode()
    langs = langs[langs != {}]

    d = langs.value_counts().rename_axis("language").to_frame('counts').reset_index()
    if filter == 1:
        d = d.iloc[1:]
    elif filter == 2:
        d = d.iloc[10:]

    # Just keep top 25 to avoid vertical scroll
    d = d.iloc[:25]

    st.write(alt.Chart(d).mark_bar().encode(
        x='counts',
        y=alt.X('language', sort=None)
    ))

    st.subheader("Raw Data")
    col1, col2 = st.columns(2)
    with col1:
        st.dataframe(df1)
    with col2:
        d = langs.value_counts().rename_axis("language").to_frame('counts').reset_index()
        st.dataframe(d)
    
    

with tab2:
    st.header("License info")

    no_license_count = data["license"].isna().sum()
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric(label="License Specified", value=total_samples-no_license_count)
    with col2:
        st.metric(label="No license Specified", value=no_license_count)
    with col3:
        st.metric(label="Total Unique Licenses", value=len(data["license"].unique()))

    st.subheader("Distribution of licenses per model repo")
    license_filter = st.selectbox(
        'All or filtered',
        ["All", "No Apache 2.0", "Remove top 10"])

    filter = 0
    if license_filter == "All":
        filter = 0
    elif license_filter == "No Apache 2.0":
        filter = 1
    else:
        filter = 2

    d = data["license"].value_counts().rename_axis("license").to_frame('counts').reset_index()
    if filter == 1:
        d = d.iloc[1:]
    elif filter == 2:
        d = d.iloc[10:]

    # Just keep top 25 to avoid vertical scroll
    d = d.iloc[:25]

    st.write(alt.Chart(d).mark_bar().encode(
        x='counts',
        y=alt.X('license', sort=None)
    ))
    st.text("There are some edge cases, as old repos using lists of licenses.")

    st.subheader("Raw Data")
    d = data["license"].value_counts().rename_axis("license").to_frame('counts').reset_index()
    st.dataframe(d)
    
with tab3:
    st.header("Pipeline info")

    tags = data["tags"].explode()
    tags = tags[tags.notna()].value_counts().rename_axis("tag").to_frame('counts').reset_index()
    s = tags["tag"]
    s = s[s.apply(type) == str]
    unique_tags = len(s.unique())

    no_pipeline_count = data["pipeline"].isna().sum()
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric(label="# models that have any pipeline", value=total_samples-no_pipeline_count)
    with col2:
        st.metric(label="No pipeline Specified", value=no_pipeline_count)
    with col3:
        st.metric(label="Total Unique Pipelines", value=len(data["pipeline"].unique()))

    pipeline_filter = st.selectbox(
        'Modalities',
        ["All", "NLP", "CV", "Audio", "RL", "Multimodal", "Tabular"])

    filter = 0
    if pipeline_filter == "All":
        filter = 0
    elif pipeline_filter == "NLP":
        filter = 1
    elif pipeline_filter == "CV":
        filter = 2
    elif pipeline_filter == "Audio":
        filter = 3
    elif pipeline_filter == "RL":
        filter = 4
    elif pipeline_filter == "Multimodal":
        filter = 5
    elif pipeline_filter == "Tabular":
        filter = 6

    st.subheader("High-level metrics")
    filtered_data = data[data['pipeline'].notna()]

    if filter == 1:
        filtered_data = data[data["modality"] == "nlp"]
    elif filter == 2:
        filtered_data = data[data["modality"] == "cv"]
    elif filter == 3:
        filtered_data = data[data["modality"] == "audio"]
    elif filter == 4:
        filtered_data = data[data["modality"] == "rl"]
    elif filter == 5:
        filtered_data = data[data["modality"] == "multimodal"]
    elif filter == 6:
        filtered_data = data[data["modality"] == "tabular"]

    col1, col2, col3 = st.columns(3)
    with col1:
        p = st.selectbox(
            'What pipeline do you want to see?',
            ["all", *filtered_data["pipeline"].unique()]
        )
    with col2:
        l = st.selectbox(
            'What library do you want to see?',
            ["all", *filtered_data["library"].unique()]
        )
    with col3:
        f = st.selectbox(
            'What framework support? (transformers)',
            ["all", "py", "tf", "jax"]
        ) 

    col1, col2 = st.columns(2)
    with col1:
        filt = st.multiselect(
            label="Tags (All by default)",
            options=s.unique(),
            default=None)
    with col2:
        o = st.selectbox(
            label="Operation (for tags)",
            options=["Any", "All", "None"]
        )

    def filter_fn(row):
        tags = row["tags"]
        tags[:] = [d for d in tags if isinstance(d, str)]
        if o == "All":
            if all(elem in tags for elem in filt):
                return True

        s1 = set(tags)
        s2 = set(filt)
        if o == "Any":
            if bool(s1 & s2):
                return True
        if o == "None":
            if len(s1.intersection(s2)) == 0:
                return True
        return False

    
    if p != "all":
        filtered_data = filtered_data[filtered_data["pipeline"] == p]
    if l != "all":
        filtered_data = filtered_data[filtered_data["library"] == l]
    if f != "all":
        if f == "py":
            filtered_data = filtered_data[filtered_data["pytorch"] == 1]
        elif f == "tf":
            filtered_data = filtered_data[filtered_data["tensorflow"] == 1]
        elif f == "jax":
            filtered_data = filtered_data[filtered_data["jax"] == 1]
    if filt != []:
        filtered_data = filtered_data[filtered_data.apply(filter_fn, axis=1)]


    d = filtered_data["pipeline"].value_counts().rename_axis("pipeline").to_frame('counts').reset_index()
    columns_of_interest = ["downloads_30d", "likes", "pytorch", "tensorflow", "jax"]
    grouped_data = filtered_data.groupby("pipeline").sum()[columns_of_interest]
    final_data = pd.merge(
        d, grouped_data, how="outer", on="pipeline"
    )
    sums = grouped_data.sum()

    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric(label="Total models", value=filtered_data.shape[0])
    with col2:
        st.metric(label="Cumulative Downloads (30d)", value=sums["downloads_30d"])
    with col3:
        st.metric(label="Cumulative likes", value=sums["likes"])

    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric(label="Total in PT", value=sums["pytorch"])
    with col2:
        st.metric(label="Total in TF", value=sums["tensorflow"])
    with col3:
        st.metric(label="Total in JAX", value=sums["jax"])
    
    st.metric(label="Unique Tags", value=unique_tags)

    

    st.subheader("Count of models per pipeline")
    st.write(alt.Chart(d).mark_bar().encode(
        x='counts',
        y=alt.X('pipeline', sort=None)
    ))

    st.subheader("Aggregated data")
    st.dataframe(final_data)

    st.subheader("Most common model types (specific to transformers")
    d = filtered_data["model_type"].value_counts().rename_axis("model_type").to_frame('counts').reset_index()
    d = d.iloc[:15]
    st.write(alt.Chart(d).mark_bar().encode(
        x='counts',
        y=alt.X('model_type', sort=None)
    ))

    st.subheader("Most common library types (Learn more in library tab)")
    d = filtered_data["library"].value_counts().rename_axis("library").to_frame('counts').reset_index().head(15)
    st.write(alt.Chart(d).mark_bar().encode(
        x='counts',
        y=alt.X('library', sort=None)
    ))
    
    st.subheader("Tags by count")
    tags = filtered_data["tags"].explode()
    tags = tags[tags.notna()].value_counts().rename_axis("tag").to_frame('counts').reset_index()
    st.write(alt.Chart(tags.head(30)).mark_bar().encode(
        x='counts',
        y=alt.X('tag', sort=None)
    ))
    
    st.subheader("Raw Data")
    columns_of_interest = [
        "repo_id", "author", "model_type", "files_per_repo", "library",
        "downloads_30d", "likes", "pytorch", "tensorflow", "jax"]
    raw_data = filtered_data[columns_of_interest]
    st.dataframe(raw_data)
    
    

    # todo : add activity metric


with tab4:
    st.header("Discussions Tab info")

    columns_of_interest = ["prs_count", "prs_open", "prs_merged", "prs_closed", "discussions_count", "discussions_open", "discussions_closed"]
    sums = data[columns_of_interest].sum()

    col1, col2, col3, col4 = st.columns(4)
    with col1:
        st.metric(label="Total PRs", value=sums["prs_count"])
    with col2:
        st.metric(label="PRs opened", value=sums["prs_open"])
    with col3:
        st.metric(label="PRs merged", value=sums["prs_merged"])
    with col4:
        st.metric(label="PRs closed", value=sums["prs_closed"])

    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric(label="Total discussions", value=sums["discussions_count"])
    with col2:
        st.metric(label="Discussions open", value=sums["discussions_open"])
    with col3:
        st.metric(label="Discussions closed", value=sums["discussions_closed"])

    filtered_data = data[["repo_id", "prs_count", "prs_open", "prs_merged", "prs_closed", "discussions_count", "discussions_open", "discussions_closed"]].sort_values("prs_count", ascending=False).reset_index(drop=True)
    st.dataframe(filtered_data)

with tab5:
    st.header("Library info")

    no_library_count = data["library"].isna().sum()
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric(label="# models that have any library", value=total_samples-no_library_count)
    with col2:
        st.metric(label="No library Specified", value=no_library_count)
    with col3:
        st.metric(label="Total Unique library", value=len(data["library"].unique()))


    st.subheader("High-level metrics")
    filtered_data = data[data['library'].notna()]

    col1, col2 = st.columns(2)
    with col1:
        lib = st.selectbox(
            'What library do you want to see? ',
            ["all", *filtered_data["library"].unique()]
        )
    with col2:
        pip = st.selectbox(
            'What pipeline do you want to see? ',
            ["all", *filtered_data["pipeline"].unique()]
        )

    if pip != "all":
        filtered_data = filtered_data[filtered_data["pipeline"] == pip]
    if lib != "all":
        filtered_data = filtered_data[filtered_data["library"] == lib]


    d = filtered_data["library"].value_counts().rename_axis("library").to_frame('counts').reset_index()
    grouped_data = filtered_data.groupby("library").sum()[["downloads_30d", "likes"]]
    final_data = pd.merge(
        d, grouped_data, how="outer", on="library"
    )
    sums = grouped_data.sum()

    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric(label="Total models", value=filtered_data.shape[0])
    with col2:
        st.metric(label="Cumulative Downloads (30d)", value=sums["downloads_30d"])
    with col3:
        st.metric(label="Cumulative likes", value=sums["likes"])

    st.subheader("Most common library types (Learn more in library tab)")
    d = filtered_data["library"].value_counts().rename_axis("library").to_frame('counts').reset_index().head(15)
    st.write(alt.Chart(d).mark_bar().encode(
        x='counts',
        y=alt.X('library', sort=None)
    ))

    

    st.subheader("Aggregated Data")
    st.dataframe(final_data)
    
    st.subheader("Raw Data")
    columns_of_interest = ["repo_id", "author", "files_per_repo", "library", "downloads_30d", "likes"]
    filtered_data = filtered_data[columns_of_interest]
    st.dataframe(filtered_data)

with tab6:
    st.header("Model cards")

    columns_of_interest = ["has_model_index", "has_metadata", "has_text", "text_length"]
    rows = data.shape[0]

    cond = data["has_model_index"] | data["has_text"]
    with_model_card = data[cond]
    c_model_card = with_model_card.shape[0]
    st.subheader("High-level metrics")
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric(label="# models with model card file", value=c_model_card)
    with col2:
        st.metric(label="# models without model card file", value=rows-c_model_card)
    
    with_index = data["has_model_index"].sum()
    with col1:
        st.metric(label="# models with model index", value=with_index)
    with col2:
        st.metric(label="# models without model index", value=rows-with_index)

    with_text = data["has_text"]
    with col1:
        st.metric(label="# models with model card text", value=with_text.sum())
    with col2:
        st.metric(label="# models without model card text", value=rows-with_text.sum())

    
    st.subheader("Length (chars) of model card content")
    fig, ax = plt.subplots() 
    ax = data["length_bins"].value_counts().plot.bar()
    st.metric(label="# average length of model card (chars)", value=data[with_text]["text_length"].mean())
    st.pyplot(fig)

    st.subheader("Tags (Read more in Pipeline tab)")
    tags = data["tags"].explode()
    tags = tags[tags.notna()].value_counts().rename_axis("tag").to_frame('counts').reset_index()
    st.write(alt.Chart(tags.head(30)).mark_bar().encode(
        x='counts',
        y=alt.X('tag', sort=None)
    ))

with tab7:
    st.header("Authors")
    st.text("This info corresponds to the repos owned by the authors")
    authors = data.groupby("author").sum().drop(["text_length", "Unnamed: 0", "language_count"], axis=1).sort_values("downloads_30d", ascending=False)
    d = data["author"].value_counts().rename_axis("author").to_frame('counts').reset_index()
    final_data = pd.merge(
        d, authors, how="outer", on="author"
    )
    st.dataframe(final_data)

with tab8:
    st.header("Raw Data")
    d = data.astype(str)
    st.dataframe(d)