src/backend/sort_queue.py

"""Sort evaluation queue by priority."""

from typing import List
from huggingface_hub import HfApi

from src.backend.manage_requests import EvalRequest


def sort_models_by_priority(api: HfApi, models: List[EvalRequest]) -> List[EvalRequest]:
    """
    Sort models by priority for evaluation.

    Current strategy: FIFO (first in, first out) based on submission time.
    Can be extended to prioritize by model popularity, size, etc.

    Args:
        api: HuggingFace API client
        models: List of evaluation requests

    Returns:
        Sorted list of evaluation requests
    """
    # Sort by submission time (oldest first)
    return sorted(models, key=lambda x: x.submitted_time)


def sort_models_by_likes(api: HfApi, models: List[EvalRequest]) -> List[EvalRequest]:
    """
    Sort models by Hub likes (most popular first).

    Args:
        api: HuggingFace API client
        models: List of evaluation requests

    Returns:
        Sorted list of evaluation requests
    """
    return sorted(models, key=lambda x: x.likes, reverse=True)


def sort_models_by_size(models: List[EvalRequest], ascending: bool = True) -> List[EvalRequest]:
    """
    Sort models by parameter count.

    Args:
        models: List of evaluation requests
        ascending: If True, smallest models first

    Returns:
        Sorted list of evaluation requests
    """
    return sorted(models, key=lambda x: x.params, reverse=not ascending)