app.py

import datetime
import os
import copy
from dataclasses import asdict, dataclass
from functools import lru_cache
from json import JSONDecodeError
from typing import Any, Dict, List, Optional, Union
from huggingface_hub.utils import GatedRepoError
import gradio as gr
from requests.exceptions import HTTPError
import requests
from diskcache import Cache
from huggingface_hub import (
    HfApi,
    hf_hub_url,
    list_repo_commits,
    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()

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:
        url = hf_hub_url(repo_id=model, filename="config.json")
        return list(requests.get(url).json()["label2id"].keys())
    except (KeyError, JSONDecodeError, AttributeError):
        return None


@dataclass
class EngagementStats:
    likes: int
    downloads: int
    created_at: datetime.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):
    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 = httpx.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
        card_text = None
        length = None
    return card_text, length


def _try_parse_card_data(hub_id):
    data = {}
    keys = ["license", "language", "datasets", "tags"]
    for key in keys:
        try:
            value = model_info(hub_id, token=token).cardData[key]
            data[key] = value
        except (KeyError, AttributeError):
            data[key] = None
    return data


@dataclass
class ModelMetadata:
    hub_id: str
    tags: Optional[List[str]]
    license: Optional[str]
    library_name: Optional[str]
    datasets: Optional[List[str]]
    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

    @classmethod
    def from_hub(cls, hub_id):
        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,
        )


COMMON_SCORES = {
    "license": {
        "required": True,
        "score": 2,
        "missing_recommendation": (
            "You have not added a license to your models metadata"
        ),
    },
    "datasets": {
        "required": False,
        "score": 1,
        "missing_recommendation": (
            "You have not added any datasets to your models metadata"
        ),
    },
    "model_card_text": {
        "required": True,
        "score": 3,
        "missing_recommendation": """You haven't created a model card for your model. It is strongly recommended to have a model card for your model. \nYou can create for your model by clicking [here](https://huggingface.co/HUB_ID/edit/main/README.md)""",
    },
    "tags": {
        "required": False,
        "score": 2,
        "missing_recommendation": (
            "You don't have any tags defined in your model metadata. Tags can help"
            " people find relevant models on the Hub. You can create for your model by"
            " clicking [here](https://huggingface.co/HUB_ID/edit/main/README.md)"
        ),
    },
}


TASK_TYPES_WITH_LANGUAGES = {
    "text-classification",
    "token-classification",
    "table-question-answering",
    "question-answering",
    "zero-shot-classification",
    "translation",
    "summarization",
    "text-generation",
    "text2text-generation",
    "fill-mask",
    "sentence-similarity",
    "text-to-speech",
    "automatic-speech-recognition",
    "text-to-image",
    "image-to-text",
    "visual-question-answering",
    "document-question-answering",
}

LABELS_REQUIRED_TASKS = {
    "text-classification",
    "token-classification",
    "object-detection",
    "audio-classification",
    "image-classification",
    "tabular-classification",
}
ALL_PIPELINES = {
    "audio-classification",
    "audio-to-audio",
    "automatic-speech-recognition",
    "conversational",
    "depth-estimation",
    "document-question-answering",
    "feature-extraction",
    "fill-mask",
    "graph-ml",
    "image-classification",
    "image-segmentation",
    "image-to-image",
    "image-to-text",
    "object-detection",
    "question-answering",
    "reinforcement-learning",
    "robotics",
    "sentence-similarity",
    "summarization",
    "table-question-answering",
    "tabular-classification",
    "tabular-regression",
    "text-classification",
    "text-generation",
    "text-to-image",
    "text-to-speech",
    "text-to-video",
    "text2text-generation",
    "token-classification",
    "translation",
    "unconditional-image-generation",
    "video-classification",
    "visual-question-answering",
    "voice-activity-detection",
    "zero-shot-classification",
    "zero-shot-image-classification",
}


@lru_cache(maxsize=None)
def generate_task_scores_dict():
    task_scores = {}
    for task in ALL_PIPELINES:
        task_dict = copy.deepcopy(COMMON_SCORES)
        if task in TASK_TYPES_WITH_LANGUAGES:
            task_dict = {
                **task_dict,
                **{
                    "languages": {
                        "required": True,
                        "score": 2,
                        "missing_recommendation": (
                            "You haven't defined any languages in your metadata. This"
                            f" is usually recommend for {task} task"
                        ),
                    }
                },
            }
        if task in LABELS_REQUIRED_TASKS:
            task_dict = {
                **task_dict,
                **{
                    "labels": {
                        "required": True,
                        "score": 2,
                        "missing_recommendation": (
                            "You haven't defined any labels in the config.json file"
                            f" these are usually recommended for {task}"
                        ),
                    }
                },
            }
        max_score = sum(value["score"] for value in task_dict.values())
        task_dict["_max_score"] = max_score
        task_scores[task] = task_dict
    return task_scores


@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


SCORES = generate_task_scores_dict()
GENERIC_SCORES = generate_common_scores()


@cache.memoize(expire=60 * 60 * 24 * 3)  # expires after 3 days
def _basic_check(hub_id):
    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"
            )
            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
                )
            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)


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,
    )


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():
        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)