grabber-ai/tools/retrain.py

"""Переобучение и проверка моделей grabber-ai (используется из tools/collect.py).

Шаги вынесены в функции, чтобы меню вызывало их по очереди:
    prepare()                   data/labels.csv -> data/processed
    train_segmenter(out_dir)    обучение сегментатора страниц
    train_main()                основная модель (тип, автор, 18+): долго
    run_tests()                 tests/run_tests.py
    compare(dirs)               правила vs сегментаторы по трекерам
    promote(new_dir)            новые веса -> models/segmenter (старые в архив)
"""

from __future__ import annotations

import re
import shutil
import subprocess
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
MODELS = ROOT / "models"
MIN_N = 30            # трекеры с меньшим числом страниц в вывод попадают, но решение не определяют

RE_CELL = re.compile(r"(\w+) n=(\d+) P=([\d.]+) R=([\d.]+)")


def run(args: list[str]) -> int:
    """Запуск шага тем же Python; вывод идёт прямо в консоль."""
    print("\n$ python " + " ".join(args), flush=True)
    return subprocess.call([sys.executable, *args], cwd=ROOT)


def prepare() -> bool:
    return run(["-m", "src.prepare_data"]) == 0


def train_segmenter(out_dir: str = "models/segmenter_v2", epochs: int = 4) -> bool:
    return run(["-m", "src.segmenter", "--train", "--out", out_dir,
                "--epochs", str(epochs)]) == 0


def backup_main() -> str | None:
    """Копия models/best перед обучением: src.train перезаписывает каталог."""
    import shutil
    import time
    src = ROOT / "models" / "best"
    if not src.exists():
        return None
    name = f"best_backup_{time.strftime('%Y%m%d_%H%M%S')}"
    shutil.copytree(src, ROOT / "models" / name)
    return name


def train_main() -> bool:
    """Основная модель: 5-fold, ~15-20 минут на GPU. Старый models/best сохраняется в models/best_backup_*."""
    saved = backup_main()
    if saved:
        print(f"копия прежней модели: models/{saved}", flush=True)
    return run(["-m", "src.train"]) == 0


PROFILES = {                      # те же имена, что в src/train.py:MODEL_CONFIGS (без импорта torch в меню)
    "tiny2": "rubert-tiny2, 30M — самая быстрая, годится для ONNX/CPU",
    "base": "ruBert-base, 178M — текущий рабочий выбор",
    "large": "ruBert-large, 427M — та же семья, втрое больше: обучение ~3× дольше, инференс на CPU медленнее",
    "modern": "RuModernBERT-base, 149M (2024) — новая архитектура; ONNX-экспорт проекта её не поддерживает",
}


def train_models(profiles: list[str], folds: int = 0, tag: str = "") -> bool:
    """Обучение профилей отдельными запусками в models/runs/: models/best и reports/ не меняются."""
    args = ["-m", "src.train", "--runs", "--models", *profiles, "--folds", str(folds)]
    if tag:
        args += ["--tag", tag]
    return run(args) == 0


def evaluate_models(out_json: Path, refresh: bool = False) -> dict | None:
    """Оценка всех моделей на текущей тестовой выборке (подпроцесс: GPU освобождается после выхода)."""
    out_json.unlink(missing_ok=True)
    args = ["-m", "src.runs", "--out", str(out_json)] + (["--refresh"] if refresh else [])
    if run(args) != 0 or not out_json.exists():
        return None
    import json
    return json.loads(out_json.read_text(encoding="utf-8"))


def run_tests() -> bool:
    return run(["tests/run_tests.py"]) == 0


def _e2e(seg_dir: str) -> dict:
    """{'rules': {трекер: (n, P, R)}, 'model': {...}} на страницах, которых модель не видела."""
    proc = subprocess.run([sys.executable, "tests/e2e_segmenter.py", seg_dir], cwd=ROOT,
                          capture_output=True, text=True, encoding="utf-8", errors="replace")
    if proc.returncode != 0:
        tail = "\n".join(proc.stderr.strip().splitlines()[-6:])
        raise RuntimeError(f"e2e_segmenter упал на {seg_dir}:\n{tail}")
    test_block = proc.stdout.split("train-страницы")[0]
    out: dict[str, dict] = {}
    for line in test_block.splitlines():
        m = re.match(r"\s+(rules|model)\s*:(.*)", line)
        if m:
            out[m.group(1)] = {t: (int(n), float(p), float(r))
                               for t, n, p, r in RE_CELL.findall(m.group(2))}
    return out


def compare(seg_dirs: list[str]) -> dict:
    """Печатает таблицу: правила и каждый сегментатор по трекерам.

    Возвращает {'rules': {...}, <dir>: {...}} для решения о подмене весов.
    """
    results: dict[str, dict] = {}
    for d in seg_dirs:
        print(f"\nзамер: {d} (несколько минут)…", flush=True)
        res = _e2e(d)
        results[d] = res.get("model", {})
        results.setdefault("rules", res.get("rules", {}))
    trackers = sorted(results["rules"])
    cols = ["rules"] + list(seg_dirs)
    print("\nPrecision / Recall на страницах, которых сегментатор не видел при обучении")
    print(f"{'трекер':10s} {'n':>4s}  " + "  ".join(f"{Path(c).name:>16s}" for c in cols))
    for t in trackers:
        n = results["rules"][t][0]
        cells = []
        for c in cols:
            v = results[c].get(t)
            cells.append(f"{v[1]:.3f} / {v[2]:.3f}" if v else "-")
        warn = "  (мало страниц)" if n < MIN_N else ""
        print(f"{t:10s} {n:4d}  " + "  ".join(f"{x:>16s}" for x in cells) + warn)
    return results


def verdict(results: dict, old: str, new: str) -> tuple[bool, str]:
    """Новые веса лучше старых? Только по трекерам с n >= MIN_N.

    Лучше: суммарная precision не ниже, recall не просел больше чем на 0.02
    ни на одном трекере.
    """
    big = [t for t, v in results["rules"].items() if v[0] >= MIN_N]
    if not big:
        return False, "нет трекеров с достаточным числом страниц"
    dp = sum(results[new][t][1] - results[old][t][1] for t in big if t in results[new] and t in results[old])
    worst_r = min(results[new][t][2] - results[old][t][2] for t in big
                  if t in results[new] and t in results[old])
    ok = dp >= 0 and worst_r >= -0.02
    return ok, (f"Δprecision (сумма по {', '.join(big)}) = {dp:+.3f}, "
                f"худшее Δrecall = {worst_r:+.3f}")


def promote(new_dir: str, current: str = "models/segmenter") -> str:
    """Старые веса -> models/segmenter_old_<время>, новые -> models/segmenter."""
    cur, new = ROOT / current, ROOT / new_dir
    archive = ""
    if cur.exists():
        archive = f"{cur.name}_old_{time.strftime('%Y%m%d_%H%M%S')}"
        cur.rename(cur.with_name(archive))
    shutil.move(str(new), str(cur))
    return archive