app.py

import asyncio
import os
import re
from typing import Dict

import gradio as gr
import httpx
from cachetools import TTLCache, cached
from cashews import NOT_NONE, cache
from dotenv import load_dotenv
from httpx import AsyncClient, Limits
from huggingface_hub import (
    ModelCard,
    ModelFilter,
    get_repo_discussions,
    hf_hub_url,
    list_models,
    logging,
)
from huggingface_hub.utils import HfHubHTTPError, RepositoryNotFoundError
from tqdm.asyncio import tqdm as atqdm
from tqdm.auto import tqdm
import random

cache.setup("mem://")


load_dotenv()
token = os.environ["HUGGINGFACE_TOKEN"]
user_agent = os.environ["USER_AGENT"]
assert token
assert user_agent

headers = {"user-agent": user_agent, "authorization": f"Bearer {token}"}

limits = Limits(max_keepalive_connections=10, max_connections=50)


def create_client():
    return AsyncClient(headers=headers, limits=limits, http2=True)


@cached(cache=TTLCache(maxsize=100, ttl=60 * 10))
def get_models(user_or_org):
    model_filter = ModelFilter(library="transformers", author=user_or_org)
    return list(
        tqdm(
            iter(
                list_models(
                    filter=model_filter,
                    sort="downloads",
                    direction=-1,
                    cardData=True,
                    full=True,
                )
            )
        )
    )


def filter_models(models):
    new_models = []
    for model in tqdm(models):
        try:
            if card_data := model.cardData:
                base_model = card_data.get("base_model", None)
                if not base_model:
                    new_models.append(model)
        except AttributeError:
            continue
    return new_models


MODEL_ID_RE_PATTERN = re.compile(
    "This model is a fine-tuned version of \[(.*?)\]\(.*?\)"
)
BASE_MODEL_PATTERN = re.compile("base_model:\s+(.+)")


@cached(cache=TTLCache(maxsize=100, ttl=60 * 3))
def has_model_card(model):
    if siblings := model.siblings:
        for sibling in siblings:
            if sibling.rfilename == "README.md":
                return True
    return False


@cached(cache=TTLCache(maxsize=100, ttl=60))
def check_already_has_base_model(text):
    return bool(re.search(BASE_MODEL_PATTERN, text))


@cached(cache=TTLCache(maxsize=100, ttl=60))
def extract_model_name(text):
    return match.group(1) if (match := re.search(MODEL_ID_RE_PATTERN, text)) else None


# semaphore = asyncio.Semaphore(10)  # Maximum number of concurrent tasks


@cache(ttl=120, condition=NOT_NONE)
async def check_readme_for_match(model):
    if not has_model_card(model):
        return None
    model_card_url = hf_hub_url(model.modelId, "README.md")
    client = create_client()
    try:
        resp = await client.get(model_card_url)
        if check_already_has_base_model(resp.text):
            return None
        else:
            return None if resp.status_code != 200 else extract_model_name(resp.text)
    except httpx.ConnectError:
        return None
    except httpx.ReadTimeout:
        return None
    except httpx.ConnectTimeout:
        return None
    except Exception as e:
        print(e)
        return None


@cache(ttl=120, condition=NOT_NONE)
async def check_model_exists(model, match):
    client = create_client()
    url = f"https://huggingface.co/api/models/{match}"
    try:
        resp = await client.get(url)
        if resp.status_code == 200:
            return {"modelid": model.modelId, "match": match}
        if resp.status_code == 401:
            return False
    except httpx.ConnectError:
        return None
    except httpx.ReadTimeout:
        return None
    except httpx.ConnectTimeout:
        return None
    except Exception as e:
        print(e)
        return None


@cache(ttl=120, condition=NOT_NONE)
async def check_model(model):
    match = await check_readme_for_match(model)
    if match:
        return await check_model_exists(model, match)


async def prep_tasks(models):
    tasks = []
    for model in models:
        task = asyncio.create_task(check_model(model))
        tasks.append(task)
    return [await f for f in atqdm.as_completed(tasks)]


def get_data_for_user(user_or_org):
    models = get_models(user_or_org)
    models = filter_models(models)
    results = asyncio.run(prep_tasks(models))
    results = [r for r in results if r is not None]
    return results


logger = logging.get_logger()

token = os.getenv("HUGGINGFACE_TOKEN")


def generate_issue_text(based_model_regex_match, opened_by=None):
    return f"""This pull request aims to enrich the metadata of your model by adding [`{based_model_regex_match}`](https://huggingface.co/{based_model_regex_match}) as a `base_model` field, situated in the `YAML` block of your model's `README.md`.

How did we find this information? We performed a regular expression match on your `README.md` file to determine the connection.

**Why add this?** Enhancing your model's metadata in this way:
- **Boosts Discoverability** - It becomes straightforward to trace the relationships between various models on the Hugging Face Hub.
- **Highlights Impact** - It showcases the contributions and influences different models have within the community.

For a hands-on example of how such metadata can play a pivotal role in mapping model connections, take a look at [librarian-bots/base_model_explorer](https://huggingface.co/spaces/librarian-bots/base_model_explorer).

This PR comes courtesy of [Librarian Bot](https://huggingface.co/librarian-bot) by request of {opened_by}"""


def update_metadata(metadata_payload: Dict[str, str], user_making_request=None):
    metadata_payload["opened_pr"] = False
    regex_match = metadata_payload["match"]
    repo_id = metadata_payload["modelid"]
    try:
        model_card = ModelCard.load(repo_id)
    except RepositoryNotFoundError:
        return metadata_payload
    model_card.data["base_model"] = regex_match
    template = generate_issue_text(regex_match, opened_by=user_making_request)
    try:
        if previous_discussions := list(get_repo_discussions(repo_id)):
            logger.info("found previous discussions")
            if prs := [
                discussion
                for discussion in previous_discussions
                if discussion.is_pull_request
            ]:
                logger.info("found previous pull requests")
                for pr in prs:
                    if pr.author == "librarian-bot":
                        logger.info("previously opened PR")
                        if (
                            pr.title
                            == "Librarian Bot: Add base_model information to model"
                        ):
                            logger.info("previously opened PR to add base_model tag")
                            metadata_payload["opened_pr"] = True
                            return metadata_payload
        model_card.push_to_hub(
            repo_id,
            token=token,
            repo_type="model",
            create_pr=True,
            commit_message="Librarian Bot: Add base_model information to model",
            commit_description=template,
        )
        metadata_payload["opened_pr"] = True
        return metadata_payload
    except HfHubHTTPError:
        return metadata_payload


def open_prs(profile: gr.OAuthProfile | None, user_or_org: str = None):
    if not profile:
        return "Please login to open PR requests"
    username = profile.preferred_username
    user_to_receive_prs = user_or_org or username
    data = get_data_for_user(user_to_receive_prs)
    if user_or_org:
        random.sample(data, min(5, len(data)))
    if not data:
        return "No PRs to open"
    results = []
    for metadata_payload in data:
        try:
            results.append(
                update_metadata(metadata_payload, user_making_request=username)
            )

        except Exception as e:
            logger.error(e)
    return f"Opened {len([r for r in results if r['opened_pr']])} PRs"


# description_text = """
# ## Welcome to the Librarian Bot Metadata Request Service

# ⭐ The Librarian Bot Metadata Request Service allows you to request metadata updates for your models on the Hugging Face Hub. ⭐

# Currently this app allows you to request for librarian bot to add metadata for the `base_model` field, situated in the `YAML` block of your model's `README.md`.

# This app will allow you to request metadata for all your models or for another user or org. If you request metadata for another user or org, librarian bot will randomly select 5 models to request metadata for.


# ### How does librarian bot know what metadata to add to your model card? 

# Librarian bot will perform a regular expression match on your `README.md` file to determine whether your model may have bene fine-tuned from another model. This model is known as the `base_model`.

# ### Why add this info to Model Cards?

# Enhancing your model's metadata in this way:
# - 🚀 **Boosts Discoverability** - It becomes straightforward to trace the relationships between various models on the Hugging Face Hub.
# - 🏆**Highlights Impact** - It showcases the contributions and influences different models have within the community.

# For a hands-on example of how such metadata can play a pivotal role in mapping model connections, take a look at [librarian-bots/base_model_explorer](https://huggingface.co/spaces/librarian-bots/base_model_explorer).

# """

description_text = """
## Enhance Your Model's Metadata with Librarian Bot! 

Welcome to the Librarian Bot Metadata Request Service. With a few clicks, enrich your Hugging Face models with key metadata!

🎯 **Purpose of this App**
- Request metadata updates for your models on the Hugging Face Hub, specifically to add or update the `base_model` field in the `YAML` section of your model's `README.md`.
- Optionally, request metadata for models belonging to another user or organization. If doing so, the bot will randomly pick 5 models for metadata addition.

🤖 **How Does Librarian Bot Determine Metadata?**
- It scans your `README.md` to try to determine if your model has been fine-tuned from another model. This original model is identified as the `base_model`.

🚀 **Benefits of Metadata Enhancement**
- **Boosts Discoverability**: Easier tracing of relationships between Hugging Face Hub models.
- **Highlights Impact**: Demonstrates the influence and contribution of different models.

💡 Explore the [librarian-bots/base_model_explorer](https://huggingface.co/spaces/librarian-bots/base_model_explorer) for a hands-on look at the significance of this `base_model` metadata.

**Note**: This app is currently in beta. If you encounter any issues, please [add to this discussion](https://huggingface.co/spaces/librarian-bots/metadata_request_service)

"""


with gr.Blocks() as demo:
    gr.HTML(
        "<h1 style='text-align:center;'><span>&#129302;</span> Librarian Bot Metadata Request Service <span>&#129302;</span></h1>"
    )
    gr.Markdown(
        """<div style='text-align:center;'><img src='https://huggingface.co/spaces/davanstrien/librarian_bot_request_metadata/resolve/main/image.png' style='display:block;margin-left:auto;margin-right:auto;width:150px;'></div><p>"""
    )
    gr.Markdown(description_text)

    with gr.Row():
        gr.Markdown(
            """
        ## How to Use the Librarian Bot Metadata Request Service

        1. **Login to Hugging Face**: Use the login button below to sign in. If you don't have an account, [create one here](https://huggingface.co/join).
        2. **Specify Target User/Organization**: Enter a username or organization name if you wish the Librarian Bot to search metadata for someone other than yourself. Leaving this blank will prompt the bot to look for metadata for your own models and make PRs when a match is found.
        3. **Initiate Metadata Enhancement**: Click the "Open Pull Requests" button. The bot will then search for `base_model` metadata and create Pull Requests for models lacking this information."""
        )
    with gr.Row():
        gr.LoginButton()
        gr.LogoutButton()
        user = gr.Textbox(
            value=None, label="(Optional) user or org to open pull requests for"
        )
    button = gr.Button(value="Open Pull Requests")
    results = gr.Markdown()
    button.click(open_prs, [user], results)


demo.queue(concurrency_count=1).launch()