src/display/utils.py

import os
import re
import json
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Any

import pandas as pd
from huggingface_hub import HfApi, hf_hub_download, list_repo_files


HF_TOKEN = os.environ.get("HF_TOKEN")
SUBMISSIONS_REPO = os.environ.get("SUBMISSIONS_REPO", "your-org/CULTURE-MT-submissions")
RESULTS_REPO = os.environ.get("RESULTS_REPO", "your-org/CULTURE-MT-results")

API = HfApi(token=HF_TOKEN)


RESULT_COLUMNS = [
    "Rank",
    "Model",
    "Organization",
    "Base Model",
    "Method",
    "Avg Score",
    "Effective Rate",
    "Ineffective Rate",
    "Cultural",
    "Fluency",
    "Adequacy",
    "Submission Time",
]


def _safe_name(text: str) -> str:
    text = text.strip().replace(" ", "_")
    text = re.sub(r"[^a-zA-Z0-9_\-.]", "", text)
    return text[:80] or "anonymous_model"


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _read_jsonl(path: str) -> List[Dict[str, Any]]:
    rows = []
    seen = set()

    with open(path, "r", encoding="utf-8") as f:
        for line_no, line in enumerate(f, start=1):
            line = line.strip()
            if not line:
                continue

            try:
                item = json.loads(line)
            except json.JSONDecodeError as e:
                raise ValueError(f"Line {line_no}: invalid JSON: {e}")

            if "id" not in item:
                raise ValueError(f"Line {line_no}: missing `id`.")

            if "translation" not in item:
                raise ValueError(f"Line {line_no}: missing `translation`.")

            sid = str(item["id"])
            if sid in seen:
                raise ValueError(f"Duplicate id: {sid}")

            if not isinstance(item["translation"], str) or not item["translation"].strip():
                raise ValueError(f"Line {line_no}: empty `translation`.")

            seen.add(sid)
            rows.append(
                {
                    "id": sid,
                    "translation": item["translation"],
                }
            )

    if not rows:
        raise ValueError("submission.jsonl is empty.")

    return rows


def submit_prediction(
    model_name: str,
    organization: str,
    base_model: str,
    method: str,
    description: str,
    predictions_file,
) -> str:
    if not HF_TOKEN:
        return "❌ `HF_TOKEN` is not configured in Space secrets."

    if predictions_file is None:
        return "❌ Please upload `submission.jsonl`."

    if not model_name.strip():
        return "❌ Please provide a model name."

    try:
        predictions = _read_jsonl(predictions_file.name)

        timestamp = _now_iso()
        submit_id = f"{_safe_name(model_name)}_{int(time.time())}"

        predictions_path = f"predictions/{submit_id}.jsonl"
        metadata_path = f"metadata/{submit_id}.json"

        normalized_jsonl = "\n".join(
            json.dumps(x, ensure_ascii=False) for x in predictions
        ).encode("utf-8")

        metadata = {
            "submission_id": submit_id,
            "model_name": model_name.strip(),
            "organization": organization.strip(),
            "base_model": base_model.strip(),
            "method": method.strip(),
            "description": description.strip(),
            "submission_time": timestamp,
            "num_predictions": len(predictions),
            "predictions_file": predictions_path,
            "status": "pending",
        }

        API.upload_file(
            path_or_fileobj=normalized_jsonl,
            path_in_repo=predictions_path,
            repo_id=SUBMISSIONS_REPO,
            repo_type="dataset",
            token=HF_TOKEN,
            commit_message=f"Add predictions for {submit_id}",
        )

        API.upload_file(
            path_or_fileobj=json.dumps(metadata, ensure_ascii=False, indent=2).encode("utf-8"),
            path_in_repo=metadata_path,
            repo_id=SUBMISSIONS_REPO,
            repo_type="dataset",
            token=HF_TOKEN,
            commit_message=f"Add metadata for {submit_id}",
        )

        return (
            f"✅ Submission received.\n\n"
            f"**Submission ID:** `{submit_id}`\n\n"
            f"**Status:** pending evaluation.\n\n"
            f"The result will appear on the leaderboard after the private evaluator finishes."
        )

    except Exception as e:
        return f"❌ Submission failed: `{str(e)}`"


def _load_result_file(file_path: str) -> Dict[str, Any]:
    local_path = hf_hub_download(
        repo_id=RESULTS_REPO,
        filename=file_path,
        repo_type="dataset",
        token=HF_TOKEN,
    )

    with open(local_path, "r", encoding="utf-8") as f:
        return json.load(f)


def load_results() -> pd.DataFrame:
    if not HF_TOKEN:
        return pd.DataFrame(columns=RESULT_COLUMNS)

    try:
        files = list_repo_files(
            repo_id=RESULTS_REPO,
            repo_type="dataset",
            token=HF_TOKEN,
        )

        result_files = [
            f for f in files
            if f.startswith("results/") and f.endswith(".json")
        ]

        rows = []

        for file_path in result_files:
            try:
                item = _load_result_file(file_path)
                summary = item.get("summary", item)

                rows.append(
                    {
                        "Model": item.get("model_name", ""),
                        "Organization": item.get("organization", ""),
                        "Base Model": item.get("base_model", ""),
                        "Method": item.get("method", ""),
                        "Avg Score": summary.get("avg_score"),
                        "Effective Rate": summary.get("effective_rate"),
                        "Ineffective Rate": summary.get("ineffective_rate"),
                        "Cultural": summary.get("avg_cultural"),
                        "Fluency": summary.get("avg_fluency"),
                        "Adequacy": summary.get("avg_adequacy"),
                        "Submission Time": item.get("submission_time") or item.get("timestamp", ""),
                    }
                )

            except Exception:
                continue

        if not rows:
            return pd.DataFrame(columns=RESULT_COLUMNS)

        df = pd.DataFrame(rows)

        df = df.sort_values(
            by=["Avg Score", "Effective Rate", "Ineffective Rate"],
            ascending=[False, False, True],
            na_position="last",
        ).reset_index(drop=True)

        df.insert(0, "Rank", range(1, len(df) + 1))

        return df[RESULT_COLUMNS]

    except Exception:
        return pd.DataFrame(columns=RESULT_COLUMNS)