hf_utils.py

import json
import os
import tempfile
import time
import uuid
from datetime import datetime, timezone
from typing import Any, Dict, List

import pandas as pd
from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.utils import HfHubHTTPError

from config import DB_REPO_ID, DB_REPO_TYPE, SUBMISSIONS_TOKEN, DAILY_SUBMISSION_CAP

# =========================
# HF API CLIENT
# =========================

_api = None
def api_client() -> HfApi:
    global _api
    if _api is None:
        _api = HfApi()
    return _api

# =========================
# DATE / CAP HELPERS
# =========================

def _today_utc_str() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%d")

def _get_cap_for_phase(phase_codename: str) -> int:
    """Return the daily submission cap for a given phase."""
    if phase_codename in ("test-challenge2024", "test-challenge2025"):
        return 1
    return DAILY_SUBMISSION_CAP

def _count_submissions_today(username: str, phase_codename: str | None = None) -> int:
    """Count today's submissions for a user, optionally filtered by phase."""
    try:
        files = api_client().list_repo_files(
            repo_id=DB_REPO_ID, repo_type=DB_REPO_TYPE, token=SUBMISSIONS_TOKEN
        )
        today = _today_utc_str()
        count = 0
        for f in files:
            if not (f.startswith("submissions/") and f.endswith("/meta.json")):
                continue
            try:
                p = hf_hub_download(
                    repo_id=DB_REPO_ID, repo_type=DB_REPO_TYPE,
                    filename=f, token=SUBMISSIONS_TOKEN
                )
                meta = json.load(open(p))
                if meta.get("username", "").lower() != username.lower():
                    continue
                if phase_codename and meta.get("phase_codename") != phase_codename:
                    continue
                ts = meta.get("timestamp", 0)
                sub_date = datetime.fromtimestamp(ts, tz=timezone.utc).strftime("%Y-%m-%d")
                if sub_date == today:
                    count += 1
            except Exception:
                continue
        return count
    except Exception:
        return 0

# =========================
# UPLOAD HELPERS
# =========================

def _upload_json(data: Any, path_in_repo: str, commit_message: str = "") -> None:
    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as tmp:
        json.dump(data, tmp, ensure_ascii=False)
        tmp_path = tmp.name
    try:
        api_client().upload_file(
            path_or_fileobj=tmp_path,
            path_in_repo=path_in_repo,
            repo_id=DB_REPO_ID,
            repo_type=DB_REPO_TYPE,
            token=SUBMISSIONS_TOKEN,
            commit_message=commit_message or f"Add {path_in_repo}",
        )
    finally:
        try:
            os.remove(tmp_path)
        except OSError:
            pass

def _create_submission_record(*, pred, team, model_name, phase_codename,
                               challenge_type, original_filename, username, email,
                               subfolder: str = "") -> str:
    """
    Write pred.json / meta.json / status.json to the dataset repo.
    subfolder: e.g. "answer-therapy" → submissions/answer-therapy/<uuid>/
    """
    if not SUBMISSIONS_TOKEN:
        raise ValueError("Missing SUBMISSIONS_TOKEN.")
    submission_id = str(uuid.uuid4())
    ts = int(time.time())
    meta = {
        "submission_id": submission_id,
        "team":              team.strip(),
        "model":             model_name.strip(),
        "phase_codename":    phase_codename,
        "challenge_type":    challenge_type,
        "timestamp":         ts,
        "original_filename": original_filename,
        "username":          username,
        "email":             email,
    }
    status = {"state": "queued", "timestamp": ts}
    prefix = f"submissions/{subfolder}/{submission_id}" if subfolder else f"submissions/{submission_id}"
    _upload_json(pred,   f"{prefix}/pred.json",   f"pred {submission_id}")
    _upload_json(meta,   f"{prefix}/meta.json",   f"meta {submission_id}")
    _upload_json(status, f"{prefix}/status.json", f"status {submission_id}")
    return submission_id

# =========================
# READ HELPERS
# =========================

def _load_user_submissions(username: str, subfolder: str = "") -> List[Dict]:
    """Load all submissions for a user from a given subfolder (or root)."""
    if not SUBMISSIONS_TOKEN:
        return []
    try:
        files = api_client().list_repo_files(
            repo_id=DB_REPO_ID, repo_type=DB_REPO_TYPE, token=SUBMISSIONS_TOKEN
        )
    except Exception:
        return []

    prefix = f"submissions/{subfolder}/" if subfolder else "submissions/"
    results = []
    for f in files:
        if not (f.startswith(prefix) and f.endswith("/meta.json")):
            continue
        try:
            parts = f.split("/")
            sid = parts[2] if subfolder else parts[1]
            meta_path = hf_hub_download(
                repo_id=DB_REPO_ID, repo_type=DB_REPO_TYPE,
                filename=f, token=SUBMISSIONS_TOKEN
            )
            meta = json.load(open(meta_path))
            if meta.get("username", "").lower() != username.lower():
                continue
            status_file = f"{prefix}{sid}/status.json"
            try:
                status_path = hf_hub_download(
                    repo_id=DB_REPO_ID, repo_type=DB_REPO_TYPE,
                    filename=status_file, token=SUBMISSIONS_TOKEN
                )
                status = json.load(open(status_path))
            except Exception:
                status = {"state": "unknown"}
            metrics = status.get("metrics", {}) if status.get("state") == "done" else {}
            error   = status.get("error",   "") if status.get("state") == "failed" else ""
            results.append({
                "submission_id":  sid,
                "team":           meta.get("team", ""),
                "model":          meta.get("model", ""),
                "phase":          meta.get("phase_codename", ""),
                "challenge_type": meta.get("challenge_type", ""),
                "timestamp":      meta.get("timestamp", 0),
                "state":          status.get("state", "unknown"),
                "error":          error[:120] if error else "",
                **metrics,
            })
        except Exception:
            continue

    results.sort(key=lambda x: x["timestamp"], reverse=True)
    return results

def _load_leaderboard_df(leaderboard_file: str, metric_cols: list) -> pd.DataFrame:
    """Generic leaderboard loader for any challenge."""
    empty = pd.DataFrame(columns=["team", "model", "phase_codename", *metric_cols, "timestamp"])
    if not SUBMISSIONS_TOKEN:
        return empty
    try:
        path = hf_hub_download(
            repo_id=DB_REPO_ID, repo_type=DB_REPO_TYPE,
            filename=leaderboard_file, token=SUBMISSIONS_TOKEN
        )
    except HfHubHTTPError as e:
        if "404" in str(e):
            return empty
        raise

    rows = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                rows.append(json.loads(line))
            except json.JSONDecodeError:
                continue

    if not rows:
        return empty

    df = pd.DataFrame(rows)
    for col in ["team", "model", "phase_codename", "timestamp", *metric_cols]:
        if col not in df.columns:
            df[col] = None
    return df