Обученный сегментатор страниц (главный путь), правила — запасной; фикстура rt-topic-6870900

Co-Authored-By: Claude Sonnet 5.5 <[email protected]>

src/segmenter.py

@@ -0,0 +1,215 @@
+"""Обучаемый сегментатор: какие строки страницы описывают дорожки озвучки.
+
+Главный путь выделения дорожек (правила из html_extract — запасной).
+Модель — rubert-tiny2 (бинарная классификация строки). Вход — строка с
+контекстом: «пред. строка [SEP] строка [SEP] след. строка».
+
+    python -m src.segmenter --train          # обучение -> models/segmenter
+    python -m src.segmenter --eval           # метрики на отложенных страницах
+
+Python API:
+    seg = LineSegmenter("models/segmenter")
+    seg.select_lines(page_text) -> list[str]   # строки-дорожки (префикс не снят)
+
+Разметка: строка страницы положительна, если содержит `raw` строки датасета
+этого релиза (или наоборот). Положительные примеры добавляются синтетически:
+`raw` всех train-строк датасета (включая nnm/kinozal, у которых нет страниц)
+вставляются в контекст реальных строк страниц. Нарезка ВНУТРИ строки
+(«Дублированный, многоголосый») в данных не размечена и остаётся на правилах.
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import hashlib
+import json
+import random
+import re
+from pathlib import Path
+
+import numpy as np
+import torch
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+from .html_extract import read_input
+
+BACKBONE = "cointegrated/rubert-tiny2"
+MAX_LEN = 160
+LINE_CHARS = 220
+CTX_CHARS = 70
+THRESHOLD = 0.5
+SEED = 7
+
+
+def norm(s: str) -> str:
+    return re.sub(r"\s+", " ", s or "").strip()
+
+
+def is_test_release(release_id: str) -> bool:
+    """Детерминированное разбиение страниц: ~20% релизов уходят в тест."""
+    return int(hashlib.md5(str(release_id).encode()).hexdigest(), 16) % 5 == 0
+
+
+def make_input(prev: str, line: str, nxt: str) -> tuple[str, str]:
+    """(текст-пара для токенизатора): строка + контекст."""
+    return norm(line)[:LINE_CHARS], f"{norm(prev)[-CTX_CHARS:]} || {norm(nxt)[:CTX_CHARS]}"
+
+
+def page_lines(text: str) -> list[str]:
+    return [norm(l) for l in text.splitlines() if norm(l)]
+
+
+# ---------------------------------------------------------------- данные
+def load_pages(pages_dir: Path = Path("data/pages")) -> dict[str, list[str]]:
+    return {p.stem: page_lines(read_input(p)) for p in sorted(pages_dir.glob("*.txt"))}
+
+
+def load_raws(csv_path: Path = Path("data/labels.csv")) -> list[dict]:
+    with open(csv_path, encoding="utf-8-sig") as f:
+        return list(csv.DictReader(f))
+
+
+def label_page(lines: list[str], raws: list[str]) -> list[int]:
+    """1, если строка содержит raw (или raw содержит строку ≥ 15 символов)."""
+    out = []
+    for line in lines:
+        hit = any(r and (r in line or (len(line) >= 15 and line in r)) for r in raws)
+        out.append(int(hit))
+    return out
+
+
+def build_examples(split: str, rng: random.Random, synth_per_page: int = 6) -> list[tuple]:
+    """Список (line, ctx, label). split: 'train' | 'test'."""
+    pages = load_pages()
+    rows = load_raws()
+    raws_by_rel: dict[str, list[str]] = {}
+    for r in rows:
+        raws_by_rel.setdefault(r["release_id"], []).append(norm(r["raw"]))
+
+    want_test = split == "test"
+    page_items = []  # (release_id, lines, labels)
+    for rid, lines in pages.items():
+        if is_test_release(rid) != want_test:
+            continue
+        page_items.append((rid, lines, label_page(lines, raws_by_rel.get(rid, []))))
+
+    # все raw train-релизов (без тестовых страниц) — материал для синтетики
+    test_pos = {l for rid, lines in pages.items() if is_test_release(rid)
+                for l, y in zip(lines, label_page(lines, raws_by_rel.get(rid, []))) if y}
+    pool = sorted({norm(r["raw"]) for r in rows
+                   if r["raw"] and not is_test_release(r["release_id"])
+                   and norm(r["raw"]) not in test_pos})
+
+    ex: list[tuple] = []
+    for rid, lines, labels in page_items:
+        for i, (line, y) in enumerate(zip(lines, labels)):
+            prev = lines[i - 1] if i else ""
+            nxt = lines[i + 1] if i + 1 < len(lines) else ""
+            if not y and not HINT.search(line) and rng.random() > 0.15:
+                continue  # простые негативы прореживаем, трудные (с намёками) берём все
+            ex.append((*make_input(prev, line, nxt), y))
+        if not want_test:
+            negs = [l for l, y in zip(lines, labels) if not y]
+            for _ in range(synth_per_page if negs else 0):
+                raw = rng.choice(pool)
+                p, n = rng.choice(negs), rng.choice(negs)
+                ex.append((*make_input(p, raw, n), 1))
+    return ex
+
+
+# слова, делающие негативную строку «трудной»: озвучка/аудио/субтитры/кодеки
+HINT = re.compile(r"(?i)(перевод|озвуч|аудио|звук|дубл|голос|субтитр|kbps|kHz|AC3|DTS|MVO|DVO|AVO)")
+
+
+# ---------------------------------------------------------------- модель
+class LineSegmenter:
+    def __init__(self, model_dir: str | Path = "models/segmenter", device: str | None = None):
+        self.dir = Path(model_dir)
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.tok = AutoTokenizer.from_pretrained(self.dir)
+        self.model = AutoModelForSequenceClassification.from_pretrained(self.dir).to(self.device).eval()
+
+    @torch.no_grad()
+    def score(self, lines: list[str]) -> np.ndarray:
+        """Вероятность «строка описывает дорожку» для каждой строки."""
+        pairs = [make_input(lines[i - 1] if i else "", l, lines[i + 1] if i + 1 < len(lines) else "")
+                 for i, l in enumerate(lines)]
+        out = []
+        for s in range(0, len(pairs), 128):
+            chunk = pairs[s:s + 128]
+            enc = self.tok([a for a, _ in chunk], [b for _, b in chunk], truncation="only_second",
+                           max_length=MAX_LEN, padding=True, return_tensors="pt").to(self.device)
+            out.append(torch.softmax(self.model(**enc).logits, -1)[:, 1].cpu().numpy())
+        return np.concatenate(out) if out else np.zeros(0)
+
+    def select_lines(self, page_text: str, threshold: float = THRESHOLD) -> list[str]:
+        lines = page_lines(page_text)
+        if not lines:
+            return []
+        probs = self.score(lines)
+        return [l for l, p in zip(lines, probs) if p >= threshold]
+
+
+# ---------------------------------------------------------------- обучение
+def train(out_dir: str = "models/segmenter", epochs: int = 4, batch: int = 32, lr: float = 5e-5):
+    rng = random.Random(SEED)
+    torch.manual_seed(SEED)
+    tr = build_examples("train", rng)
+    te = build_examples("test", rng)
+    print(f"train {len(tr)} (pos {sum(y for *_, y in tr)}), test {len(te)} (pos {sum(y for *_, y in te)})")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    tok = AutoTokenizer.from_pretrained(BACKBONE)
+    model = AutoModelForSequenceClassification.from_pretrained(BACKBONE, num_labels=2).to(device)
+    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
+    steps = epochs * ((len(tr) + batch - 1) // batch)
+    sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=lr, total_steps=steps, pct_start=0.1)
+
+    def enc(b):
+        return tok([x[0] for x in b], [x[1] for x in b], truncation="only_second",
+                   max_length=MAX_LEN, padding=True, return_tensors="pt").to(device)
+
+    def evaluate(data):
+        model.eval()
+        pred = []
+        with torch.no_grad():
+            for s in range(0, len(data), 128):
+                pred += model(**enc(data[s:s + 128])).logits.argmax(-1).tolist()
+        y = [x[2] for x in data]
+        tp = sum(p and t for p, t in zip(pred, y))
+        prec = tp / max(sum(pred), 1)
+        rec = tp / max(sum(y), 1)
+        model.train()
+        return prec, rec
+
+    model.train()
+    for ep in range(epochs):
+        rng.shuffle(tr)
+        for s in range(0, len(tr), batch):
+            b = tr[s:s + batch]
+            loss = model(**enc(b), labels=torch.tensor([x[2] for x in b], device=device)).loss
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            opt.step()
+            sched.step()
+            opt.zero_grad()
+        p, r = evaluate(te)
+        print(f"epoch {ep + 1}: test line precision {p:.3f} recall {r:.3f}")
+    Path(out_dir).mkdir(parents=True, exist_ok=True)
+    model.save_pretrained(out_dir)
+    tok.save_pretrained(out_dir)
+    return p, r
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--train", action="store_true")
+    ap.add_argument("--out", default="models/segmenter")
+    ap.add_argument("--epochs", type=int, default=4)
+    args = ap.parse_args()
+    if args.train:
+        train(args.out, args.epochs)
+
+
+if __name__ == "__main__":
+    main()