app update

app.py

@@ -3,47 +3,73 @@ import streamlit as st
 import math
 
 
-def setup_page():
-    st.title("Huggingface model explorer")
-    st.text(f"search {len(models_df)} models by name or readme")
-    st.text(
-        "note that there are many more models but here we only show those with readme"
-    )
+class ModelFinder:
+    def __init__(self, models_df):
+        self.setup_inputs()
+        self.models_df = models_df
+        self.n_per_page = 10
 
+    def setup_page(self):
+        st.title("Huggingface model explorer")
+        st.text(f"search {len(models_df)} models by name or readme")
+        st.text(
+            "note that there are many more models but here we only show those with readme"
+        )
 
-def setup_inputs():
-    col1, col2, col3 = st.columns(3)
-    query = col1.text_input("model name query", value="")
-    readme_query = col2.text_input("readme query", value="")
-    return query, readme_query, col3 
+    def setup_inputs(self):
+        col1, col2, col3, col4, col5 = st.columns(5)
+        self.query_input = col1.text_input("model name query", value="")
+        self.author_query_input = col2.text_input("author query", value="")
+        self.id_query_input = col3.text_input("modelId query", value="")
+        self.readme_query_input = col4.text_input("readme query", value="")
+        self.page = col5
 
+    def get_selected_models_df(self, query, readme_query, id_query, author_query):
+        return self.models_df[
+            self.models_df["readme"].str.lower().str.contains(readme_query)
+            & self.models_df["modelId"].str.lower().str.contains(id_query)
+            & self.models_df["author"].str.lower().str.contains(author_query)
+            & self.models_df["model_name"].str.lower().str.contains(query)
+        ]
 
-def get_selected_models_df(query, readme_query):
-    return models_df[
-        models_df["readme"].str.lower().str.contains(readme_query)
-        & models_df["modelId"].str.lower().str.contains(query)
-    ]
+    def show_paged_selected_model_info(self, selected_models_df):
+        page = self.page.number_input("page", 0, math.ceil(len(selected_models_df) / 10))
+        selected_models_df_subset = selected_models_df.iloc[
+            page * self.n_per_page : (page + 1) * self.n_per_page
+        ]
+        st.write(f"found {len(selected_models_df)} models")
+        for (model_name, tag, readme) in selected_models_df_subset[
+            ["modelId", "pipeline_tag", "readme"]
+        ].itertuples(index=False):
+            with st.expander(f"{model_name} ({tag})"):
+                st.write(readme)
 
+    def run(self):
+        self.setup_page()
+        selected_models_df = self.get_selected_models_df(
+            self.query_input,
+            self.readme_query_input,
+            self.id_query_input,
+            self.author_query_input,
+        )
+        self.show_paged_selected_model_info(selected_models_df)
 
-def show_paged_selected_model_info(models_df, page_col):
-    page = page_col.number_input("page", 0, math.ceil(len(selected_models_df) / 10))
-    selected_models_df_subset = models_df.iloc[
-        page * n_per_page : (page+ 1) * n_per_page
-    ]
-    st.write(f"found {len(selected_models_df)} models")
-    for (model_name, tag, readme) in selected_models_df_subset[
-        ["modelId", "pipeline_tag", "readme"]
-    ].itertuples(index=False):
-        with st.expander(f"{model_name} ({tag})"):
-            st.write(readme)
 
+def prepare_models_df(path):
+    df = pd.read_parquet(path).dropna(subset=["readme"])
+    sep_tuples = [
+        tp if len(tp) == 2 else ("", tp[0])
+        for tp in df["modelId"].str.split("/").to_list()
+    ]
+    authors, model_names = zip(*sep_tuples)
+    df["author"] = authors
+    df["model_name"] = model_names
+    return df
 
-models_df = pd.read_csv("models_with_readmes.csv").dropna(subset=["readme"])
 
-setup_page()
-query, readme_query, page_col = setup_inputs()
-selected_models_df = get_selected_models_df(query, readme_query)
+model_path = "models_with_readmes.parquet"
+models_df = prepare_models_df(model_path)
 
-n_per_page = 10
+app = ModelFinder(models_df)
 
-show_paged_selected_model_info(selected_models_df, page_col)
+app.run()