grabber-ai/src/runs.py

"""Реестр обученных моделей: честное сравнение на одной и той же тестовой выборке и выбор models/best.

Зачем отдельный модуль. Метрики, записанные при обучении, относятся к тестовой выборке того момента;
после пересборки датасета (другие тексты, другие группы) выборка меняется, и «0.966 тогда» нельзя
сравнивать с «0.97 сейчас». Поэтому любая модель (запуск из models/runs, models/best, архивные копии)
оценивается на ТЕКУЩИХ тестовых строках, результат кэшируется в <каталог модели>/eval.json под
отпечатком выборки. Модель получает вход в том формате, на котором обучена: без «[RUS AC3 …]»,
если в её label_maps нет audio_context.

Сравнение парное: по каждой строке видно, что новая модель починила и что сломала; при малой
разнице показывается, что она в пределах шума (знаковый критерий).
"""

from __future__ import annotations

import hashlib
import json
import math
import re
import shutil
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
MODELS = ROOT / "models"
RUNS = MODELS / "runs"
DATASET = ROOT / "data" / "processed" / "dataset.json"
EVAL_FILE = "eval.json"
UNSCORABLE = "__unscorable__"
RE_TAIL = re.compile(r"\s\[(?:[A-Z]{3}|\?)(?: [^\]\n]*)?\]$")      # «[RUS AC3 6ch 448kbps 18+]» в конце входа


# ---------------------------------------------------------------- данные
def plain_text(row: dict) -> str:
    """Вход без аудио-контекста (для моделей, обученных на чистом тексте дорожки)."""
    tail = row.get("tail")
    if tail is None:                                   # датасет собран до появления поля tail
        return RE_TAIL.sub("", row["text"])
    suffix = " " + tail
    return row["text"][:-len(suffix)] if tail and row["text"].endswith(suffix) else row["text"]


def test_rows(dataset: Path | str | None = None) -> list[dict]:
    rows = json.loads(Path(dataset or DATASET).read_text(encoding="utf-8"))
    return [r for r in rows if r["split"] == "test"]


def fingerprint(rows: list[dict]) -> str:
    """Отпечаток тестовой выборки: те же строки с теми же метками (порядок не важен)."""
    keys = sorted(json.dumps([r["text"], r["translation_type"], r.get("spans") or r.get("span"), bool(r.get("is_18plus")),
                              r.get("lang")], ensure_ascii=False) for r in rows)
    return hashlib.sha1("\n".join(keys).encode("utf-8")).hexdigest()[:10]


# ---------------------------------------------------------------- статистика
def wilson(k: int, n: int, z: float = 1.96) -> tuple[float, float]:
    """95% интервал доли k из n."""
    if not n:
        return 0.0, 0.0
    p = k / n
    d = 1 + z * z / n
    c = (p + z * z / (2 * n)) / d
    h = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / d
    return max(0.0, c - h), min(1.0, c + h)


def sign_test(fixed: int, broken: int) -> float:
    """Двусторонний точный знаковый критерий: случайна ли разница между «починил» и «сломал»."""
    n = fixed + broken
    if not n:
        return 1.0
    k = min(fixed, broken)
    return min(1.0, 2 * sum(math.comb(n, i) for i in range(k + 1)) / 2 ** n)


def paired(base: list[int], new: list[int]) -> dict:
    """Построчное сравнение двух моделей на одних и тех же строках."""
    fixed = sum(1 for a, b in zip(base, new) if not a and b)
    broken = sum(1 for a, b in zip(base, new) if a and not b)
    return {"fixed": fixed, "broken": broken, "p": sign_test(fixed, broken)}


# ---------------------------------------------------------------- оценка чекпоинта
def store_eval(model_dir: Path | str, fp: str, entry: dict) -> None:
    path = Path(model_dir) / EVAL_FILE
    data = json.loads(path.read_text(encoding="utf-8")) if path.exists() else {}
    data[fp] = entry
    path.write_text(json.dumps(data, ensure_ascii=False), encoding="utf-8")


def load_eval(model_dir: Path | str, fp: str) -> dict | None:
    path = Path(model_dir) / EVAL_FILE
    if not path.exists():
        return None
    try:
        return json.loads(path.read_text(encoding="utf-8")).get(fp)
    except ValueError:
        return None


def eval_checkpoint(model_dir: Path | str, rows: list[dict], device: str | None = None) -> dict:
    """Оценка чекпоинта на строках rows (тестовая выборка). Нужен GPU/CPU и torch."""
    import torch
    from torch.utils.data import DataLoader
    from transformers import AutoTokenizer

    from .model import load_model
    from .train import Rows, evaluate, make_collate

    d = Path(model_dir)
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    maps = json.loads((d / "label_maps.json").read_text(encoding="utf-8"))
    type2id = {t: i for i, t in enumerate(maps["types"])}
    type2id[UNSCORABLE] = len(type2id)                 # типы, которых модель не знает, считаются ошибкой
    id2type = {i: t for t, i in type2id.items()}
    model = load_model(d, device)
    lang2id = ({l: i for i, l in enumerate(maps.get("langs") or [])} or None) if model.lang_head is not None else None
    use_ctx = bool(maps.get("audio_context"))
    view, n_unscorable = [], 0
    for r in rows:
        known = r["translation_type"] in type2id
        n_unscorable += not known
        view.append({**r, "text": r["text"] if use_ctx else plain_text(r),
                     "translation_type": r["translation_type"] if known else UNSCORABLE})
    tokenizer = AutoTokenizer.from_pretrained(d)
    loader = DataLoader(Rows(view), batch_size=64, collate_fn=make_collate(tokenizer, type2id, lang2id=lang2id))
    per_row: list[int] = []
    metrics = evaluate(model, loader, device, id2type, per_row=per_row)
    del model
    if device == "cuda":
        torch.cuda.empty_cache()
    return {"metrics": metrics, "per_row": per_row, "n_unscorable": n_unscorable}


def cached_eval(model_dir: Path | str, rows: list[dict], fp: str, refresh: bool = False,
                device: str | None = None) -> dict:
    hit = None if refresh else load_eval(model_dir, fp)
    if hit:
        return hit
    entry = eval_checkpoint(model_dir, rows, device)
    store_eval(model_dir, fp, entry)
    return entry


# ---------------------------------------------------------------- каталог моделей
def _read_json(path: Path) -> dict:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return {}


def describe(model_dir: Path | str) -> dict:
    d = Path(model_dir)
    cfg = _read_json(d / "dub_config.json")
    run = _read_json(d / "run.json") or _read_json(d / "train_config.json")
    maps = _read_json(d / "label_maps.json")
    tok_cfg = _read_json(d / "tokenizer_config.json")
    weights = d / "model.pt"
    return {
        "dir": d, "name": d.name, "profile": run.get("name", ""), "trained_fp": run.get("fingerprint"),
        "backbone": cfg.get("backbone", "?"),
        "mb": round(weights.stat().st_size / 1e6) if weights.exists() else 0,
        "created": run.get("created") or time.strftime("%Y-%m-%d %H:%M", time.localtime(d.stat().st_mtime)),
        "seconds": (run.get("final") or {}).get("train_seconds"),
        "audio_context": bool(maps.get("audio_context")), "lang_head": bool(cfg.get("n_langs")),
        "token_type_ids": "token_type_ids" in (tok_cfg.get("model_input_names") or ["token_type_ids"]),
    }


def candidates() -> list[Path]:
    """Модели, из которых можно выбирать: best, запуски, архивные копии, прежние models/<имя>."""
    found: list[Path] = []
    best = MODELS / "best"
    if (best / "model.pt").exists():
        found.append(best)
    others = [p for p in list(RUNS.glob("*")) + list(MODELS.glob("best_backup_*")) + [MODELS / "tiny2", MODELS / "base"]
              if (p / "model.pt").exists()]
    others.sort(key=lambda p: p.stat().st_mtime, reverse=True)
    seen = {(p / "train_config.json").read_bytes() for p in found if (p / "train_config.json").exists()}
    for p in others:
        cfg = p / "train_config.json"
        key = cfg.read_bytes() if cfg.exists() else None
        if key is not None and key in seen:            # models/base — копия models/best и т.п.
            continue
        if key is not None:
            seen.add(key)
        found.append(p)
    return found


def compare(dirs: list[Path | str], rows: list[dict] | None = None, baseline: Path | str | None = None,
            refresh: bool = False, on_step=None) -> list[dict]:
    """Метрики каждой модели на ТЕКУЩЕЙ тестовой выборке + парное сравнение с baseline (models/best)."""
    rows = rows if rows is not None else test_rows()
    fp = fingerprint(rows)
    base_dir = Path(baseline) if baseline else None
    base_entry = None
    if base_dir and (base_dir / "model.pt").exists():
        base_entry = cached_eval(base_dir, rows, fp, refresh, None)
    out = []
    for i, d in enumerate(dirs):
        d = Path(d)
        if on_step:
            on_step(i, len(dirs), d)
        entry = base_entry if base_dir and d.resolve() == base_dir.resolve() else cached_eval(d, rows, fp, refresh)
        m = entry["metrics"]
        k = sum(entry["per_row"])
        item = {**describe(d), "metrics": m, "n": len(entry["per_row"]), "n_unscorable": entry.get("n_unscorable", 0),
                "ci": wilson(k, len(entry["per_row"])), "per_row": entry["per_row"], "fingerprint": fp, "vs": None}
        # модель обучена не на этом разбиении: часть текущих тестовых строк она могла видеть при обучении
        item["leak_risk"] = item["trained_fp"] != fp
        if base_entry and d.resolve() != base_dir.resolve():
            item["vs"] = paired(base_entry["per_row"], entry["per_row"])
        out.append(item)
    return out


def verdict(item: dict) -> str:
    """Короткий вывод про модель относительно models/best."""
    vs = item.get("vs")
    if not vs:
        return "текущая" if item["name"] == "best" else "—"
    net = vs["fixed"] - vs["broken"]
    if vs["p"] > 0.05:
        return f"в пределах шума (+{vs['fixed']}/-{vs['broken']}, p={vs['p']:.2f})"
    return f"{'лучше' if net > 0 else 'хуже'} (+{vs['fixed']}/-{vs['broken']}, p={vs['p']:.3f})"


def warnings_for(model_dir: Path | str) -> list[str]:
    """Что учесть перед заменой models/best этой моделью."""
    info = describe(model_dir)
    best = describe(MODELS / "best") if (MODELS / "best" / "model.pt").exists() else None
    out = []
    if best and info["backbone"] != best["backbone"]:
        repo = "models--" + info["backbone"].replace("/", "--")
        cached = (Path.home() / ".cache" / "huggingface" / "hub" / repo).exists()
        out.append(f"другая основа ({info['backbone']}): на машине сервиса её нужно скачать с Hugging Face "
                   f"(у вас {'уже в кэше' if cached else 'в кэше нет'}).")
    if best and info["mb"] > best["mb"] * 1.5:
        out.append(f"вес {info['mb']} МБ против {best['mb']} МБ: инференс на CPU заметно медленнее.")
    if not info["token_type_ids"]:
        out.append("у токенизатора нет token_type_ids: src/export_onnx.py и клиенты ONNX (Rust/Node) "
                   "рассчитаны на BERT и такую модель не примут без доработки.")
    if (MODELS / "best-onnx").exists():
        out.append("models/best-onnx останется от прежней модели: пересоздайте `python -m src.export_onnx`.")
    return out


def promote(model_dir: Path | str) -> str | None:
    """Копирует модель в models/best; прежняя best уходит в models/best_backup_<время>."""
    src, best = Path(model_dir), MODELS / "best"
    if best.exists() and src.resolve() == best.resolve():
        return None
    tmp = MODELS / ".best_new"
    shutil.rmtree(tmp, ignore_errors=True)
    shutil.copytree(src, tmp)
    archive = None
    if best.exists():
        archive = f"best_backup_{time.strftime('%Y%m%d_%H%M%S')}"
        best.rename(MODELS / archive)
    try:
        tmp.rename(best)
    except OSError:
        if archive:
            (MODELS / archive).rename(best)
        raise
    return archive


def deletable(model_dir: Path | str) -> bool:
    """Удалять можно запуски и архивные копии, но не models/best и не models/<профиль>."""
    p = Path(model_dir)
    return p.name != "best" and (p.parent == RUNS or p.name.startswith("best_backup_"))


def delete(model_dir: Path | str) -> None:
    if not deletable(model_dir):
        raise ValueError(f"{model_dir}: удалять можно только запуски и архивные копии")
    shutil.rmtree(model_dir)


def _jsonable(items: list[dict]) -> list[dict]:
    out = []
    for it in items:
        d = {k: v for k, v in it.items() if k != "per_row"}
        d["dir"] = str(it["dir"])
        d["ci"] = list(it["ci"])
        out.append(d)
    return out


def main() -> None:
    import argparse
    ap = argparse.ArgumentParser(description="Оценка моделей на текущей тестовой выборке (кэш в <модель>/eval.json)")
    ap.add_argument("--out", required=True, help="куда записать JSON с результатами")
    ap.add_argument("--dirs", nargs="*", help="каталоги моделей (по умолчанию все кандидаты)")
    ap.add_argument("--data", default=None)
    ap.add_argument("--refresh", action="store_true", help="пересчитать, игнорируя кэш")
    args = ap.parse_args()
    dirs = [Path(d) for d in args.dirs] if args.dirs else candidates()
    rows = test_rows(args.data)
    print(f"тестовая выборка: {len(rows)} строк, отпечаток {fingerprint(rows)}", flush=True)
    items = compare(dirs, rows, baseline=MODELS / "best", refresh=args.refresh,
                    on_step=lambda i, n, d: print(f"оценка {i + 1}/{n}: {d.name}", flush=True))
    Path(args.out).write_text(json.dumps({"fingerprint": fingerprint(rows), "n": len(rows),
                                          "items": _jsonable(items)}, ensure_ascii=False), encoding="utf-8")


if __name__ == "__main__":
    main()