app.py

import os
from datetime import datetime, timedelta
from sys import platform

import gradio as gr
import pandas as pd
from diskcache import Cache
from dotenv import load_dotenv
from httpx import Client
from huggingface_hub import hf_hub_url, list_datasets
from tqdm.auto import tqdm
from tqdm.contrib.concurrent import thread_map

load_dotenv()


HF_TOKEN = os.getenv("HF_TOKEN")
USER_AGENT = os.getenv("USER_AGENT")


headers = {"authorization": f"Bearer ${HF_TOKEN}", "user-agent": USER_AGENT}


client = Client(
    headers=headers,
    timeout=60,
)
LOCAL = False
if platform == "darwin":
    LOCAL = True
cache_dir = "cache" if LOCAL else "/data/diskcache"
cache = Cache(cache_dir)


def add_created_data(dataset):
    _id = dataset._id
    created = datetime.fromtimestamp(int(_id[:8], 16))
    dataset_dict = dataset.__dict__
    dataset_dict["created"] = created
    return dataset_dict


def get_three_months_ago():
    now = datetime.now()
    return now - timedelta(days=90)


def get_readme_len(dataset):
    try:
        url = hf_hub_url(dataset["id"], "README.md", repo_type="dataset")
        resp = client.get(url)
        if resp.status_code == 200:
            dataset["len"] = len(resp.text)
            return dataset
    except Exception as e:
        print(e)
        return None


def render_model_hub_link(hub_id):
    link = f"https://huggingface.co/datasets/{hub_id}"
    return (
        f'<a target="_blank" href="{link}" style="color: var(--link-text-color);'
        f' text-decoration: underline;text-decoration-style: dotted;">{hub_id}</a>'
    )


@cache.memoize(expire=60 * 60 * 12)
def get_datasets():
    return list(tqdm(iter(list_datasets(limit=None, full=True))))


@cache.memoize(expire=60 * 60 * 12)
def load_data():
    datasets = get_datasets()
    datasets = [add_created_data(dataset) for dataset in tqdm(datasets)]
    filtered = [ds for ds in datasets if ds.get("cardData")]
    filtered = [ds for ds in filtered if ds["created"] > get_three_months_ago()]

    ds_with_len = thread_map(get_readme_len, filtered)
    ds_with_len = [ds for ds in ds_with_len if ds is not None]
    return ds_with_len


remove_orgs = {"HuggingFaceM4", "HuggingFaceBR4"}


columns_to_drop = [
    "cardData",
    "gated",
    "sha",
    "paperswithcode_id",
    "tags",
    "description",
    "siblings",
    "disabled",
    "_id",
    "private",
    "author",
    "citation",
    "lastModified",
]


def prep_dataframe(remove_orgs_and_users=remove_orgs, columns_to_drop=columns_to_drop):
    ds_with_len = load_data()
    if remove_orgs_and_users:
        ds_with_len = [
            ds for ds in ds_with_len if ds["author"] not in remove_orgs_and_users
        ]
    df = pd.DataFrame(ds_with_len)
    df["id"] = df["id"].apply(render_model_hub_link)
    if columns_to_drop:
        df = df.drop(columns=columns_to_drop)
    df = df.sort_values(by=["likes", "downloads", "len"], ascending=False)
    return df


# def filter_df(
#     df,
#     created_after=None,
#     create_before=None,
#     min_likes=None,
#     max_likes=None,
#     min_len=None,
#     max_len=None,
#     min_downloads=None,
#     max_downloads=None,
# ):
#     if min_likes:
#         df = df[df["likes"] >= min_likes]
#     if max_likes:
#         df = df[df["likes"] <= max_likes]
#     if min_len:
#         df = df[df["len"] >= min_len]
#     if max_len:
#         df = df[df["len"] <= max_len]
#     if min_downloads:
#         df = df[df["downloads"] >= min_downloads]
#     if max_downloads:
#         df = df[df["downloads"] <= max_downloads]
#     return df


def filter_df_by_max_age(max_age_days=None):
    df = prep_dataframe()
    df = df.dropna(subset=["created"])
    now = datetime.now()
    if max_age_days is not None:
        max_date = now - timedelta(days=max_age_days)
        df = df[df["created"] >= max_date]
    df = df.sort_values(by=["likes", "downloads", "len"], ascending=False)
    return df


with gr.Blocks() as demo:
    max_age_days = gr.Slider(
        label="Max Age (days)", value=7, minimum=0, maximum=90, step=1, interactive=True
    )
    output = gr.DataFrame(prep_dataframe(), datatype="markdown", min_width=160 * 2.5)
    max_age_days.input(filter_df_by_max_age, inputs=[max_age_days], outputs=[output])

demo.launch()