"""Обучение multi-task модели (тип озвучки + автор) и сравнение конфигураций. Запуск: python -m src.train # обе модели: 5-fold CV + финальные python -m src.train --models tiny2 # только rubert-tiny2 python -m src.train --models base --folds 3 --epochs-override 5 Этапы для каждой модели: 1. 5-fold CV на train-части (группы: релиз + одинаковый текст) -> средние метрики; 2. финальное обучение на train (early stopping на 10% val-группах); 3. однократная оценка на held-out test -> reports/final_<имя>.json, confusion matrix и кривые обучения -> reports/*.png; 4. чекпоинт -> models/<имя>/. Лучшая модель копируется в models/best. """ from __future__ import annotations import argparse import json import random import re import shutil import time from pathlib import Path import numpy as np import torch from sklearn.metrics import accuracy_score, confusion_matrix, f1_score from torch.utils.data import DataLoader, Dataset from transformers import AutoTokenizer, get_linear_schedule_with_warmup from .model import DubTagger, predict_with_lang, save_model from .spans import MAX_LEN, decode_span # noqa: F401 (прежнее место; сервису не нужны sklearn/matplotlib) MODEL_CONFIGS = { "tiny2": {"backbone": "cointegrated/rubert-tiny2", "lr": 1e-4, "batch": 64, "epochs": 20}, "base": {"backbone": "ai-forever/ruBert-base", "lr": 3e-5, "batch": 32, "epochs": 10}, # крупные профили (RTX 5070 Ti 16 ГБ: bf16 + batch 16 помещаются); в обучение по умолчанию не входят "large": {"backbone": "ai-forever/ruBert-large", "lr": 2e-5, "batch": 16, "epochs": 8}, "modern": {"backbone": "deepvk/RuModernBERT-base", "lr": 5e-5, "batch": 32, "epochs": 8}, } DEFAULT_MODELS = ["tiny2", "base"] # как в docs/training-windows.md: python -m src.train PROFILE_INFO = { "tiny2": "rubert-tiny2, 30M: самая быстрая, годится для ONNX/CPU", "base": "ruBert-base, 178M: текущий рабочий выбор", "large": "ruBert-large, 427M: та же семья, но втрое больше; ~3× дольше и тяжелее на CPU", "modern": "RuModernBERT-base, 149M (2024): новая архитектура; ONNX-экспорт проекта её пока не поддерживает", } SEED = 42 def set_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) class Rows(Dataset): def __init__(self, rows: list[dict]): self.rows = rows def __len__(self): return len(self.rows) def __getitem__(self, i): return self.rows[i] def make_collate(tokenizer, type2id: dict[str, int], max_len: int = MAX_LEN, lang2id: dict[str, int] | None = None): """Токенизация батча + BIO-метки по символьным спанам (через offsets). Возвращает (enc, types, bio, offsets): offsets — список списков (a, b) по токенам, нужен для декодирования спана при оценке/инференсе. """ def collate(batch: list[dict]): texts = [b["text"] for b in batch] enc = tokenizer(texts, truncation=True, max_length=max_len, padding=True, return_tensors="pt", return_offsets_mapping=True) offsets = enc.pop("offset_mapping").tolist() # O по умолчанию (включая спецтокены — модель учится говорить «не автор»); # -100 только для паддинга, он не входит в лосс bio = torch.zeros_like(enc["input_ids"]) bio[enc["attention_mask"] == 0] = -100 for i, b in enumerate(batch): span = b.get("span") if not span: continue s, e = span first = True for t, (a, z) in enumerate(offsets[i]): if z > a and a < e and z > s: # токен пересекает спан bio[i, t] = 1 if first else 2 # B / I first = False types = torch.tensor([type2id[b["translation_type"]] for b in batch]) ages = torch.tensor([int(bool(b.get("is_18plus"))) for b in batch]) if lang2id: # неизвестный язык (None) = -100: в лосс и метрику не входит langs = torch.tensor([lang2id.get(b.get("lang"), -100) if b.get("lang") else -100 for b in batch]) return enc, types, bio, ages, offsets, langs return enc, types, bio, ages, offsets return collate def _NORM(s: str | None) -> str | None: return re.sub(r"\s+", " ", s).strip().lower() if s else None @torch.no_grad() def evaluate(model, loader, device, id2type, per_row: list | None = None) -> dict: """Метрики на выборке: тип (acc, macro-F1), спан (exact), 18+ (acc, F1). Exact match спана сравнивает нормализованный текст предсказанного спана с текстом золотого спана (None == None считается совпадением). joint_exact — строка распознана полностью: тип + автор + 18+. """ model.eval() y_ids, p_ids = [], [] span_ok, y_age, p_age = [], [], [] y_lang, p_lang = [], [] for enc, types, _, ages, offsets, *rest in loader: enc_dev = {k: v.to(device) for k, v in enc.items()} with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device == "cuda"): if model.lang_head is not None and rest: type_logits, bio_logits, age_logits, lang_logits = predict_with_lang(model, **enc_dev) y_lang += rest[0].tolist() p_lang += lang_logits.argmax(-1).tolist() else: type_logits, bio_logits, age_logits = model.predict(**enc_dev) p_ids += type_logits.argmax(-1).tolist() y_ids += types.tolist() p_age += age_logits.argmax(-1).tolist() y_age += ages.tolist() bio_pred = bio_logits.argmax(-1).cpu().numpy() attn = enc["attention_mask"].tolist() base = len(span_ok) rows = loader.dataset.rows for i in range(len(types)): pred = decode_span(bio_pred[i], offsets[i], attn[i]) gold = rows[base + i].get("span") pred_txt = rows[base + i]["text"][pred[0]:pred[1]] if pred else None gold_txt = rows[base + i]["text"][gold[0]:gold[1]] if gold else None span_ok.append(_NORM(pred_txt) == _NORM(gold_txt)) y = [id2type[v] for v in y_ids] p = [id2type[v] for v in p_ids] type_ok = [a == b for a, b in zip(y, p)] age_ok = [a == b for a, b in zip(y_age, p_age)] if per_row is not None: # 1 = строка распознана полностью (тип + автор + 18+), по порядку loader per_row.extend(int(a and b and c) for a, b, c in zip(type_ok, span_ok, age_ok)) lang_pairs = [(a, b) for a, b in zip(y_lang, p_lang) if a != -100] out_lang = ({"lang_accuracy": float(np.mean([a == b for a, b in lang_pairs])), "lang_n": len(lang_pairs)} if lang_pairs else {}) return { **out_lang, "type_accuracy": float(np.mean(type_ok)) if type_ok else 0.0, "type_macro_f1": f1_score(y, p, average="macro", zero_division=0), "span_exact": float(np.mean(span_ok)) if span_ok else 0.0, "age_accuracy": float(np.mean(age_ok)) if age_ok else 0.0, "age_f1": f1_score(y_age, p_age, average="binary", zero_division=0) if any(y_age) or any(p_age) else 1.0, "joint_exact": float(np.mean([a and b and c for a, b, c in zip(type_ok, span_ok, age_ok)])) if type_ok else 0.0, } def class_weights(rows, type2id: dict[str, int], power: float = 0.5) -> torch.Tensor: """Веса классов типа: (N / (K · n_c)) ** power, нормированные так, что средний вес по примерам = 1. power 0.5 (корень из обратной частоты) смягчает дисбаланс, не раздувая шум редчайших классов (у «авт. многоголосый» несколько строк): полная обратная частота дала бы им вес в сотни раз больше. Потолок веса — 4. """ counts = np.ones(len(type2id)) for r in rows: counts[type2id[r["translation_type"]]] += 1 w = (counts.sum() / (len(counts) * counts)) ** power w = w / (w * counts).sum() * counts.sum() w = np.minimum(w, 4.0) # потолок: классы из единиц строк не должны перетягивать обучение return torch.tensor(w, dtype=torch.float32) def train_model(name: str, cfg: dict, train_rows, val_rows, test_rows, type2id, device, epochs=None, seed=SEED, log_prefix="", lang2id=None, use_class_weights: bool = True): """Обучение с early stopping на val; возвращает (модель, токенизатор, отчёт).""" set_seed(seed) tokenizer = AutoTokenizer.from_pretrained(cfg["backbone"]) if not tokenizer.is_fast: # спан автора считается по offsets: нужен быстрый токенизатор raise RuntimeError(f"{cfg['backbone']}: нет быстрого токенизатора, offsets недоступны") model = DubTagger(cfg["backbone"], n_types=len(type2id), n_langs=len(lang2id or {})).to(device) collate = make_collate(tokenizer, type2id, lang2id=lang2id) train_loader = DataLoader(Rows(train_rows), batch_size=cfg["batch"], shuffle=True, collate_fn=collate) val_loader = DataLoader(Rows(val_rows), batch_size=128, collate_fn=collate) test_loader = DataLoader(Rows(test_rows), batch_size=128, collate_fn=collate) id2type = {i: t for t, i in type2id.items()} if use_class_weights: model.set_type_weights(class_weights(train_rows, type2id)) epochs = epochs or cfg["epochs"] steps = len(train_loader) * epochs optim = torch.optim.AdamW(model.parameters(), lr=cfg["lr"], weight_decay=0.01) sched = get_linear_schedule_with_warmup(optim, int(steps * 0.1), steps) best, best_state, history = -1.0, None, [] patience_left = 3 for epoch in range(epochs): model.train() losses = [] for enc, types, bio, ages, _, *rest in train_loader: enc_dev = {k: v.to(device) for k, v in enc.items()} with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device == "cuda"): loss, _, _, _ = model(**enc_dev, type_labels=types.to(device), bio_labels=bio.to(device), age_labels=ages.to(device), lang_labels=rest[0].to(device) if rest else None) optim.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optim.step() sched.step() losses.append(loss.item()) m = evaluate(model, val_loader, device, id2type) score = (m["type_macro_f1"] + m["span_exact"] + m["age_f1"]) / 3 history.append({"epoch": epoch, "loss": float(np.mean(losses)), **m}) print(f" {log_prefix}[{name}] epoch {epoch}: " f"loss={np.mean(losses):.4f} type_f1={m['type_macro_f1']:.3f} " f"span={m['span_exact']:.3f} age={m['age_accuracy']:.3f}" + (f" lang={m['lang_accuracy']:.3f} (n={m['lang_n']})" if "lang_accuracy" in m else "")) if score > best: best = score best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()} patience_left = 3 else: patience_left -= 1 if patience_left <= 0: break if best_state: model.load_state_dict(best_state) per_row: list = [] report = {"history": history, "test": evaluate(model, test_loader, device, id2type, per_row=per_row), "per_row": per_row} return model, tokenizer, report def main() -> None: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--models", nargs="+", default=DEFAULT_MODELS, choices=list(MODEL_CONFIGS)) ap.add_argument("--data", default="data/processed/dataset.json") ap.add_argument("--label-maps", default=None, help="по умолчанию label_maps.json рядом с --data") ap.add_argument("--folds", type=int, default=5, help="0 — без кросс-валидации (втрое-впятеро быстрее)") ap.add_argument("--epochs-override", type=int, default=None) ap.add_argument("--lr", type=float, default=None, help="заменяет lr профиля") ap.add_argument("--batch", type=int, default=None, help="заменяет batch профиля") ap.add_argument("--seed", type=int, default=SEED) ap.add_argument("--no-class-weights", action="store_true", help="потери типа без весов классов (по умолчанию: корень из обратной частоты)") ap.add_argument("--runs", action="store_true", help="каждая модель — отдельный запуск в --runs-dir; models/best, reports/final_* " "и summary.md не трогаются (выбор лучшей — tools/collect.py, пункт 7)") ap.add_argument("--runs-dir", default="models/runs") ap.add_argument("--tag", default="", help="суффикс имени запуска") args = ap.parse_args() device = "cuda" if torch.cuda.is_available() else "cpu" print(f"device: {device}" + (f" ({torch.cuda.get_device_name(0)})" if device == "cuda" else "")) rows = json.loads(Path(args.data).read_text(encoding="utf-8")) maps_path = Path(args.label_maps) if args.label_maps else Path(args.data).with_name("label_maps.json") label_maps = json.loads(maps_path.read_text(encoding="utf-8")) type2id = {t: i for i, t in enumerate(label_maps["types"])} lang2id = {l: i for i, l in enumerate(label_maps.get("langs") or [])} or None id2type = {i: t for t, i in type2id.items()} train_rows = [r for r in rows if r["split"] == "train"] test_rows = [r for r in rows if r["split"] == "test"] print(f"rows: train {len(train_rows)} / test {len(test_rows)}, " f"classes {len(type2id)}") from . import runs fp = runs.fingerprint(test_rows) Path("models").mkdir(exist_ok=True) summary: dict[str, dict] = {} for name in args.models: cfg = dict(MODEL_CONFIGS[name]) if args.lr: cfg["lr"] = args.lr if args.batch: cfg["batch"] = args.batch t0 = time.time() # --- 5-fold CV на train-части --------------------------------- fold_ids = sorted({r["fold"] for r in train_rows}) cv_metrics = [] if args.folds: print(f"\n=== {name} ({cfg['backbone']}): CV ===") for f in fold_ids[:args.folds]: tr = [r for r in train_rows if r["fold"] != f] va = [r for r in train_rows if r["fold"] == f] _, _, res = train_model(name, cfg, tr, va, va, type2id, device, epochs=args.epochs_override, seed=args.seed + f, log_prefix=f"cv{f} ", lang2id=lang2id, use_class_weights=not args.no_class_weights) cv_metrics.append(res["test"]) cv = ({k: float(np.mean([m[k] for m in cv_metrics])) for k in cv_metrics[0]} if cv_metrics else None) # --- финальная модель: train (90%) + val (10% групп) ------------ print(f"=== {name}: финальное обучение ===") groups = sorted({r["group_id"] for r in train_rows}) rng = random.Random(args.seed) rng.shuffle(groups) val_groups = set(groups[:max(1, len(groups) // 10)]) tr = [r for r in train_rows if r["group_id"] not in val_groups] va = [r for r in train_rows if r["group_id"] in val_groups] model, tokenizer, res = train_model( name, cfg, tr, va, test_rows, type2id, device, epochs=args.epochs_override, seed=args.seed, lang2id=lang2id, use_class_weights=not args.no_class_weights) res["cv_mean"] = cv res["train_seconds"] = round(time.time() - t0, 1) per_row = res.pop("per_row") summary[name] = res if args.runs: # отдельный запуск: ничего существующего не перезаписывается run_dir = Path(args.runs_dir) / f"{time.strftime('%Y%m%d_%H%M%S')}_{name}{args.tag}" save_model(model, tokenizer, run_dir) shutil.copy(maps_path, run_dir / "label_maps.json") info = {"name": name, **cfg, "max_len": MAX_LEN, "seed": args.seed, "folds": args.folds, "class_weights": not args.no_class_weights, "created": time.strftime("%Y-%m-%d %H:%M:%S"), "fingerprint": fp, "rows": {"train": len(train_rows), "test": len(test_rows)}, "final": {k: v for k, v in res.items() if k != "history"}} (run_dir / "run.json").write_text(json.dumps(info, ensure_ascii=False, indent=1), encoding="utf-8") (run_dir / "train_config.json").write_text(json.dumps(info, ensure_ascii=False, indent=1), encoding="utf-8") runs.store_eval(run_dir, fp, {"metrics": res["test"], "per_row": per_row, "n_unscorable": 0}) _plots(name, res, model, tokenizer, test_rows, type2id, id2type, device, lang2id, out=run_dir) print(f"запуск сохранён: {run_dir}") del model if device == "cuda": torch.cuda.empty_cache() continue model_dir = Path("models") / name save_model(model, tokenizer, model_dir) shutil.copy(maps_path, model_dir) (model_dir / "train_config.json").write_text(json.dumps( {"name": name, **cfg, "max_len": MAX_LEN, "seed": args.seed, "class_weights": not args.no_class_weights, "final": {k: v for k, v in res.items() if k != "history"}}, ensure_ascii=False, indent=1), encoding="utf-8") runs.store_eval(model_dir, fp, {"metrics": res["test"], "per_row": per_row, "n_unscorable": 0}) Path(f"reports/final_{name}.json").write_text( json.dumps(res, ensure_ascii=False, indent=1), encoding="utf-8") _plots(name, res, model, tokenizer, test_rows, type2id, id2type, device, lang2id) if args.runs: return # --- выбор лучшей и итоговый отчёт --------------------------------- best_name = max(summary, key=lambda n: (summary[n]["test"]["joint_exact"], summary[n]["test"]["type_macro_f1"])) if Path("models/best").exists(): shutil.rmtree("models/best") shutil.copytree(Path("models") / best_name, "models/best") lines = ["# Сравнение моделей", "", "| модель | type acc | type macro-F1 | span exact | age acc | joint exact | время, с |", "|---|---|---|---|---|---|---|"] for n, r in summary.items(): t = r["test"] lines.append(f"| {n} | {t['type_accuracy']:.3f} | {t['type_macro_f1']:.3f} " f"| {t['span_exact']:.3f} | {t['age_accuracy']:.3f} " f"| {t['joint_exact']:.3f} | {r['train_seconds']} |") lines += ["", f"Лучшая модель: **{best_name}** -> `models/best`"] Path("reports/summary.md").write_text("\n".join(lines), encoding="utf-8") print("\n" + "\n".join(lines)) def _plots(name, res, model, tokenizer, test_rows, type2id, id2type, device, lang2id=None, out="reports"): """Кривые обучения и confusion matrix на test (в каталог out).""" out = Path(out) out.mkdir(parents=True, exist_ok=True) import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt hist = res["history"] if hist: fig, ax = plt.subplots(1, 2, figsize=(10, 4)) ax[0].plot([h["epoch"] for h in hist], [h["loss"] for h in hist]) ax[0].set_title("train loss") ax[0].set_xlabel("epoch") ax[1].plot([h["epoch"] for h in hist], [h["type_macro_f1"] for h in hist], label="type macro-F1") ax[1].plot([h["epoch"] for h in hist], [h["span_exact"] for h in hist], label="span exact") ax[1].set_xlabel("epoch") ax[1].legend() fig.tight_layout() fig.savefig(out / f"history_{name}.png", dpi=120) plt.close(fig) loader = DataLoader(Rows(test_rows), batch_size=128, collate_fn=make_collate(tokenizer, type2id, lang2id=lang2id)) y, p = [], [] with torch.no_grad(): for enc, types, _, _, _, *_ in loader: enc_dev = {k: v.to(device) for k, v in enc.items()} tl, _, _ = model.predict(**enc_dev) y += types.tolist() p += tl.argmax(-1).tolist() y = [id2type[v] for v in y] p = [id2type[v] for v in p] labels = sorted(set(y) | set(p), key=lambda t: -y.count(t)) cm = confusion_matrix(y, p, labels=labels) fig, ax = plt.subplots(figsize=(9, 7)) im = ax.imshow(cm, cmap="Blues") ax.set_xticks(range(len(labels)), labels, rotation=45, ha="right", fontsize=8) ax.set_yticks(range(len(labels)), labels, fontsize=8) for i in range(len(labels)): for j in range(len(labels)): if cm[i, j]: ax.text(j, i, cm[i, j], ha="center", va="center", fontsize=7, color="white" if cm[i, j] > cm.max() / 2 else "black") fig.colorbar(im) ax.set_title(f"confusion matrix: {name} (test)") ax.set_xlabel("предсказано") ax.set_ylabel("истина") fig.tight_layout() fig.savefig(out / f"confusion_{name}.png", dpi=120) plt.close(fig) if __name__ == "__main__": main()