src/submission/check_validity.py

import json
import logging
import os
from collections import defaultdict

import huggingface_hub
from huggingface_hub import ModelCard
from huggingface_hub.hf_api import ModelInfo
from transformers import AutoConfig
from transformers.models.auto.tokenization_auto import AutoTokenizer

logger = logging.getLogger(__name__)


def check_model_card(repo_id: str) -> tuple[bool, str]:
    """Checks if the model card and license exist and have been filled"""
    logger.debug(f"Checking model card for {repo_id}")
    try:
        card = ModelCard.load(repo_id)
    except huggingface_hub.utils.EntryNotFoundError:
        logger.error(f"Model card not found for {repo_id}")
        return False, "Please add a model card to your model to explain how you trained/fine-tuned it."

    # Enforce license metadata
    if card.data.license is None:
        if not ("license_name" in card.data and "license_link" in card.data):
            return False, (
                "License not found. Please add a license to your model card using the `license` metadata or a"
                " `license_name`/`license_link` pair."
            )

    # Enforce card content
    if len(card.text) < 200:
        logger.error(f"Model card is too short for {repo_id}")
        return False, "Please add a description to your model card, it is too short."

    return True, ""


def is_model_on_hub(
    model_name: str, revision: str, token: str = None, trust_remote_code=False, test_tokenizer=False
) -> tuple[bool, str]:
    """Checks if the model model_name is on the hub, and whether it (and its tokenizer) can be loaded with AutoClasses."""
    logger.debug(f"Checking if model {model_name} is on the hub with revision {revision}")
    try:
        config = AutoConfig.from_pretrained(
            model_name, revision=revision, trust_remote_code=trust_remote_code, token=token
        )
        if test_tokenizer:
            try:
                AutoTokenizer.from_pretrained(
                    model_name, revision=revision, trust_remote_code=trust_remote_code, token=token
                )
            except ValueError as e:
                return (False, f"uses a tokenizer which is not in a transformers release: {e}", None)
            except Exception as e:
                logger.error(f"Error loading tokenizer for {model_name}: {e}")
                return (
                    False,
                    "'s tokenizer cannot be loaded. Is your tokenizer class in a stable transformers release, and correctly configured?",
                    None,
                )
            # Check safetensors format for non-GGUF models
        safetensors_check, safetensors_msg = check_safetensors_format(model_name, revision, token)
        if not safetensors_check:
            return False, safetensors_msg, None

        return True, None, config

    except ValueError:
        return (
            False,
            "needs to be launched with `trust_remote_code=True`. For safety reason, we do not allow these models to be automatically submitted to the leaderboard.",
            None,
        )

    except Exception as e:
        return False, f"was not found on hub: {str(e)}", None
        return False, "was not found on hub!", None


def get_model_size(model_info: ModelInfo, precision: str):
    """Gets the model size from the configuration, or the model name if the configuration does not contain the information."""
    logger.debug(f"Getting model size for {model_info.modelId} with precision {precision}")
    try:
        model_size = round(model_info.safetensors["total"] / 1e9, 3)
    except (AttributeError, TypeError):
        logger.error(f"Error getting model size for {model_info.modelId} with precision {precision}")
        return 0  # Unknown model sizes are indicated as 0, see NUMERIC_INTERVALS in app.py

    size_factor = 8 if (precision == "GPTQ" or "gptq" in model_info.modelId.lower()) else 1
    model_size = size_factor * model_size
    return model_size


def get_model_arch(model_info: ModelInfo):
    """Gets the model architecture from the configuration"""
    logger.debug(f"Getting model architecture for {model_info.modelId}")
    return model_info.config.get("architectures", "Unknown")


def already_submitted_models(requested_models_dir: str) -> set[str]:
    """Gather a list of already submitted models to avoid duplicates"""
    logger.debug(f"Getting already submitted models from {requested_models_dir}")
    depth = 1
    file_names = []
    users_to_submission_dates = defaultdict(list)

    for root, _, files in os.walk(requested_models_dir):
        current_depth = root.count(os.sep) - requested_models_dir.count(os.sep)
        if current_depth == depth:
            for file in files:
                if not file.endswith(".json"):
                    continue
                with open(os.path.join(root, file), "r") as f:
                    info = json.load(f)
                    file_names.append(f"{info['model']}_{info['revision']}_{info['precision']}")

                    # Select organisation
                    if info["model"].count("/") == 0 or "submitted_time" not in info:
                        continue
                    organisation, _ = info["model"].split("/")
                    users_to_submission_dates[organisation].append(info["submitted_time"])

    logger.debug(
        f"Returning already submitted models: {set(file_names)} and users to submission dates: {users_to_submission_dates}"
    )
    return set(file_names), users_to_submission_dates


def check_safetensors_format(model_name: str, revision: str, token: str = None) -> tuple[bool, str]:
    """Checks if the model uses safetensors format"""
    logger.debug(f"Checking safetensors format for {model_name} with revision {revision}")
    try:
        # Use HF API to list repository files
        api = huggingface_hub.HfApi()
        files = api.list_repo_files(model_name, revision=revision, token=token)

        # Check for any .safetensors files in the repository
        if any(f.endswith(".safetensors") for f in files):
            logger.debug(f"Model {model_name} with revision {revision} uses safetensors format")
            return True, ""

        logger.error(f"Model {model_name} with revision {revision} does not use safetensors format")
        return False, (
            "Model weights must be in safetensors format. Please convert your model using: \n"
            "```python\n"
            "from transformers import AutoModelForCausalLM\n"
            "from safetensors.torch import save_file\n\n"
            "model = AutoModelForCausalLM.from_pretrained('your-model')\n"
            "state_dict = model.state_dict()\n"
            "save_file(state_dict, 'model.safetensors')\n"
            "```"
        )
    except Exception as e:
        logger.error(f"Error checking safetensors format: {str(e)}")
        return False, f"Error checking safetensors format: {str(e)}"