update hf_search

st_utils.py

@@ -1,10 +1,12 @@
 import json
 from huggingface_hub import HfApi, ModelFilter, DatasetFilter, ModelSearchArguments
 from pprint import pprint
-from hf_search import hf_search
+from hf_search import HFSearch
 import streamlit as st
 import itertools
 
+hf_search = HFSearch(top_k=32)
+
 @st.cache
 def hf_api(query, limit=5, filters={}):
     print("query", query)
@@ -31,7 +33,6 @@ def hf_api(query, limit=5, filters={}):
     count = len(hits)
     if len(hits) > limit:
         hits = hits[:limit]
-    pprint(hits)
     return {"hits": hits, "count": count}
 
 
@@ -41,7 +42,7 @@ def semantic_search(query, limit=5, filters={}):
     print("filters", filters)
     print("limit", limit)
 
-    hits = hf_search(query=query, method="retrieve & rerank", limit=limit, filters=filters)
+    hits = hf_search.search(query=query, method="retrieve & rerank", limit=limit, filters=filters)
     hits = [
         {
             "modelId": hit["modelId"],
@@ -62,13 +63,13 @@ def bm25_search(query, limit=5, filters={}):
     print("limit", limit)
 
     # TODO: filters
-    hits = hf_search(query=query, method="bm25", limit=limit)
+    hits = hf_search.search(query=query, method="bm25", limit=limit, filters=filters)
     hits = [
         {
             "modelId": hit["modelId"],
             "tags": hit["tags"],
-            "downloads": hit["downloads"],
-            "likes": hit["likes"],
+            "downloads": int(hit["downloads"]),
+            "likes": int(hit["likes"]),
             "readme": hit.get("readme", None),
         }
         for hit in hits