src/tools/collections.py

import pandas as pd
from huggingface_hub import add_collection_item, delete_collection_item, get_collection, update_collection_item
from huggingface_hub.utils._errors import HfHubHTTPError
from pandas import DataFrame

from src.display.utils import AutoEvalColumn, ModelType
from src.envs import H4_TOKEN, PATH_TO_COLLECTION

# Specific intervals for the collections
intervals = {
    "1B": pd.Interval(0, 1.5, closed="right"),
    "3B": pd.Interval(2.5, 3.5, closed="neither"),
    "7B": pd.Interval(6, 8, closed="neither"),
    "13B": pd.Interval(10, 14, closed="neither"),
    "30B": pd.Interval(25, 35, closed="neither"),
    "65B": pd.Interval(60, 70, closed="neither"),
}


def _filter_by_type_and_size(df, model_type, size_interval):
    """Filter DataFrame by model type and parameter size interval."""
    type_emoji = model_type.value.symbol[0]
    filtered_df = df[df[AutoEvalColumn.model_type_symbol.name] == type_emoji]
    params_column = pd.to_numeric(df[AutoEvalColumn.params.name], errors="coerce")
    mask = params_column.apply(lambda x: x in size_interval)
    return filtered_df.loc[mask]


def _add_models_to_collection(collection, models, model_type, size):
    """Add best models to the collection and update positions."""
    cur_len_collection = len(collection.items)
    for ix, model in enumerate(models, start=1):
        try:
            collection = add_collection_item(
                PATH_TO_COLLECTION,
                item_id=model,
                item_type="model",
                exists_ok=True,
                note=f"Best {model_type.to_str(' ')} model of around {size} on the leaderboard today!",
                token=H4_TOKEN,
            )
            # Ensure position is correct if item was added
            if len(collection.items) > cur_len_collection:
                item_object_id = collection.items[-1].item_object_id
                update_collection_item(collection_slug=PATH_TO_COLLECTION, item_object_id=item_object_id, position=ix)
                cur_len_collection = len(collection.items)
            break  # assuming we only add the top model
        except HfHubHTTPError:
            continue


def update_collections(df: DataFrame):
    """Update collections by filtering and adding the best models."""
    collection = get_collection(collection_slug=PATH_TO_COLLECTION, token=H4_TOKEN)
    cur_best_models = []

    for model_type in ModelType:
        if not model_type.value.name:
            continue
        for size, interval in intervals.items():
            filtered_df = _filter_by_type_and_size(df, model_type, interval)
            best_models = list(
                filtered_df.sort_values(AutoEvalColumn.average.name, ascending=False)[AutoEvalColumn.dummy.name][:10]
            )
            print(model_type.value.symbol, size, best_models)
            _add_models_to_collection(collection, best_models, model_type, size)
            cur_best_models.extend(best_models)

    # Cleanup
    existing_models = {item.item_id for item in collection.items}
    to_remove = existing_models - set(cur_best_models)
    for item_id in to_remove:
        try:
            delete_collection_item(collection_slug=PATH_TO_COLLECTION, item_object_id=item_id, token=H4_TOKEN)
        except HfHubHTTPError:
            continue