add hf_search

.gitattributes

@@ -0,0 +1 @@
+*.pkl filter=lfs diff=lfs merge=lfs -text

requirements.txt

@@ -2,4 +2,5 @@ altair
 pandas
 streamlit
 huggingface_hub
-numerize
+numerize
+git+https://github.com/NouamaneTazi/hf_search hf_search

server/api.py

@@ -2,6 +2,7 @@ from flask import Flask, request
 import json
 from huggingface_hub import HfApi, ModelFilter, DatasetFilter, ModelSearchArguments
 from pprint import pprint
+from hf_search import hf_search
 
 app = Flask(__name__)
 
@@ -27,10 +28,10 @@ def hf_api():
         library=filters["library"],
     )
     models = api.list_models(search=query, filter=filt, limit=limit, full=True)
-    res = []
+    hits = []
     for model in models:
         model = model.__dict__
-        res.append(
+        hits.append(
             {
                 "modelId": model.get("modelId"),
                 "tags": model.get("tags"),
@@ -38,11 +39,11 @@ def hf_api():
                 "likes": model.get("likes"),
             }
         )
-    count = len(res)
-    if len(res) > limit:
-        res = res[:limit]
-    pprint(res)
-    return json.dumps({"value": res, "count": count})
+    count = len(hits)
+    if len(hits) > limit:
+        hits = hits[:limit]
+    pprint(hits)
+    return json.dumps({"value": hits, "count": count})
 
 
 @app.route("/custom/search", methods=["POST"])
@@ -55,10 +56,10 @@ def main():
     print("filters", filters)
     print("limit", limit)
 
-    # records, count_filtered = search_query(query=request_data["query"], filters=filters, top=request_data["top"])
-    # assert len(set([record["id"] for record in records])) == len(records), "ids of results are not unique"
-    # res = {"value": records, "count": count_filtered}
-    # return json.dumps(res, indent=2)
+    # TODO: filters
+    hits = hf_search(query=query, method="retrieve & rerank", limit=limit)
+
+    return json.dumps({"value": hits, "count": len(hits)})
 
 
 if __name__ == "__main__":