evaluator.py

"""FonBench — évaluateur de fond du Space.

Un unique thread démon : il prend une demande dans la file Supabase,
charge le corpus de test PRIVÉ, transcrit par tranches et republie les
scores. Aucun GPU : tout se fait sur le CPU du Space.

Trois précautions valent d'être expliquées, parce qu'elles dictent la
structure du fichier.

1. Le corpus est chargé AVANT le modèle, et le token qui y donne accès est
   retiré de l'environnement dès l'import. Le Space charge des modèles
   écrits par des inconnus : aucun ne doit tourner dans un processus où
   traîne encore de quoi lire le corpus. S'y ajoute `trust_remote_code=False`
   sans exception — c'est ce qui empêche le code personnalisé d'un dépôt de
   s'exécuter.

2. On reprend après redémarrage. Un Space gratuit redémarre souvent et une
   évaluation dure des heures : sans reprise, les gros modèles ne
   finiraient jamais. Ce qu'on sauvegarde à chaque tranche, ce sont des
   COMPTEURS d'erreurs, jamais des transcriptions — sinon on recopierait le
   corpus privé dans la base. La somme des compteurs redonne exactement les
   mêmes scores (voir fonbench_eval).

3. Le Space n'a pas la clé service de la base. Il écrit via quatre
   procédures protégées par un jeton dédié : au pire, un jeton volé permet
   de polluer la file publique, pas de toucher aux comptes ni aux scores.
"""

from __future__ import annotations

import os
import threading
import time
import traceback

import requests

import fonbench_eval as fe

# ZeroGPU : le GPU n'est attaché que le temps d'une fonction décorée. En
# dehors d'un Space ZeroGPU (tests locaux), le décorateur devient neutre.
try:
    import spaces

    GPU = spaces.GPU
    ON_ZEROGPU = True
except Exception:  # noqa: BLE001
    ON_ZEROGPU = False

    def GPU(*args, **kwargs):  # noqa: D103
        def deco(fn):
            return fn
        return deco(args[0]) if args and callable(args[0]) else deco

# --- configuration ------------------------------------------------------

def env(name: str, default: str = "") -> str:
    """Variable d'environnement nettoyée.

    Un secret collé dans l'interface du Space embarque volontiers un retour
    à la ligne, ce qui rend l'en-tête HTTP invalide et fait échouer tous les
    appels — sans indice évident dans les logs. On tranche à la source.
    """
    return (os.environ.get(name) or default).strip()


SUPABASE_URL = env(
    "SUPABASE_URL", "https://cqdimvcnmhrsdcoobkmd.supabase.co"
).rstrip("/")
# Publique par conception : la Row Level Security fait le tri, pas le secret.
ANON_KEY = env(
    "SUPABASE_ANON_KEY", "sb_publishable_MapYll-_Y0hNoLYOfaDR3w_tsTEHyiz"
)
SPACE_TOKEN = env("SPACE_TOKEN")

# Sorti de l'environnement immédiatement : un modèle inconnu chargé plus
# tard ne doit pas pouvoir le lire dans os.environ.
_DATA_TOKEN = (os.environ.pop("HF_DATA_TOKEN", None)
               or os.environ.pop("HF_TOKEN", None) or "").strip() or None

REST = f"{SUPABASE_URL}/rest/v1"
HEADERS = {"apikey": ANON_KEY, "Content-Type": "application/json"}

WORKER_VERSION = "space-1.0"
POLL_SECONDS = 60

# Une tranche = un créneau GPU. ZeroGPU plafonne la durée d'un créneau, et
# plus on en demande, plus on attend : on vise ~10 s de calcul réel.
# Sur A10G, un CTC tourne autour de 100× temps réel (100 énoncés ≈ 700 s
# d'audio ≈ 7 s), un Whisper autour de 5× (20 énoncés ≈ 140 s ≈ 30 s).
GPU_DURATION = 120
CHUNK_CTC = 100
CHUNK_SEQ2SEQ = 20
BATCH_CTC = 8
# Quota journalier épuisé : on attend au lieu d'abandonner. Le point de
# reprise est déjà enregistré, rien n'est perdu.
QUOTA_WAIT_SECONDS = 1800

# Architectures autorégressives : pipeline + découpage à 30 s, sinon les
# énoncés longs font échouer Whisper.
SEQ2SEQ_TYPES = {
    "whisper", "speech_to_text", "speech-encoder-decoder",
    "speech_encoder_decoder", "seamless_m4t", "seamless_m4t_v2",
}

_state = {"message": "démarrage…", "model": None, "done": 0, "total": 0}


def log(msg: str) -> None:
    print(f"[evaluator] {msg}", flush=True)


def status() -> dict:
    return dict(_state)


# --- accès à la file ----------------------------------------------------

def rpc(name: str, payload: dict):
    r = requests.post(
        f"{REST}/rpc/{name}", headers=HEADERS, json=payload, timeout=60
    )
    if r.status_code >= 400:
        raise RuntimeError(f"{name}: {r.status_code} {r.text[:200]}")
    return r.json() if r.text.strip() else None


# --- audio --------------------------------------------------------------

def decode_audio(cell, target_sr: int = 16000):
    """Décode un audio (WAV/FLAC/OGG/Opus/WebM/MP3…) en mono float32 16 kHz.

    PyAV plutôt que soundfile : le corpus contient des conteneurs que
    libsndfile ne sait pas ouvrir (WebM, Opus) et des fichiers dont
    l'en-tête le fait échouer (« array is too big »).
    """
    import io

    import av
    import numpy as np

    raw = cell["bytes"] if isinstance(cell, dict) else cell
    if raw is None and isinstance(cell, dict) and cell.get("path"):
        with open(cell["path"], "rb") as f:
            raw = f.read()

    with av.open(io.BytesIO(raw)) as container:
        stream = container.streams.audio[0]
        resampler = av.audio.resampler.AudioResampler(
            format="flt", layout="mono", rate=target_sr
        )
        chunks: list = []

        def _emit(frame):
            res = resampler.resample(frame)
            for rf in res if isinstance(res, list) else ([res] if res else []):
                chunks.append(rf.to_ndarray().reshape(-1))

        for frame in container.decode(stream):
            _emit(frame)
        _emit(None)  # flush du resampler

    if not chunks:
        return np.zeros(1, dtype="float32")
    return np.concatenate(chunks).astype("float32")


# --- corpus -------------------------------------------------------------

_dataset_cache: dict = {}


def load_corpus(bench: dict):
    """Charge le corpus de test. Gardé en mémoire entre deux évaluations."""
    from datasets import Audio, load_dataset

    key = (bench["hf_dataset_id"], bench.get("hf_revision"), bench["hf_split"])
    if key in _dataset_cache:
        return _dataset_cache[key]

    ds = load_dataset(
        bench["hf_dataset_id"],
        split=bench["hf_split"],
        revision=bench.get("hf_revision") or None,
        token=_DATA_TOKEN,
    )
    if bench.get("sample_size"):
        ds = ds.shuffle(seed=bench.get("sample_seed") or 42)
        ds = ds.select(range(min(bench["sample_size"], len(ds))))
    ds = ds.cast_column("audio", Audio(decode=False))
    _dataset_cache.clear()          # 16 Go de RAM : un seul corpus à la fois
    _dataset_cache[key] = ds
    return ds


# --- modèle -------------------------------------------------------------

class Transcriber:
    """Charge un modèle du Hub et transcrit des tableaux 16 kHz."""

    def __init__(self, model_id: str, revision: str):
        import torch
        from transformers import AutoConfig

        torch.set_num_threads(max(1, (os.cpu_count() or 2)))
        self.torch = torch
        self.model_id = model_id

        cfg = AutoConfig.from_pretrained(
            model_id, revision=revision, trust_remote_code=False
        )
        self.architecture = cfg.model_type
        self.seq2seq = cfg.model_type in SEQ2SEQ_TYPES or any(
            "ConditionalGeneration" in a or "Seq2Seq" in a
            for a in (getattr(cfg, "architectures", None) or [])
        )
        self.chunk_size = CHUNK_SEQ2SEQ if self.seq2seq else CHUNK_CTC
        self.decoder_type = "encoder-decoder" if self.seq2seq else "ctc"

        if self.seq2seq:
            self._load_seq2seq(model_id, revision)
        else:
            self._load_ctc(model_id, revision)

        self.model_params = sum(p.numel() for p in self.model.parameters())

    def _load_seq2seq(self, model_id, revision):
        from transformers import (
            AutoModelForSpeechSeq2Seq,
            AutoProcessor,
            pipeline,
        )

        self.processor = AutoProcessor.from_pretrained(
            model_id, revision=revision, trust_remote_code=False
        )
        self.model = AutoModelForSpeechSeq2Seq.from_pretrained(
            model_id, revision=revision, trust_remote_code=False,
        ).eval()
        self.pipe = pipeline(
            "automatic-speech-recognition",
            model=self.model,
            tokenizer=self.processor.tokenizer,
            feature_extractor=self.processor.feature_extractor,
            # Sans ce découpage, tout énoncé de plus de 30 s fait échouer
            # Whisper sur une incompatibilité de dimensions.
            chunk_length_s=30,
        )

    def _load_ctc(self, model_id, revision):
        from transformers import AutoModelForCTC, AutoProcessor

        self.processor = AutoProcessor.from_pretrained(
            model_id, revision=revision, trust_remote_code=False
        )
        self.model = AutoModelForCTC.from_pretrained(
            model_id, revision=revision, trust_remote_code=False,
        ).eval()
        self.pipe = None

        # MMS multilingue : les poids fongbe vivent dans un adaptateur séparé,
        # sans quoi le modèle transcrit dans une tout autre langue.
        try:
            if getattr(self.model.config, "adapter_attn_dim", None):
                self.model.load_adapter("fon")
                self.processor.tokenizer.set_target_lang("fon")
                log("adaptateur MMS « fon » chargé")
        except Exception as exc:  # noqa: BLE001 — modèle déjà spécialisé
            log(f"pas d'adaptateur fon ({type(exc).__name__}) — poids tels quels")

    def transcribe(self, arrays: list) -> list[str]:
        """Transcrit une tranche. À n'appeler que depuis un créneau GPU."""
        torch = self.torch
        dev = "cuda" if torch.cuda.is_available() else "cpu"
        self.model.to(dev)
        if self.pipe is not None:
            self.pipe.device = torch.device(dev)

        if self.seq2seq:
            return [(self.pipe(a)["text"] or "").strip() for a in arrays]

        out: list[str] = []
        for i in range(0, len(arrays), BATCH_CTC):
            batch = arrays[i:i + BATCH_CTC]
            inputs = self.processor(
                batch, sampling_rate=16000, return_tensors="pt", padding=True
            )
            inputs = {k: v.to(dev) for k, v in inputs.items()}
            with torch.inference_mode():
                logits = self.model(**inputs).logits
            ids = torch.argmax(logits, dim=-1).cpu()
            out.extend(t.strip() for t in self.processor.batch_decode(ids))
        return out

    def hardware(self) -> str:
        try:
            if self.torch.cuda.is_available():
                return f"HF Space ZeroGPU ({self.torch.cuda.get_device_name(0)})"
        except Exception:  # noqa: BLE001
            pass
        return "HF Space CPU"


# Le modèle reste une variable de module : ZeroGPU exécute la fonction
# décorée dans un processus dérivé, qui hérite de la mémoire sans rien
# sérialiser. Le passer en argument recopierait tout le modèle à chaque appel.
_CURRENT: Transcriber | None = None


@GPU(duration=GPU_DURATION)
def gpu_transcribe(arrays: list) -> list[str]:
    return _CURRENT.transcribe(arrays)


def _is_quota_error(exc: Exception) -> bool:
    """Le quota ZeroGPU est-il épuisé (par opposition à une vraie panne) ?"""
    texte = f"{type(exc).__name__} {exc}".lower()
    return any(m in texte for m in ("quota", "gpu task aborted",
                                    "no gpu is currently available",
                                    "exceeded your"))


# --- une évaluation -----------------------------------------------------

def check_model(model_id: str) -> str:
    """Révision figée du modèle. Lève si le dépôt est inutilisable."""
    from huggingface_hub import HfApi

    # Sans token : un modèle soumis doit être public pour être vérifiable.
    api = HfApi(token=False)
    try:
        info = api.model_info(model_id, files_metadata=False)
    except Exception:  # noqa: BLE001
        raise RuntimeError("Modèle introuvable ou privé sur le Hub.")

    # On accepte aussi les poids `.bin` : la moitié des modèles fongbe
    # publics datent d'avant safetensors, les exclure viderait le classement.
    # Le risque du pickle est couvert ailleurs — transformers charge les .bin
    # avec `weights_only=True` (donc sans exécution de code) depuis la 4.36,
    # et `trust_remote_code=False` bloque le code personnalisé du dépôt.
    names = [s.rfilename for s in (info.siblings or [])]
    if not any(n.endswith((".safetensors", ".bin", ".ckpt", ".pt"))
               for n in names):
        raise RuntimeError(
            "Aucun fichier de poids trouvé dans le dépôt "
            "(.safetensors ou .bin attendu)."
        )
    return info.sha


def run_job(job: dict) -> None:
    request, bench = job["request"], job["benchmark"]
    rid, model_id = request["id"], request["model_id"]
    _state.update(model=model_id, message="préparation", done=0, total=0)
    log(f"▶ {model_id} sur {bench['id']}")

    revision = check_model(model_id)

    # Corpus d'abord : le token de lecture ne doit jamais coexister avec un
    # modèle inconnu déjà chargé en mémoire.
    ds = load_corpus(bench)
    total = len(ds)

    verdict = rpc("space_begin", {
        "p_token": SPACE_TOKEN, "p_id": rid,
        "p_revision": revision, "p_total": total,
    })
    if verdict == "duplicate":
        log(f"↷ {model_id} : déjà évalué à cette révision")
        _state.update(message="doublon ignoré", model=None)
        return

    global _CURRENT
    tr = _CURRENT = Transcriber(model_id, revision)
    chunk = tr.chunk_size

    saved = job.get("progress") or {}
    counters = saved.get("counters") or fe.new_counters()
    counters = {k: int(counters.get(k, 0)) for k in fe.COUNTER_KEYS}
    next_chunk = int(saved.get("next_chunk") or 0)
    audio_s = float(saved.get("audio_seconds") or 0.0)
    compute_s = float(saved.get("compute_seconds") or 0.0)
    if next_chunk:
        log(f"↻ reprise à la tranche {next_chunk} ({next_chunk * chunk} énoncés)")

    # Chaque tranche tient dans un créneau GPU et vaut battement de cœur :
    # elle reste bien en deçà des 20 min au-delà desquelles la base
    # considère la tâche abandonnée et la redonne.
    n_chunks = (total + chunk - 1) // chunk
    skipped = 0

    ci = next_chunk
    while ci < n_chunks:
        rows = ds.select(range(ci * chunk, min((ci + 1) * chunk, total)))
        arrays, refs = [], []
        for row in rows:
            try:
                arrays.append(decode_audio(row["audio"]))
                refs.append(row["transcription"])
            except Exception:  # noqa: BLE001 — énoncé illisible : pas la faute du modèle
                skipped += 1

        if arrays:
            t0 = time.time()
            try:
                hyps = gpu_transcribe(arrays)
            except Exception as exc:  # noqa: BLE001
                if not _is_quota_error(exc):
                    raise
                # Quota ZeroGPU épuisé : on patiente et on refait la même
                # tranche. Le point de reprise est déjà en base, rien n'est
                # perdu — c'est précisément ce que le découpage permet.
                log(f"⏸ quota GPU épuisé ({exc}) — reprise dans "
                    f"{QUOTA_WAIT_SECONDS // 60} min à la tranche {ci}")
                _state.update(message="quota GPU épuisé — en attente")
                time.sleep(QUOTA_WAIT_SECONDS)
                continue
            compute_s += time.time() - t0
            audio_s += sum(len(a) for a in arrays) / 16000.0
            fe.accumulate(counters, refs, hyps)

        done = min((ci + 1) * chunk, total)
        _state.update(message="évaluation", done=done, total=total)
        rpc("space_checkpoint", {
            "p_token": SPACE_TOKEN, "p_id": rid, "p_next_chunk": ci + 1,
            "p_done": done, "p_counters": counters,
            "p_audio": round(audio_s, 2), "p_compute": round(compute_s, 2),
        })
        rtfx = audio_s / compute_s if compute_s else 0
        log(f"  {done}/{total} — RTFx {rtfx:.2f}")
        ci += 1

    if skipped:
        log(f"⚠ {skipped} énoncés illisibles écartés")
    if not counters["n_scored"]:
        raise RuntimeError("Aucune transcription exploitable produite.")

    metrics = fe.finalize(counters)
    rtfx = round(audio_s / compute_s, 3) if compute_s else None
    rpc("space_finish", {
        "p_token": SPACE_TOKEN, "p_id": rid, "p_status": "done",
        "p_metrics": {k: metrics[k] for k in (
            "wer", "cer", "mer", "wil", "wer_seg", "cer_seg", "wer_ton", "twer"
        )},
        "p_meta": {
            "architecture": tr.architecture,
            "decoder_type": tr.decoder_type,
            "model_params": tr.model_params,
            "rtfx": rtfx,
            "rtf": round(1 / rtfx, 4) if rtfx else None,
            "eval_seconds": round(compute_s, 1),
            "hardware": tr.hardware(),
            "worker_version": WORKER_VERSION,
        },
        "p_error": None,
    })
    log(f"✔ {model_id} — WER_seg {metrics['wer_seg']:.1%} "
        f"T-WER {metrics['twer']} RTFx {rtfx}")
    _state.update(message="terminé", model=None, done=0, total=0)


def loop() -> None:
    if not SPACE_TOKEN:
        _state["message"] = "SPACE_TOKEN absent — évaluateur à l'arrêt"
        log("SPACE_TOKEN absent : aucune évaluation ne sera lancée.")
        return
    if not _DATA_TOKEN:
        _state["message"] = "HF_DATA_TOKEN absent — évaluateur à l'arrêt"
        log("HF_DATA_TOKEN absent : le corpus privé est illisible.")
        return

    # Un token de lecture mal choisi ne se voit qu'au premier échec, après
    # qu'une demande a déjà été prise et perdue : autant le dire au démarrage.
    try:
        from huggingface_hub import HfApi, whoami

        who = whoami(token=_DATA_TOKEN)
        role = (who.get("auth") or {}).get("accessToken", {}).get("role")
        orgs = [o["name"] for o in who.get("orgs", [])]
        log(f"token de lecture : compte « {who['name']} », portée « {role} », "
            f"orgs {orgs or '(aucune)'}")

        # Appartenir à l'organisation ne suffit pas : un token fine-grained
        # ne voit que les dépôts explicitement cochés. On teste l'accès réel.
        api = HfApi(token=_DATA_TOKEN)
        for b in requests.get(
            f"{REST}/benchmarks", headers=HEADERS,
            params={"is_active": "eq.true", "select": "id,hf_dataset_id"},
            timeout=30,
        ).json():
            try:
                api.dataset_info(b["hf_dataset_id"])
                log(f"  ✓ {b['hf_dataset_id']} lisible")
            except Exception as exc:  # noqa: BLE001
                log(f"  ✗ {b['hf_dataset_id']} INACCESSIBLE "
                    f"({type(exc).__name__}) — corriger HF_DATA_TOKEN : il "
                    f"faut un token « read », ou un token fine-grained avec "
                    f"« Read access to contents » cochée sur ce dépôt.")
    except Exception as exc:  # noqa: BLE001
        log(f"⚠ HF_DATA_TOKEN refusé par le Hub : {type(exc).__name__}")

    log("évaluateur démarré")
    while True:
        job = None
        try:
            job = rpc("space_claim", {"p_token": SPACE_TOKEN})
        except Exception as exc:  # noqa: BLE001 — réseau : on réessaie
            log(f"⚠ space_claim : {exc}")

        if not job:
            _state.update(message="aucune tâche en attente", model=None)
            time.sleep(POLL_SECONDS)
            continue

        try:
            run_job(job)
        except Exception as exc:  # noqa: BLE001 — tout échec doit être publié
            log(f"✘ échec : {exc}")
            traceback.print_exc()
            _state.update(message=f"échec : {exc}", model=None)
            try:
                rpc("space_finish", {
                    "p_token": SPACE_TOKEN, "p_id": job["request"]["id"],
                    "p_status": "failed", "p_metrics": None, "p_meta": None,
                    # Message court et sans extrait du corpus.
                    "p_error": f"{type(exc).__name__}: {exc}"[:300],
                })
            except Exception:  # noqa: BLE001
                traceback.print_exc()
        finally:
            import gc

            gc.collect()


_thread: threading.Thread | None = None


def start() -> threading.Thread:
    """Démarre l'évaluateur. Sans effet s'il tourne déjà.

    Le lanceur du Space peut importer app.py plus d'une fois : sans cette
    garde, deux threads se disputeraient la même tâche.
    """
    global _thread
    if _thread is None or not _thread.is_alive():
        _thread = threading.Thread(target=loop, name="fonbench-eval",
                                   daemon=True)
        _thread.start()
    return _thread