draft

app.py

@@ -1,9 +1,10 @@
 import datetime
 import os
+import copy
 from dataclasses import asdict, dataclass
 from functools import lru_cache
 from json import JSONDecodeError
-from typing import List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 from huggingface_hub.utils import GatedRepoError
 import gradio as gr
 from requests.exceptions import HTTPError
@@ -16,9 +17,18 @@ from huggingface_hub import (
     logging,
     model_info,
 )
+from tqdm.auto import tqdm
+from tqdm.contrib.concurrent import thread_map
+import backoff
 from huggingface_hub.utils import EntryNotFoundError, disable_progress_bars
 import httpx
 import orjson
+import httpx
+from functools import lru_cache
+
+from sys import platform
+
+CACHE_DIR = "./cache" if platform == "darwin" else "/data/"
 
 disable_progress_bars()
 
@@ -26,6 +36,10 @@ logging.set_verbosity_error()
 
 token = os.getenv("HF_TOKEN")
 
+cache = Cache(
+    CACHE_DIR, eviction_policy="least-frequently-used", size_limit=4e9
+)  # 4gb in bytes
+
 
 def get_model_labels(model):
     try:
@@ -54,11 +68,11 @@ def _get_engagement_stats(hub_id):
 
 def _try_load_model_card(hub_id):
     try:
-        url = hf_hub_url(repo_id=hub_id, filename="README.md")
+        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 = httpx.get(url).text
         length = len(card_text)
-        # card_text = ModelCard.load(hub_id, token=token).text
-        # length = len(card_text)
     except EntryNotFoundError:
         card_text = None
         length = None
@@ -97,7 +111,6 @@ class ModelMetadata:
     model_card_length: Optional[int] = None
 
     @classmethod
-    @lru_cache()
     def from_hub(cls, hub_id):
         try:
             model = model_info(hub_id)
@@ -109,10 +122,6 @@ class ModelMetadata:
             library_name = model.library_name
         except AttributeError:
             library_name = None
-        # try:
-        #     tags = model.tags
-        # except AttributeError:
-        #     tags = None
         try:
             pipeline_tag = model.pipeline_tag
         except AttributeError:
@@ -236,7 +245,7 @@ ALL_PIPELINES = {
 def generate_task_scores_dict():
     task_scores = {}
     for task in ALL_PIPELINES:
-        task_dict = COMMON_SCORES.copy()
+        task_dict = copy.deepcopy(COMMON_SCORES)
         if task in TASK_TYPES_WITH_LANGUAGES:
             task_dict = {
                 **task_dict,
@@ -271,81 +280,328 @@ def generate_task_scores_dict():
     return task_scores
 
 
-SCORES = generate_task_scores_dict()
+@lru_cache(maxsize=None)
+def generate_common_scores():
+    GENERIC_SCORES = copy.deepcopy(COMMON_SCORES)
+    GENERIC_SCORES["_max_score"] = sum(
+        value["score"] for value in GENERIC_SCORES.values()
+    )
+    return GENERIC_SCORES
 
 
-# cache = Cache("/data/")
+SCORES = generate_task_scores_dict()
+GENERIC_SCORES = generate_common_scores()
 
 
-# @cache.memoize(expire=60 * 60 * 24 * 3)  # expires after 3 days
+@cache.memoize(expire=60 * 60 * 24 * 3)  # expires after 3 days
 def _basic_check(hub_id):
-    try:
-        data = ModelMetadata.from_hub(hub_id)
-        score = 0
-        if task := data.pipeline_tag:
-            task_scores = SCORES[task]
-            to_fix = {}
-            data_dict = asdict(data)
-            for k, v in task_scores.items():
-                if k.startswith("_"):
-                    continue
-                if data_dict[k] is None:
-                    to_fix[k] = task_scores[k]["missing_recommendation"].replace(
-                        "HUB_ID", hub_id
-                    )
-                if data_dict[k] is not None:
-                    score += v["score"]
-            max_score = task_scores["_max_score"]
-            score = score / max_score
-            (
-                f"Your model's metadata score is {round(score*100)}% based on suggested"
-                f" metadata for {task}. \n"
+    data = ModelMetadata.from_hub(hub_id)
+    score = 0
+    if data is None:
+        return None
+    to_fix = {}
+    if task := data.pipeline_tag:
+        task_scores = SCORES[task]
+        data_dict = asdict(data)
+        for k, v in task_scores.items():
+            if k.startswith("_"):
+                continue
+            if data_dict[k] is None:
+                to_fix[k] = task_scores[k]["missing_recommendation"].replace(
+                    "HUB_ID", hub_id
+                )
+            if data_dict[k] is not None:
+                score += v["score"]
+        max_score = task_scores["_max_score"]
+        score = score / max_score
+        (
+            f"Your model's metadata score is {round(score*100)}% based on suggested"
+            f" metadata for {task}. \n"
+        )
+        if to_fix:
+            recommendations = (
+                "Here are some suggestions to improve your model's metadata for"
+                f" {task}: \n"
             )
-            # recommendations = []
-            if to_fix:
-                recommendations = (
-                    "Here are some suggestions to improve your model's metadata for"
-                    f" {task}: \n"
+            for v in to_fix.values():
+                recommendations += f"\n- {v}"
+            data_dict["recommendations"] = recommendations
+        data_dict["score"] = score * 100
+    else:
+        data_dict = asdict(data)
+        for k, v in GENERIC_SCORES.items():
+            if k.startswith("_"):
+                continue
+            if data_dict[k] is None:
+                to_fix[k] = GENERIC_SCORES[k]["missing_recommendation"].replace(
+                    "HUB_ID", hub_id
                 )
-                for v in to_fix.values():
-                    recommendations += f"\n- {v}"
-                data_dict["score"] = score
-                data_dict["recommendations"] = recommendations
-                return orjson.dumps(data_dict)
-            
-    except Exception as e:
-        print(e)
-        return None
+            if data_dict[k] is not None:
+                score += v["score"]
+        score = score / GENERIC_SCORES["_max_score"]
+        data_dict["score"] = max(
+            0, (score / 2) * 100
+        )  # TODO currently setting a manual penalty for not having a task
+
+    return orjson.dumps(data_dict)
 
 
 def basic_check(hub_id):
     return _basic_check(hub_id)
 
 
-# print("caching models...")
-# print("getting top 5,000 models")
-# models = list_models(sort="downloads", direction=-1, limit=5_000)
-# model_ids = [model.modelId for model in models]
-# print("calculating metadata scores...")
-# thread_map(basic_check, model_ids)
+def create_query_url(query, skip=0):
+    return f"https://huggingface.co/api/search/full-text?q={query}&limit=100&skip={skip}&type=model"
+
+
+@cache.memoize(expire=60 * 60 * 24 * 3)  # expires after 3 days
+def get_results(query) -> Dict[Any, Any]:
+    url = create_query_url(query)
+    r = httpx.get(url)
+    return r.json()
+
+
+# result = {
+#     "repoId": "621ffdc036468d709f175eb5",
+#     "repoOwnerId": "60d099234330bad169e611f0",
+#     "isPrivate": False,
+#     "type": "model",
+#     "likes": 0,
+#     "isReadmeFile": True,
+#     "readmeStartLine": 8,
+#     "updatedAt": 1687806057107,
+#     "repoName": "hate_speech_en",
+#     "repoOwner": "IMSyPP",
+#     "tags": "pytorch, bert, text-classification, en, transformers, license:mit, has_space",
+#     "name": "IMSyPP/hate_speech_en",
+#     "fileName": "README.md",
+#     "formatted": {
+#         "repoName": [{"text": "hate_speech_en", "type": "text"}],
+#         "repoOwner": [{"text": "IMSyPP", "type": "text"}],
+#         "fileContent": [
+#             {"text": "\n# ", "type": "text"},
+#             {"text": "Hate", "type": "highlight"},
+#             {"text": " ", "type": "text"},
+#             {"text": "Speech", "type": "highlight"},
+#             {
+#                 "text": " Classifier for Social Media Content in English Language\n\nA monolingual model for ",
+#                 "type": "text",
+#             },
+#             {"text": "hate", "type": "highlight"},
+#             {"text": " ", "type": "text"},
+#             {"text": "speech", "type": "highlight"},
+#             {
+#                 "text": " classification of social media content in English language. The model was trained on 103190 YouTube comments and tested on an independent test set of 20554 YouTube comments. It is based on English BERT base pre-trained language model.\n\n## Tokenizer\n\nDuring training the text was preprocessed using the original English BERT base tokenizer. We suggest the same tokenizer is used for inference.\n\n## Model output\n\nThe model classifies each input into one of four distinct classes:\n* 0 - acceptable\n* 1 - inappropriate\n* 2 - offensive\n* 3 - violent",
+#                 "type": "text",
+#             },
+#         ],
+#         "tags": [
+#             {
+#                 "text": "pytorch, bert, text-classification, en, transformers, license:mit, has_space",
+#                 "type": "text",
+#             }
+#         ],
+#         "name": [{"text": "IMSyPP/hate_speech_en", "type": "text"}],
+#         "fileName": [{"text": "README.md", "type": "text"}],
+#     },
+#     "authorData": {
+#         "avatarUrl": "https://aeiljuispo.cloudimg.io/v7/https://s3.amazonaws.com/moonup/production/uploads/1624284535629-60d08803565dd1d0867f7a37.png?w=200&h=200&f=face",
+#         "fullname": "IMSyPP EU REC AG project 875263 - Innovative Monitoring Systems and Prevention Policies of Online Hate Speech",
+#         "name": "IMSyPP",
+#         "type": "org",
+#         "isHf": False,
+#     },
+# }
+
+
+@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
+    }
+
+
+def filter_search_results(
+    results: List[Dict[Any, Any]], min_score=None, min_model_card_length=None
+):  # TODO make code more intuitive
+    results = thread_map(parse_single_result, results)
+    for i, parsed_result in tqdm(enumerate(results)):
+        # parsed_result = parse_single_result(result)
+        if parsed_result is None:
+            continue
+        if (
+            min_score is None
+            and min_model_card_length is not None
+            and parsed_result["model_card_length"] > min_model_card_length
+            or min_score is None
+            and min_model_card_length is None
+        ):
+            yield parsed_result
+        elif min_score is not None:
+            if parsed_result["metadata_score"] <= min_score:
+                continue
+            if (
+                min_model_card_length is not None
+                and parsed_result["model_card_length"] > min_model_card_length
+                or min_model_card_length is None
+            ):
+                parsed_result["original_position"] = i
+                yield parsed_result
+
+
+def sort_search_results(filtered_search_results):
+    return sorted(
+        list(filtered_search_results),
+        key=lambda x: (x["metadata_score"], x["original_position"]),
+        reverse=True,
+    )
 
-with gr.Blocks() as demo:
-    gr.Markdown(
-        """
-# Model Metadata Checker
 
-This app will check your model's metadata for a few common issues."""
+def find_context(text, query, window_size):
+    # Split the text into words
+    words = text.split()
+    # Find the index of the query token
+    try:
+        index = words.index(query)
+        # Get the start and end indices of the context window
+        start = max(0, index - window_size)
+        end = min(len(words), index + window_size + 1)
+
+        return " ".join(words[start:end])
+    except ValueError:
+        return " ".join(words[:window_size])
+
+
+# single_result[
+#     "text"
+# ] = "lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua."
+
+# results = [single_result] * 3
+
+
+def create_markdown(results):
+    rows = []
+    for result in results:
+        row = f"""# [{result['name']}]({result['repo_hub_url']})
+| Metadata Quality Score | Model card length | Licensed |
+|------------------------|-------------------|----------|
+| {result['metadata_score']:.0f}%                  |    {result['model_card_length']}  | {"&#9989;" if result['is_licensed'] else "&#10060;"} |
+\n
+*{result['text']}*
+
+<hr>
+\n"""
+        rows.append(row)
+    return "\n".join(rows)
+
+
+def get_result_card_snippet(result):
+    try:
+        result_text = httpx.get(result["search_result_file_url"]).text
+        result["text"] = find_context(result_text, query, 100)
+    except httpx.ConnectError:
+        result["text"] = "Could not load model card"
+    return result
+
+
+def _search_hub(
+    query: str,
+    min_score: Optional[int] = None,
+    min_model_card_length: Optional[int] = None,
+):
+    results = get_results(query)
+    print(f"Found {len(results['hits'])} results")
+    results = results["hits"]
+    number_original_results = len(results)
+    filtered_results = filter_search_results(
+        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)
+    percent_of_original = round(
+        len(final_results) / number_original_results * 100, ndigits=0
+    )
+    filtered_vs_og = f"""
+| Number of original results | Number of results after filtering | Percentage of results after filtering        |
+| -------------------------- | --------------------------------- | -------------------------------------------- |
+| {number_original_results}  | {len(final_results)}              | {percent_of_original}%                       |
+
+"""
+    print(final_results)
+    return filtered_vs_og, create_markdown(final_results)
+
+
+def search_hub(query: str, min_score=None, min_model_card_length=None):
+    return _search_hub(query, min_score, min_model_card_length)
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("#  &#129303; Hub model search with metadata quality filters")
     with gr.Row():
-        text = gr.Text(label="Model ID")
-        button = gr.Button(label="Check", type="submit")
-    with gr.Row():
-        gr.Markdown("Results")
-        markdown = gr.JSON()
-        button.click(_basic_check, text, markdown)
+        with gr.Column():
+            query = gr.Textbox("x-ray", label="Search query")
+        with gr.Column():
+            button = gr.Button("Search")
+            with gr.Row():
+                # gr.Checkbox(False, label="Must have licence?")
+                mim_model_card_length = gr.Number(
+                    None, label="Minimum model card length"
+                )
+                min_metadata_score = gr.Slider(0, label="Minimum metadata score")
+    filter_results = gr.Markdown("Filter results vs original search")
+    results_markdown = gr.Markdown("Search results")
+
+    button.click(
+        search_hub,
+        [query, min_metadata_score, mim_model_card_length],
+        [filter_results, results_markdown],
+    )
 
 demo.queue(concurrency_count=32)
 demo.launch()
 
 
+# with gr.Blocks() as demo:
+#     gr.Markdown(
+#         """
+# # Model Metadata Checker
+
+# This app will check your model's metadata for a few common issues."""
+#     )
+#     with gr.Row():
+#         text = gr.Text(label="Model ID")
+#         button = gr.Button(label="Check", type="submit")
+#     with gr.Row():
+#         gr.Markdown("Results")
+#         markdown = gr.JSON()
+#         button.click(_basic_check, text, markdown)
+
+# demo.queue(concurrency_count=32)
+# demo.launch()
+
+
 # gr.Interface(fn=basic_check, inputs="text", outputs="markdown").launch(debug=True)

requirements.txt

@@ -1,2 +1,146 @@
+#
+# This file is autogenerated by pip-compile with Python 3.11
+# by the following command:
+#
+#    pip-compile --resolver=backtracking
+#
+anyio==3.7.0
+    # via httpcore
+appnope==0.1.3
+    # via
+    #   ipykernel
+    #   ipython
+asttokens==2.2.1
+    # via stack-data
+backcall==0.2.0
+    # via ipython
+backoff==2.2.1
+    # via -r requirements.in
+certifi==2023.5.7
+    # via
+    #   httpcore
+    #   httpx
+    #   requests
+charset-normalizer==3.1.0
+    # via requests
+comm==0.1.3
+    # via ipykernel
+debugpy==1.6.7
+    # via ipykernel
+decorator==5.1.1
+    # via ipython
 diskcache==5.6.1
-orjson
+    # via -r requirements.in
+executing==1.2.0
+    # via stack-data
+filelock==3.12.2
+    # via huggingface-hub
+fsspec==2023.6.0
+    # via huggingface-hub
+h11==0.14.0
+    # via httpcore
+httpcore==0.17.2
+    # via httpx
+httpx==0.24.1
+    # via -r requirements.in
+huggingface-hub==0.15.1
+    # via -r requirements.in
+idna==3.4
+    # via
+    #   anyio
+    #   httpx
+    #   requests
+ipykernel==6.23.3
+    # via ipywidgets
+ipython==8.14.0
+    # via
+    #   ipykernel
+    #   ipywidgets
+ipywidgets==8.0.6
+    # via -r requirements.in
+jedi==0.18.2
+    # via ipython
+jupyter-client==8.3.0
+    # via ipykernel
+jupyter-core==5.3.1
+    # via
+    #   ipykernel
+    #   jupyter-client
+jupyterlab-widgets==3.0.7
+    # via ipywidgets
+matplotlib-inline==0.1.6
+    # via
+    #   ipykernel
+    #   ipython
+nest-asyncio==1.5.6
+    # via ipykernel
+orjson==3.9.1
+    # via -r requirements.in
+packaging==23.1
+    # via
+    #   huggingface-hub
+    #   ipykernel
+parso==0.8.3
+    # via jedi
+pexpect==4.8.0
+    # via ipython
+pickleshare==0.7.5
+    # via ipython
+platformdirs==3.8.0
+    # via jupyter-core
+prompt-toolkit==3.0.38
+    # via ipython
+psutil==5.9.5
+    # via ipykernel
+ptyprocess==0.7.0
+    # via pexpect
+pure-eval==0.2.2
+    # via stack-data
+pygments==2.15.1
+    # via ipython
+python-dateutil==2.8.2
+    # via jupyter-client
+pyyaml==6.0
+    # via huggingface-hub
+pyzmq==25.1.0
+    # via
+    #   ipykernel
+    #   jupyter-client
+requests==2.31.0
+    # via huggingface-hub
+six==1.16.0
+    # via
+    #   asttokens
+    #   python-dateutil
+sniffio==1.3.0
+    # via
+    #   anyio
+    #   httpcore
+    #   httpx
+stack-data==0.6.2
+    # via ipython
+toolz==0.12.0
+    # via -r requirements.in
+tornado==6.3.2
+    # via
+    #   ipykernel
+    #   jupyter-client
+tqdm==4.65.0
+    # via huggingface-hub
+traitlets==5.9.0
+    # via
+    #   comm
+    #   ipykernel
+    #   ipython
+    #   ipywidgets
+    #   jupyter-client
+    #   jupyter-core
+    #   matplotlib-inline
+typing-extensions==4.7.0
+    # via huggingface-hub
+urllib3==2.0.3
+    # via requests
+wcwidth==0.2.6
+    # via prompt-toolkit
+widgetsnbextension==4.0.7
+    # via ipywidgets