async version

app.py

@@ -1,98 +1,73 @@
+import asyncio
 import copy
 import os
 from dataclasses import asdict, dataclass
 from datetime import datetime, timedelta
-from functools import lru_cache
 from json import JSONDecodeError
 from typing import Any, Dict, List, Optional, Union
-
-import backoff
 import gradio as gr
 import httpx
 import orjson
-import requests
-from cachetools import TTLCache, cached
-from httpx import Client
-from httpx_caching import CachingClient, OneDayCacheHeuristic
-# from diskcache import Cache
-from huggingface_hub import (HfApi, hf_hub_url, list_repo_commits, logging,
-                             model_info)
-from huggingface_hub.utils import (EntryNotFoundError, GatedRepoError,
-                                   disable_progress_bars)
-from requests.exceptions import HTTPError
+from cashews import NOT_NONE, cache
+from httpx import AsyncClient
+from huggingface_hub import hf_hub_url, logging
+from huggingface_hub.utils import disable_progress_bars
+from rich import print
 from tqdm.auto import tqdm
-from tqdm.contrib.concurrent import thread_map
-
-cache = TTLCache(maxsize=500_000, ttl=timedelta(hours=24), timer=datetime.now)
-
-client = Client()
-
-
-client = CachingClient(client, heuristic=OneDayCacheHeuristic())
+from httpx import Client
+from datetime import datetime, timedelta
 
+cache.setup(
+    "mem://"
+) 
 
-# CACHE_DIR = "./cache" if platform == "darwin" else "/data/"
 
 disable_progress_bars()
 
 logging.set_verbosity_error()
 
 token = os.getenv("HF_TOKEN")
-
-# cache = Cache(CACHE_DIR)
+headers = {"authorization": f"Bearer {token}"}
 
 
-def get_model_labels(model):
+async def get_model_labels(model, client):
     try:
         url = hf_hub_url(repo_id=model, filename="config.json")
-        return list(requests.get(url).json()["label2id"].keys())
+        resp = await client.get(url, timeout=2)
+        return list(resp.json()["label2id"].keys())
     except (KeyError, JSONDecodeError, AttributeError):
         return None
 
 
-@dataclass
-class EngagementStats:
-    likes: int
-    downloads: int
-    created_at: datetime
-
-
-def _get_engagement_stats(hub_id):
-    api = HfApi(token=token)
-    repo = api.repo_info(hub_id)
-    return EngagementStats(
-        likes=repo.likes,
-        downloads=repo.downloads,
-        created_at=list_repo_commits(hub_id, repo_type="model")[-1].created_at,
-    )
-
-
-def _try_load_model_card(hub_id):
+async def _try_load_model_card(hub_id, client=None):
+    if not client:
+        client = AsyncClient(headers=headers)
     try:
         url = hf_hub_url(
             repo_id=hub_id, filename="README.md"
         )  # We grab card this way rather than via client library to improve performance
-        card_text = client.get(url).text
-        length = len(card_text)
-    except EntryNotFoundError:
-        card_text = None
-        length = None
-    except (
-        GatedRepoError
-    ):  # TODO return different values to reflect gating rather than no card
+        resp = await client.get(url)
+        if resp.status_code == 200:
+            card_text = resp.text
+            length = len(card_text)
+        elif resp.status_code == 404:
+            card_text = None
+            length = 0
+    except httpx.ConnectError:
         card_text = None
         length = None
     return card_text, length
 
-
-def _try_parse_card_data(hub_id):
+def _try_parse_card_data(hub_json_data):
     data = {}
-    keys = ["license", "language", "datasets", "tags"]
+    keys = ["license", "language", "datasets"]
     for key in keys:
-        try:
-            value = model_info(hub_id, token=token).cardData[key]
-            data[key] = value
-        except (KeyError, AttributeError):
+        if card_data := hub_json_data.get("cardData"):
+            try:
+                data[key] = card_data.get(key)
+            except (KeyError, AttributeError):
+                data[key] = None
+        else:
             data[key] = None
     return data
 
@@ -107,39 +82,46 @@ class ModelMetadata:
     pipeline_tag: Optional[str]
     labels: Optional[List[str]]
     languages: Optional[Union[str, List[str]]]
-    engagement_stats: Optional[EngagementStats] = None
     model_card_text: Optional[str] = None
     model_card_length: Optional[int] = None
+    likes: Optional[int] = None
+    downloads: Optional[int] = None
+    created_at: Optional[datetime] = None
 
     @classmethod
-    def from_hub(cls, hub_id):
+    @cache(ttl=timedelta(hours=3), condition=NOT_NONE)
+    async def from_hub(cls, hub_id, client=None):
         try:
-            model = model_info(hub_id)
-        except (GatedRepoError, HTTPError):
-            return None  # TODO catch gated repos and handle properly
-        card_text, length = _try_load_model_card(hub_id)
-        data = _try_parse_card_data(hub_id)
-        try:
-            library_name = model.library_name
-        except AttributeError:
-            library_name = None
-        try:
-            pipeline_tag = model.pipeline_tag
-        except AttributeError:
-            pipeline_tag = None
-        return ModelMetadata(
-            hub_id=hub_id,
-            languages=data["language"],
-            tags=data["tags"],
-            license=data["license"],
-            library_name=library_name,
-            datasets=data["datasets"],
-            pipeline_tag=pipeline_tag,
-            labels=get_model_labels(hub_id),
-            engagement_stats=_get_engagement_stats(hub_id),
-            model_card_text=card_text,
-            model_card_length=length,
-        )
+            if not client:
+                client = httpx.AsyncClient()
+            url = f"https://huggingface.co/api/models/{hub_id}"
+            resp = await client.get(url)
+            hub_json_data = resp.json()
+            card_text, length = await _try_load_model_card(hub_id)
+            data = _try_parse_card_data(hub_json_data)
+            library_name = hub_json_data.get("library_name")
+            pipeline_tag = hub_json_data.get("pipeline_tag")
+            downloads = hub_json_data.get("downloads")
+            likes = hub_json_data.get("likes")
+            tags = hub_json_data.get("tags")
+            labels = await get_model_labels(hub_id, client)
+            return ModelMetadata(
+                hub_id=hub_id,
+                languages=data["language"],
+                tags=tags,
+                license=data["license"],
+                library_name=library_name,
+                datasets=data["datasets"],
+                pipeline_tag=pipeline_tag,
+                labels=labels,
+                model_card_text=card_text,
+                downloads=downloads,
+                likes=likes,
+                model_card_length=length,
+            )
+        except Exception as e:
+            print(f"Failed to create ModelMetadata for model {hub_id}: {str(e)}")
+            return None
 
 
 COMMON_SCORES = {
@@ -242,7 +224,6 @@ ALL_PIPELINES = {
 }
 
 
-@lru_cache(maxsize=None)
 def generate_task_scores_dict():
     task_scores = {}
     for task in ALL_PIPELINES:
@@ -281,7 +262,6 @@ def generate_task_scores_dict():
     return task_scores
 
 
-@lru_cache(maxsize=None)
 def generate_common_scores():
     GENERIC_SCORES = copy.deepcopy(COMMON_SCORES)
     GENERIC_SCORES["_max_score"] = sum(
@@ -294,13 +274,11 @@ SCORES = generate_task_scores_dict()
 GENERIC_SCORES = generate_common_scores()
 
 
-# @cache.memoize(expire=60 * 60 * 24 * 3)  # expires after 3 days
-@cached(cache)
-def _basic_check(hub_id):
-    data = ModelMetadata.from_hub(hub_id)
+def _basic_check(data: Optional[ModelMetadata]):
     score = 0
     if data is None:
         return None
+    hub_id = data.hub_id
     to_fix = {}
     if task := data.pipeline_tag:
         task_scores = SCORES[task]
@@ -356,38 +334,48 @@ def create_query_url(query, skip=0):
     return f"https://huggingface.co/api/search/full-text?q={query}&limit=100&skip={skip}&type=model"
 
 
-@cached(cache)
-def get_results(query) -> Dict[Any, Any]:
+def get_results(query,sync_client=None) -> Dict[Any, Any]:
+    if not sync_client:
+        sync_client = Client(http2=True, headers=headers)
     url = create_query_url(query)
-    r = client.get(url)
+    r = sync_client.get(url)
     return r.json()
 
 
-@backoff.on_exception(
-    backoff.expo,
-    Exception,
-    max_time=2,
-    raise_on_giveup=False,
-)
 def parse_single_result(result):
     name, filename = result["name"], result["fileName"]
     search_result_file_url = hf_hub_url(name, filename)
     repo_hub_url = f"https://huggingface.co/{name}"
-    score = _basic_check(name)
-    if score is None:
-        return None
-    score = orjson.loads(score)
     return {
         "name": name,
         "search_result_file_url": search_result_file_url,
         "repo_hub_url": repo_hub_url,
-        "metadata_score": score["score"],
-        "model_card_length": score["model_card_length"],
-        "is_licensed": bool(score["license"]),
-        # "metadata_report": score
     }
 
 
+@cache(ttl=timedelta(hours=3), condition=NOT_NONE)
+async def get_hub_models(results, client=None):
+    parsed_results = [parse_single_result(result) for result in results]
+    if not client:
+        client = AsyncClient(http2=True, headers=headers)
+    model_ids = [result["name"] for result in parsed_results]
+    model_objs = [ModelMetadata.from_hub(model, client=client) for model in model_ids]
+    models = await asyncio.gather(*model_objs)
+    results = []
+    for result, model in zip(parsed_results, models):
+        score = _basic_check(model)
+        # print(f"score for {model} is {score}")
+        if score is not None:
+            score = orjson.loads(score)
+            result["metadata_score"] = score["score"]
+            result["model_card_length"] = score["model_card_length"]
+            result["is_licensed"] = (bool(score["license"]),)
+            results.append(result)
+        else:
+            results.append(None)
+    return results
+
+
 def filter_for_license(results):
     for result in results:
         if result["is_licensed"]:
@@ -406,7 +394,7 @@ def filter_search_results(
     min_model_card_length=None,
 ):  # TODO make code more intuitive
     # TODO setup filters as separate functions and chain results
-    results = thread_map(parse_single_result, results)
+    results = asyncio.run(get_hub_models(results))
     for i, parsed_result in tqdm(enumerate(results)):
         # parsed_result = parse_single_result(result)
         if parsed_result is None:
@@ -418,6 +406,7 @@ def filter_search_results(
             or min_score is None
             and min_model_card_length is None
         ):
+            parsed_result["original_position"] = i
             yield parsed_result
         elif min_score is not None:
             if parsed_result["metadata_score"] <= min_score:
@@ -472,22 +461,36 @@ def create_markdown(results):  # TODO move to separate file
         rows.append(row)
     return "\n".join(rows)
 
-
-def get_result_card_snippet(result):
+async def get_result_card_snippet(result, query=None, client=None):
+    if not client:
+        client = AsyncClient(http2=True, headers=headers)
     try:
-        result_text = httpx.get(result["search_result_file_url"]).text
+        resp = await client.get(result["search_result_file_url"])
+        result_text = resp.text
         result["text"] = find_context(result_text, query, 100)
     except httpx.ConnectError:
         result["text"] = "Could not load model card"
     return result
 
-@cached(cache)
+@cache(ttl=timedelta(hours=3), condition=NOT_NONE)
+async def get_result_card_snippets(results, query=None, client=None):
+    if not client:
+        client = AsyncClient(http2=True, headers=headers)
+    result_snippets = [
+        get_result_card_snippet(result, query=query, client=client)
+        for result in results
+    ]
+    results = await asyncio.gather(*result_snippets)
+    return results
+
+sync_client = Client(http2=True, headers=headers)
+
 def _search_hub(
     query: str,
     min_score: Optional[int] = None,
     min_model_card_length: Optional[int] = None,
 ):
-    results = get_results(query)
+    results = get_results(query, sync_client)
     print(f"Found {len(results['hits'])} results")
     results = results["hits"]
     number_original_results = len(results)
@@ -495,12 +498,7 @@ def _search_hub(
         results, min_score=min_score, min_model_card_length=min_model_card_length
     )
     filtered_results = sort_search_results(filtered_results)
-    # final_results = []
-    # for result in filtered_results:
-    #     result_text = httpx.get(result["search_result_file_url"]).text
-    #     result["text"] = find_context(result_text, query, 100)
-    #     final_results.append(result)
-    final_results = thread_map(get_result_card_snippet, filtered_results)
+    final_results = asyncio.run(get_result_card_snippets(filtered_results, query=query))
     percent_of_original = round(
         len(final_results) / number_original_results * 100, ndigits=0
     )
@@ -510,7 +508,6 @@ def _search_hub(
 | {number_original_results}  | {len(final_results)}              | {percent_of_original}%                       |
 
 """
-    print(final_results)
     return filtered_vs_og, create_markdown(final_results)
 
 

requirements.txt

@@ -1,24 +1,36 @@
 #
-# This file is autogenerated by pip-compile with Python 3.11
+# This file is autogenerated by pip-compile with Python 3.10
 # by the following command:
 #
 #    pip-compile --resolver=backtracking
 #
+aiofiles==23.1.0
+    # via gradio
+aiohttp==3.8.4
+    # via gradio
+aiosignal==1.3.1
+    # via aiohttp
+altair==5.0.1
+    # via gradio
 anyio==3.7.0
     # via
     #   httpcore
-    #   httpx-caching
+    #   starlette
 appnope==0.1.3
     # via
     #   ipykernel
     #   ipython
 asttokens==2.2.1
     # via stack-data
+async-timeout==4.0.2
+    # via aiohttp
+attrs==23.1.0
+    # via
+    #   aiohttp
+    #   jsonschema
 backcall==0.2.0
     # via ipython
-backoff==2.2.1
-    # via -r requirements.in
-cachetools==5.3.1
+cashews==6.2.0
     # via -r requirements.in
 certifi==2023.5.7
     # via
@@ -27,35 +39,73 @@ certifi==2023.5.7
     #   requests
 charset-normalizer==3.1.0
     # via
+    #   aiohttp
     #   httpx
     #   requests
+click==8.1.3
+    # via uvicorn
 comm==0.1.3
     # via ipykernel
+contourpy==1.1.0
+    # via matplotlib
+cycler==0.11.0
+    # via matplotlib
 debugpy==1.6.7
     # via ipykernel
 decorator==5.1.1
     # via ipython
+exceptiongroup==1.1.1
+    # via anyio
 executing==1.2.0
     # via stack-data
+fastapi==0.99.1
+    # via gradio
+ffmpy==0.3.0
+    # via gradio
 filelock==3.12.2
     # via huggingface-hub
+fonttools==4.40.0
+    # via matplotlib
+frozenlist==1.3.3
+    # via
+    #   aiohttp
+    #   aiosignal
 fsspec==2023.6.0
-    # via huggingface-hub
+    # via
+    #   gradio-client
+    #   huggingface-hub
+gradio==3.35.2
+    # via -r requirements.in
+gradio-client==0.2.7
+    # via gradio
 h11==0.12.0
-    # via httpcore
+    # via
+    #   httpcore
+    #   uvicorn
+h2==4.1.0
+    # via httpx
+hpack==4.0.0
+    # via h2
 httpcore==0.14.7
     # via httpx
-httpx==0.22.0
-    # via httpx-caching
-httpx-caching==0.1a2
-    # via -r requirements.in
+httpx[http2]==0.22.0
+    # via
+    #   -r requirements.in
+    #   gradio
+    #   gradio-client
 huggingface-hub==0.15.1
-    # via -r requirements.in
+    # via
+    #   -r requirements.in
+    #   gradio
+    #   gradio-client
+hyperframe==6.0.1
+    # via h2
 idna==3.4
     # via
     #   anyio
     #   requests
     #   rfc3986
+    #   yarl
 ipykernel==6.23.3
     # via ipywidgets
 ipython==8.14.0
@@ -66,6 +116,12 @@ ipywidgets==8.0.6
     # via -r requirements.in
 jedi==0.18.2
     # via ipython
+jinja2==3.1.2
+    # via
+    #   altair
+    #   gradio
+jsonschema==4.17.3
+    # via altair
 jupyter-client==8.3.0
     # via ipykernel
 jupyter-core==5.3.1
@@ -74,28 +130,66 @@ jupyter-core==5.3.1
     #   jupyter-client
 jupyterlab-widgets==3.0.7
     # via ipywidgets
+kiwisolver==1.4.4
+    # via matplotlib
+linkify-it-py==2.0.2
+    # via markdown-it-py
+markdown-it-py[linkify]==2.2.0
+    # via
+    #   gradio
+    #   mdit-py-plugins
+    #   rich
+markupsafe==2.1.3
+    # via
+    #   gradio
+    #   jinja2
+matplotlib==3.7.1
+    # via gradio
 matplotlib-inline==0.1.6
     # via
     #   ipykernel
     #   ipython
-msgpack==1.0.5
-    # via httpx-caching
-multimethod==1.9.1
-    # via httpx-caching
+mdit-py-plugins==0.3.3
+    # via gradio
+mdurl==0.1.2
+    # via markdown-it-py
+multidict==6.0.4
+    # via
+    #   aiohttp
+    #   yarl
 nest-asyncio==1.5.6
     # via ipykernel
+numpy==1.25.0
+    # via
+    #   altair
+    #   contourpy
+    #   gradio
+    #   matplotlib
+    #   pandas
 orjson==3.9.1
-    # via -r requirements.in
+    # via
+    #   -r requirements.in
+    #   gradio
 packaging==23.1
     # via
+    #   gradio-client
     #   huggingface-hub
     #   ipykernel
+    #   matplotlib
+pandas==2.0.3
+    # via
+    #   altair
+    #   gradio
 parso==0.8.3
     # via jedi
 pexpect==4.8.0
     # via ipython
 pickleshare==0.7.5
     # via ipython
+pillow==10.0.0
+    # via
+    #   gradio
+    #   matplotlib
 platformdirs==3.8.0
     # via jupyter-core
 prompt-toolkit==3.0.38
@@ -106,20 +200,49 @@ ptyprocess==0.7.0
     # via pexpect
 pure-eval==0.2.2
     # via stack-data
+pydantic==1.10.10
+    # via
+    #   fastapi
+    #   gradio
+pydub==0.25.1
+    # via gradio
 pygments==2.15.1
-    # via ipython
+    # via
+    #   gradio
+    #   ipython
+    #   rich
+pyparsing==3.1.0
+    # via matplotlib
+pyrsistent==0.19.3
+    # via jsonschema
 python-dateutil==2.8.2
-    # via jupyter-client
+    # via
+    #   jupyter-client
+    #   matplotlib
+    #   pandas
+python-multipart==0.0.6
+    # via gradio
+pytz==2023.3
+    # via pandas
 pyyaml==6.0
-    # via huggingface-hub
+    # via
+    #   gradio
+    #   huggingface-hub
 pyzmq==25.1.0
     # via
     #   ipykernel
     #   jupyter-client
 requests==2.31.0
-    # via huggingface-hub
+    # via
+    #   gradio
+    #   gradio-client
+    #   huggingface-hub
 rfc3986[idna2008]==1.5.0
     # via httpx
+rich==13.4.2
+    # via -r requirements.in
+semantic-version==2.10.0
+    # via gradio
 six==1.16.0
     # via
     #   asttokens
@@ -131,8 +254,12 @@ sniffio==1.3.0
     #   httpx
 stack-data==0.6.2
     # via ipython
+starlette==0.27.0
+    # via fastapi
 toolz==0.12.0
-    # via -r requirements.in
+    # via
+    #   -r requirements.in
+    #   altair
 tornado==6.3.2
     # via
     #   ipykernel
@@ -151,10 +278,27 @@ traitlets==5.9.0
     #   jupyter-core
     #   matplotlib-inline
 typing-extensions==4.7.0
-    # via huggingface-hub
+    # via
+    #   altair
+    #   fastapi
+    #   gradio-client
+    #   huggingface-hub
+    #   pydantic
+tzdata==2023.3
+    # via pandas
+uc-micro-py==1.0.2
+    # via linkify-it-py
 urllib3==2.0.3
     # via requests
+uvicorn==0.22.0
+    # via gradio
 wcwidth==0.2.6
     # via prompt-toolkit
+websockets==11.0.3
+    # via
+    #   gradio
+    #   gradio-client
 widgetsnbextension==4.0.7
     # via ipywidgets
+yarl==1.9.2
+    # via aiohttp