added collections back to main

app.py

@@ -3,7 +3,7 @@ import logging
 import time
 import gradio as gr
 import datasets
-from apscheduler.schedulers.background import BackgroundScheduler
+from src.tools.collections import update_collections
 from huggingface_hub import snapshot_download
 from gradio_leaderboard import Leaderboard, ColumnFilter, SelectColumns
 from gradio_space_ci import enable_space_ci
@@ -105,6 +105,8 @@ def init_space(full_init: bool = True):
         cols=COLS,
         benchmark_cols=BENCHMARK_COLS,
     )
+    if full_init:
+        update_collections(leaderboard_df)
 
     # Evaluation queue DataFrame retrieval is independent of initialization detail level
     eval_queue_dfs = get_evaluation_queue_df(EVAL_REQUESTS_PATH, EVAL_COLS)

src/tools/collections.py

@@ -0,0 +1,76 @@
+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 HF_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=HF_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=HF_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.fullname.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=HF_TOKEN)
+        except HfHubHTTPError:
+            continue