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 from huggingface_hub import get_discussion_details 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 was requested via the [Librarian Bot](https://huggingface.co/librarian-bot) [metadata request service](https://huggingface.co/spaces/librarian-bots/metadata_request_service) by request of [{opened_by}](https://huggingface.co/{opened_by}) """ PR_FROM_COMMIT_PATTERN = re.compile(r"pr%2F(\d{1,3})/README.md") def get_pr_url_from_commit_url(commit_url, repo_id): re_match = re.search(PR_FROM_COMMIT_PATTERN, commit_url) pr_number = int(re_match.groups()[0]) return get_discussion_details(repo_id=repo_id, discussion_num=pr_number).url 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 commit_url = 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 metadata_payload["pr_url"] = get_pr_url_from_commit_url( commit_url=commit_url, repo_id=repo_id ) 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 is not None: data = 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) if not results: return "No PRs to open" if not any(r["opened_pr"] for r in results): return "No PRs to open" message = "# ✨ Librarian Bot Metadata Request Summary ✨ \n\n" message += ( f"Librarian bot has {len([r for r in results if r['opened_pr']])} PRs open" " against your repos \n\n" ) message += "# URLs for newly opened PRs\n" for result in results: if result["opened_pr"]: print(result) try: message += f"- {result['pr_url']}\n" except KeyError: continue return message # 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. **Note**: The is currently in beta. If you encounter any issues, please [add to this discussion](https://huggingface.co/spaces/librarian-bots/metadata_request_service/discussions/1)
🤖 **How Does Librarian Bot Determine Metadata?** - It scans the `README.md` of the model to check 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.
💡 **See an Example of base_model Metadata in Action** 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). """ with gr.Blocks() as demo: gr.HTML( "

🤖 Librarian Bot Metadata" " Request Service 🤖

" ) gr.Markdown( """

""" ) 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. **Note**: If you specify a target user/organization, the bot will randomly select 5 models to request metadata for. If you do not specify a target user/organization, the bot will try and find `base_model` metadata for all your models.""" ) 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()