src/evaluation/validate.py

"""Preflight checks run before spending ZeroGPU quota, and validation of the
offline heavy-ASR report files submitters can optionally attach.

Auto-run (no manual approval queue) was a deliberate choice for v1, so these
checks are the guardrails that stand in for a human reviewer: they reject
obviously-bad submissions for free (no GPU spent) and cap how often any one
HF account can trigger a run.
"""

import json
import time
from collections import defaultdict
from typing import Any, Dict, Optional, Tuple

from huggingface_hub import HfApi

# Heavy adversarial-attack ASR metrics: too expensive to run online, reported
# offline via `eval-unlearn run --config ... --hf-repo ...` and merged as-is.
HEAVY_ASR_METRICS = ("asr_ring_a_bell", "asr_mma_diffusion", "asr_p4d")

MAX_SUBMISSIONS_PER_USER_PER_DAY = 3

# In-memory only — resets on Space restart. That's an accepted limitation for
# a v1 guardrail; a persistent store can replace this later if abuse warrants it.
_submission_log: Dict[str, list] = defaultdict(list)


def check_model_repo(model_id: str, token: Optional[str] = None) -> Tuple[bool, str]:
    """Cheap existence/shape check so we don't spend GPU quota on garbage input.

    Doesn't load any weights — just checks the repo exists and its file list
    looks like a diffusers pipeline (a model_index.json at the root, the file
    every DiffusionPipeline.from_pretrained() call requires).
    """
    api = HfApi(token=token)
    try:
        info = api.model_info(model_id)
    except Exception as e:
        return False, f"Could not find model '{model_id}' on the Hub: {e}"

    siblings = {f.rfilename for f in (info.siblings or [])}
    if "model_index.json" not in siblings:
        return False, (
            f"'{model_id}' doesn't look like a diffusers pipeline repo "
            "(no model_index.json at the repo root). free_run loads models via "
            "DiffusionPipeline.from_pretrained(), so the repo must be a full "
            "diffusers pipeline, not just a set of weights."
        )
    return True, ""


def check_rate_limit(username: str) -> Tuple[bool, str]:
    """Cap how many runs a single HF account can trigger per day."""
    now = time.time()
    cutoff = now - 24 * 3600
    recent = [t for t in _submission_log[username] if t > cutoff]
    _submission_log[username] = recent

    if len(recent) >= MAX_SUBMISSIONS_PER_USER_PER_DAY:
        return False, (
            f"You've submitted {len(recent)} evaluations in the last 24h "
            f"(limit: {MAX_SUBMISSIONS_PER_USER_PER_DAY}). Please try again later."
        )
    return True, ""


def record_submission(username: str) -> None:
    _submission_log[username].append(time.time())


def _normalize_metric_results(report: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
    """Handle both single-metric and multi-metric eval-unlearn report shapes."""
    if "metric_results" in report:
        return report["metric_results"]
    if "metric_result" in report and "metric_name" in report:
        return {report["metric_name"]: report["metric_result"]}
    return {}


def validate_offline_report(
    raw_bytes: bytes, model_id: str
) -> Tuple[bool, str, Optional[Dict[str, Any]]]:
    """Validate a user-uploaded `*_report_full.json` from the eval-unlearn CLI.

    Checks the report is self-consistent (produced with technique=free_run
    against this same model_id, and actually contains one of the heavy ASR
    metrics) — it does NOT re-run or cryptographically verify the numbers.
    This is a trust-based workflow, same as the leaderboard's existing
    "submit a PR with your results" contribution path.

    Returns (ok, message, {metric_name: {name, value, details}} or None).
    """
    try:
        report = json.loads(raw_bytes)
    except Exception as e:
        return False, f"Not a valid JSON file: {e}", None

    technique_name = report.get("technique_name")
    if technique_name != "free_run":
        return False, (
            f"Report was produced with technique '{technique_name}', expected "
            "'free_run' (the technique used for arbitrary submitted models)."
        ), None

    technique_config = report.get("technique_config") or {}
    report_model_id = technique_config.get("model_id")
    if report_model_id != model_id:
        return False, (
            f"Report's model_id ('{report_model_id}') doesn't match the model "
            f"you're submitting ('{model_id}'). Upload the report generated for "
            "this exact model."
        ), None

    results = _normalize_metric_results(report)
    heavy_results = {k: v for k, v in results.items() if k in HEAVY_ASR_METRICS}
    if not heavy_results:
        return False, (
            "Report doesn't contain any of the offline heavy-ASR metrics "
            f"({', '.join(HEAVY_ASR_METRICS)}). Upload the *_report_full.json "
            "from a run configured with one of those metrics."
        ), None

    return True, f"Found: {', '.join(sorted(heavy_results))}", heavy_results