src/submission/submit.py

import json
import os
from datetime import datetime, timezone

from src.display.formatting import styled_error, styled_message, styled_warning
from src.envs import (
    API,
    EVAL_REQUESTS_NAMESPACED_PATH,
    LOCAL_DEV,
    PROJECT_NAMESPACE,
    QUEUE_REPO,
    TOKEN,
)
from src.submission.check_validity import (
    already_submitted_models,
    check_model_card,
    get_model_size,
    is_model_on_hub,
)

REQUESTED_MODELS = None
USERS_TO_SUBMISSION_DATES = None


def add_new_eval(
    model: str,
    base_model: str,
    revision: str,
    precision: str,
    weight_type: str,
    model_type: str,
):
    global REQUESTED_MODELS
    global USERS_TO_SUBMISSION_DATES

    if not REQUESTED_MODELS:
        REQUESTED_MODELS, USERS_TO_SUBMISSION_DATES = already_submitted_models(EVAL_REQUESTS_NAMESPACED_PATH)
    print("[SUBMIT] Loaded submitted-model cache.")

    user_name = ""
    model_path = model
    if "/" in model:
        user_name, model_path = model.split("/", 1)

    precision = precision.split(" ")[0]
    current_time = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")

    if model_type is None or model_type == "":
        return styled_error("Please select a model type.")

    if revision == "":
        revision = "main"

    model_info = None
    model_size = 0
    license_name = "unknown"

    if LOCAL_DEV:
        print("[LOCAL_DEV] Skip hub/model-card validations and upload. Keeping local request only.")
    else:
        if weight_type in ["Delta", "Adapter"]:
            base_model_on_hub, error, _ = is_model_on_hub(
                model_name=base_model,
                revision=revision,
                token=TOKEN,
                test_tokenizer=True,
            )
            if not base_model_on_hub:
                return styled_error(f'Base model "{base_model}" {error}')

        if weight_type != "Adapter":
            model_on_hub, error, _ = is_model_on_hub(
                model_name=model,
                revision=revision,
                token=TOKEN,
                test_tokenizer=True,
            )
            if not model_on_hub:
                return styled_error(f'Model "{model}" {error}')

        try:
            model_info = API.model_info(repo_id=model, revision=revision)
        except Exception:
            return styled_error("Could not get your model information. Please fill it up properly.")

        model_size = get_model_size(model_info=model_info, precision=precision)
        try:
            license_name = model_info.cardData["license"]
        except Exception:
            return styled_error("Please select a license for your model")

        modelcard_ok, error_msg = check_model_card(model)
        if not modelcard_ok:
            return styled_error(error_msg)

    print("[SUBMIT] Creating eval request entry ...")
    eval_entry = {
        "project_id": PROJECT_NAMESPACE,
        "model": model,
        "base_model": base_model,
        "revision": revision,
        "precision": precision,
        "weight_type": weight_type,
        "status": "PENDING",
        "submitted_time": current_time,
        "model_type": model_type,
        "likes": model_info.likes if model_info is not None else 0,
        "params": model_size,
        "license": license_name,
        "private": False,
    }

    if f"{model}_{revision}_{precision}" in REQUESTED_MODELS:
        return styled_warning("This model has been already submitted.")

    out_dir = f"{EVAL_REQUESTS_NAMESPACED_PATH}/{user_name}"
    os.makedirs(out_dir, exist_ok=True)
    out_path = f"{out_dir}/{model_path}_eval_request_False_{precision}_{weight_type}.json"
    print(f"[SUBMIT] Writing request file: {out_path}")
    with open(out_path, "w", encoding="utf-8") as f:
        f.write(json.dumps(eval_entry, ensure_ascii=False, indent=2))

    if LOCAL_DEV:
        return styled_message("LOCAL_DEV mode: request saved locally to eval-queue (no HF upload).")

    print("[SUBMIT] Uploading request to HF dataset queue ...")
    API.upload_file(
        path_or_fileobj=out_path,
        path_in_repo=out_path.split("eval-queue/")[1],
        repo_id=QUEUE_REPO,
        repo_type="dataset",
        commit_message=f"Add {model} to eval queue",
    )

    os.remove(out_path)
    return styled_message(
        "Your request has been submitted to the evaluation queue!\nPlease wait for up to an hour for the model to show in the PENDING list."
    )