Upload 2 files

loop_retrieve_cards.py

@@ -0,0 +1,52 @@
+import threading
+import time
+
+from utils import fetch_models, fetch_datasets, fetch_readme
+from tqdm import tqdm
+
+
+# Define global variables
+models = None
+datasets = None
+readme_dict = {}
+
+
+# Provide an API to get models
+def get_models():
+    return models
+
+
+# Provide an API to get datasets
+def get_datasets():
+    return datasets
+
+
+# Provide an API to get READMEs
+def get_readme_dict():
+    return readme_dict
+
+
+# Start a thread to continuously update cards
+def run():
+    global models, datasets, readme_dict, cnt
+
+    while True:
+        new_models = fetch_models()
+        new_datasets = fetch_datasets()
+
+        # Add READMEs
+        new_readme_dict = {}
+        for model in new_models:
+            new_readme_dict[model] = fetch_readme(model, "model")
+
+        for dataset in new_datasets:
+            new_readme_dict[dataset] = fetch_readme(dataset, "dataset")
+
+        # Update global variables
+        models = new_models
+        datasets = new_datasets
+        readme_dict = new_readme_dict
+
+
+t = threading.Thread(target=run)
+t.start()

utils.py

@@ -0,0 +1,86 @@
+import requests
+import re
+
+
+def fetch_models(author: str = "SaProtHub") -> list:
+    """
+    Retrieve models belonging to a specific author
+
+    Args:
+        author: Author name
+
+    Returns:
+        models: List of models
+    """
+
+    url = f"https://hf-mirror.com/api/models?author={author}"
+    response = requests.get(url)
+    models_dict = response.json()
+    models = [item["id"] for item in models_dict]
+
+    return models
+
+
+def fetch_datasets(author: str = "SaProtHub") -> list:
+    """
+    Retrieve datasets belonging to a specific author
+
+    Args:
+        author: Author name
+
+    Returns:
+        datasets: List of datasets
+    """
+
+    url = f"https://hf-mirror.com/api/datasets?author={author}"
+    response = requests.get(url)
+    datasets_dict = response.json()
+    datasets = [item["id"] for item in datasets_dict]
+
+    return datasets
+
+
+def fetch_readme(card_id: str, card_type: str) -> str:
+    """
+    Retrieve the README file of a model or dataset
+
+    Args:
+        card_id: Model or dataset ID
+        card_type: Type of card, either "model" or "dataset"
+
+    Returns:
+        readme: README text
+    """
+    if card_type == "model":
+        url = f"https://hf-mirror.com/{card_id}/raw/main/README.md"
+    else:
+        url = f"https://hf-mirror.com/datasets/{card_id}/raw/main/README.md"
+
+    response = requests.get(url)
+    readme = response.text.split("---")[-1]
+
+    return readme
+
+
+def set_text_bg_color(pattern: str, text: str, color: str = "yellow") -> str:
+    """
+    Set the background color of a pattern in a text
+
+    Args:
+        pattern: Pattern to highlight
+        text: Text to search
+        color: Background color
+
+    Returns:
+        text: Text with highlighted pattern
+    """
+
+    # Find all matches, ignoring case
+    matches = set(re.findall(pattern, text, flags=re.IGNORECASE))
+
+    # Highlight all matches
+    for match in matches:
+        text = text.replace(match, f'<span style="background-color:{color}">{match}</span>')
+
+    return text
+