grabber-ai/src/segmenter.py

"""Обучаемый сегментатор: какие строки страницы описывают дорожки озвучки.

Главный путь выделения дорожек (правила из 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()


from .manual import REPEAT, manual_ids  # noqa: E402


def is_test_release(release_id: str) -> bool:
    """Детерминированное разбиение страниц: ~20% релизов уходят в тест.

    Релизы, добавленные вручную по ссылке, всегда в train.
    """
    if str(release_id) in manual_ids():
        return False
    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
        item = (rid, lines, label_page(lines, raws_by_rel.get(rid, [])))
        page_items.extend([item] * (REPEAT if rid in manual_ids() else 1))

    # все 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="longest_first",
                           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="longest_first",
                   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()