src/submission/submit.py

import json
import os
from datetime import datetime, timezone

from huggingface_hub import snapshot_download
from src.submission.check_validity import get_model_tags
from src.display.formatting import styled_error, styled_message, styled_warning
from src.envs import API, EVAL_REQUESTS_PATH, DYNAMIC_INFO_PATH, DYNAMIC_INFO_FILE_PATH, DYNAMIC_INFO_REPO, TOKEN, QUEUE_REPO, RATE_LIMIT_PERIOD, RATE_LIMIT_QUOTA
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 submit_eval_complete(
    model_name: str,
    revision_commit: str,
    model_api_url: str,
    model_api_key: str,
    online_api_model_name: str,
    runsh_file,
    adapter_file
):
    """
    Complete evaluation submission - integrates all three parts of information
    """
    # Validate model information
    if not model_name or not model_name.strip():
        return styled_error("Please enter model name")
    
    if not revision_commit or not revision_commit.strip():
        revision_commit = "main"
    
    # Validate API information (if provided)
    if model_api_url and model_api_key and online_api_model_name:
        if not model_api_url.startswith(('http://', 'https://')):
            return styled_error("API URL format is incorrect, please start with http:// or https://")
    
    # Validate inference files (if provided)
    if runsh_file and adapter_file:
        max_size = 5 * 1024 * 1024  # 5MB
        if os.path.getsize(runsh_file.name) > max_size:
            return styled_error("run.sh file size cannot exceed 5MB")
        
        if os.path.getsize(adapter_file.name) > max_size:
            return styled_error("model_adapter.py file size cannot exceed 5MB")
    
    # Call the original add_new_eval function
    try:
        result = add_new_eval(
            model=model_name,
            model_api_url=model_api_url or "",
            model_api_key=model_api_key or "",
            model_api_name=online_api_model_name or "",
            base_model="",  # Can be set as needed
            revision=revision_commit,
            precision="float16",  # Default precision
            private="false",
            weight_type="Original",  # Default weight type
            model_type="",  # Can be set as needed
            runsh=runsh_file,
            adapter=adapter_file
        )
        return result
    except Exception as e:
        return styled_error(f"Submission failed: {str(e)}")

def add_new_eval(
    model: str,
    model_api_url: str,
    model_api_key: str,
    model_api_name: str,
    base_model: str,
    revision: str,
    precision: str,
    private: str,
    weight_type: str,
    model_type: str,
    runsh,
    adapter
):
    global REQUESTED_MODELS
    global USERS_TO_SUBMISSION_DATES
    if not REQUESTED_MODELS:
        REQUESTED_MODELS, USERS_TO_SUBMISSION_DATES = already_submitted_models(EVAL_REQUESTS_PATH)

    user_name = ""
    model_path = model
    if "/" in model:
        user_name = model.split("/")[0]
        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: 
        model_type = ""
        #return styled_error("Please select a model type.")

    # Does the model actually exist?
    if revision == "":
        revision = "main"

    architecture = "?"
    downloads = 0
    created_at = ""
    # Is the model on the hub?
    if len(model_api_url)==0:
        # Is the model info correctly filled?
        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)

        modelcard_OK, error_msg = check_model_card(model)
        if not modelcard_OK:
            return styled_error(error_msg)
        tags = []

        likes = model_info.likes
    else:
        model_size = 0
        license = ""
        likes = 0
        tags = []
        downloads = 0

    # Seems good, creating the eval
    print("Adding new eval", runsh)
    max_size = 5 * 1024 * 1024  # 5MB
    if (runsh is not None) and (adapter is not None):
        if os.path.getsize(runsh.name) > max_size:
            return "Error: File size cannot exceed 5MB!"
        if os.path.getsize(adapter.name) > max_size:
            return "Error: File size cannot exceed 5MB!"
        with open(runsh.name, "r") as f:
            runsh = f.read()
        with open(adapter.name, "r") as f:
            adapter = f.read()
    else:
        runsh = ""
        adapter = ""
    eval_entry = {
        "model": model,
        "model_api_url": model_api_url,
        "model_api_key": model_api_key,
        "model_api_name": model_api_name,
        "base_model": base_model,
        "revision": revision,
        "precision": precision,
        "private": private,
        "weight_type": weight_type,
        "status": "PENDING",
        "submitted_time": current_time,
        "model_type": model_type,
        "params": model_size,
        "private": False,
        "runsh": runsh,
        "adapter": adapter,
    }

    supplementary_info = {
        "likes": 0,
        "license": '',
        "still_on_hub": True,
        "tags": tags,
        "downloads": downloads,
        "created_at": created_at
    }

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

    print("Creating eval file")
    OUT_DIR = f"{EVAL_REQUESTS_PATH}/{user_name}"
    os.makedirs(OUT_DIR, exist_ok=True)
    out_path = f"{OUT_DIR}/{model_path}_eval_request_False_{precision}_{weight_type}.json"

    with open(out_path, "w") as f:
        f.write(json.dumps(eval_entry))

    print("Uploading eval file")
    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",
    )
  
    # We want to grab the latest version of the submission file to not accidentally overwrite it
    snapshot_download(
        repo_id=DYNAMIC_INFO_REPO, local_dir=DYNAMIC_INFO_PATH, repo_type="dataset", tqdm_class=None, etag_timeout=30
    )

    with open(DYNAMIC_INFO_FILE_PATH) as f:
        all_supplementary_info = json.load(f)

    all_supplementary_info[model] = supplementary_info
    with open(DYNAMIC_INFO_FILE_PATH, "w") as f:
        json.dump(all_supplementary_info, f, indent=2)

    API.upload_file(
        path_or_fileobj=DYNAMIC_INFO_FILE_PATH,
        path_in_repo=DYNAMIC_INFO_FILE_PATH.split("/")[-1],
        repo_id=DYNAMIC_INFO_REPO,
        repo_type="dataset",
        commit_message=f"Add {model} to dynamic info queue",
    )

    # Remove the local file
    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."
    )