lambdaofgod/huggingface_explorer

1import pandas as pd2import streamlit as st3import math4 5 6class ModelFinder:7    def __init__(self, models_df):8        self.setup_inputs()9        self.models_df = models_df10        self.n_per_page = 1011 12    def setup_page(self):13        st.title("Huggingface model explorer")14        st.text(f"search {len(models_df)} models by name or readme")15        st.text(16            "note that there are many more models but here we only show those with readme"17        )18 19    def setup_inputs(self):20        col1, col2, col3, col4, col5 = st.columns(5)21        self.query_input = col1.text_input("model name query", value="")22        self.author_query_input = col2.text_input("author query", value="")23        self.id_query_input = col3.text_input("modelId query", value="")24        self.readme_query_input = col4.text_input("readme query", value="")25        self.page = col526 27    def get_selected_models_df(self, query, readme_query, id_query, author_query):28        return self.models_df[29            self.models_df["readme"].str.lower().str.contains(readme_query)30            & self.models_df["modelId"].str.lower().str.contains(id_query)31            & self.models_df["author"].str.lower().str.contains(author_query)32            & self.models_df["model_name"].str.lower().str.contains(query)33        ]34 35    def show_paged_selected_model_info(self, selected_models_df):36        page = self.page.number_input("page", 0, math.ceil(len(selected_models_df) / 10))37        selected_models_df_subset = selected_models_df.iloc[38            page * self.n_per_page : (page + 1) * self.n_per_page39        ]40        st.write(f"found {len(selected_models_df)} models")41        for (model_name, tag, readme) in selected_models_df_subset[42            ["modelId", "pipeline_tag", "readme"]43        ].itertuples(index=False):44            model_url = f"http://huggingface.co/{model_name}"45            with st.expander(f"[{model_name}]({model_url}) ({tag})"):46                st.write(readme)47 48    def run(self):49        self.setup_page()50        selected_models_df = self.get_selected_models_df(51            self.query_input,52            self.readme_query_input,53            self.id_query_input,54            self.author_query_input,55        )56        self.show_paged_selected_model_info(selected_models_df)57 58 59def prepare_models_df(path):60    df = pd.read_parquet(path).dropna(subset=["readme"])61    sep_tuples = [62        tp if len(tp) == 2 else ("", tp[0])63        for tp in df["modelId"].str.split("/").to_list()64    ]65    authors, model_names = zip(*sep_tuples)66    df["author"] = authors67    df["model_name"] = model_names68    return df69 70 71model_path = "models_with_readmes.parquet"72models_df = prepare_models_df(model_path)73 74app = ModelFinder(models_df)75 76app.run()77