"""Обучаемый сегментатор: какие строки страницы описывают дорожки озвучки. Главный путь выделения дорожек (правила из 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` (или `audio_raw`) строки датасета этого релиза (или наоборот); неразмеченные кодек-строки русских дорожок в обучение не идут (см. `label_page`). Текст kinozal собирается из data/raw/kinozal/*.json (`kinozal_text`). Положительные примеры добавляются синтетически: `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 RE_SCENE, 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)] # ---------------------------------------------------------------- данные RAW_KINOZAL = Path("data/raw/kinozal") def kinozal_text(release_id: str, raw_dir: Path = RAW_KINOZAL) -> str | None: """Текст страницы kinozal так, как его видит сервис: заголовок + «Техданные» («Перевод: …») + «Релиз» («Аудио N: … (MVO, LostFilm)»). В data/pages/kz-*.txt нет ни заголовка, ни «Техданных», а метки kinozal наполовину взяты именно оттуда, поэтому для kinozal берётся raw JSON. """ path = raw_dir / f"{str(release_id).removeprefix('kz-')}.json" if not path.exists(): return None j = json.loads(path.read_text(encoding="utf-8")) return "\n".join(x for x in (j.get("h1"), j.get("tech_tab"), j.get("release_tab")) if x) def page_text(path: Path) -> str: """Текст сохранённой страницы релиза (kinozal — из raw JSON, если он есть).""" if path.stem.startswith("kz-"): text = kinozal_text(path.stem, path.parent.parent / "raw" / "kinozal") if text: return text return read_input(path) def load_pages(pages_dir: Path = Path("data/pages")) -> dict[str, list[str]]: return {p.stem: page_lines(page_text(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)) # кодек-строка русской/украинской дорожки: «Аудио 2: AC3 … Русский (MVO, X)», # «Аудио Rus: AC3 … | LostFilm», «Аудиодорожка 1: … | Russian | TVShows»; # сводка «Аудио: Русский, английский (AC3 …)» без номера — не дорожка RE_AUDIO_TRACK = re.compile( r"(?i)^\s*(?:(?:Аудио|Audio)\s*(?:rus|ukr)\s*[:.]" r"|(?:Аудиодорожка\s*\d*|(?:Аудио|Audio)\s*(?:№|#)?\s*\d+)\s*[:.]" r".*(?:русск|russian|\brus\b|украин|ukrain|\bukr\b))") RE_TYPE_WORD = re.compile( r"(?i)\b(?:MVO|DVO|AVO|DUB|VO)\b|дубл|многоголос|двухголос|одноголос|закадр|авторск") # технические слова кодек-строки: всё, что после их удаления осталось, — имя студии/автора RE_TECH_WORD = re.compile( r"(?i)^(?:аудио\w*|audio|e?-?ac-?3|eac3|ac|dts|hd|ma|mpeg|mp3|aac|lc|opus|flac|pcm|lpcm|truehd" r"|atmos|joc|dolby|digital|surround|stereo|mono|cbr|vbr|avg|constant|core|bit|ch|сh|channels?" r"|каналов|канала|kbps|kb|кбит|бит|сек|khz|кгц|hz|lfe|dn|db|rus|ukr|eng|russian|русск\w*" r"|украин\w*|english|англ\w*|l|c|r|sl|sr|ls|rs|lr|cbr|a)$") def _informative(line: str) -> bool: """Кодек-строка называет тип озвучки или студию/автора, а не только параметры. «Аудио 1: AC3, 6 ch, 384 kbps - Русский» — нет; «… | LostFilm», «… - Русский, Юрий Живов», «… (MVO)» — да. """ if RE_TYPE_WORD.search(line): return True body = re.sub(r"(?i)^\s*(?:Аудио\w*|Audio)\s*(?:№|#)?\s*(?:\d+|rus|ukr)?\s*[:.]", "", line) words = [w for w in re.findall(r"[^\W\d_]+", body) if not RE_TECH_WORD.match(w)] return any(len(w) >= 3 for w in words) def label_page(lines: list[str], raws: list[str]) -> list[int]: """1 — строка содержит raw (или raw содержит строку ≥ 15 символов); -1 — неразмеченная кодек-строка русской дорожки с типом или студией: в обучение не идёт; пустая кодек-строка («AC3, 6 ch, 384 kbps - Русский») — 0. Кодек-строка «Аудио N: … (MVO, LostFilm)» — дорожка, но у rutracker `raw` — строка «Перевод N:», а кодек-строка в разметку не попадает. Как негатив она учила модель отбрасывать такие строки (на kinozal студии только в них), поэтому неразмеченные кодек-строки пропускаются. """ out = [] for line in lines: hit = any(r and (r in line or (len(line) >= 15 and line in r)) for r in raws) skip = RE_AUDIO_TRACK.match(line) and _informative(line) out.append(1 if hit else (-1 if skip else 0)) 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: # audio_raw — кодек-строка той же дорожки («AC3 … Русский (MVO, X)»), тоже позитив for text in (r["raw"], r.get("audio_raw")): if norm(text): raws_by_rel.setdefault(r["release_id"], []).append(norm(text)) 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 == 1} 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)): if y < 0: continue # неразмеченная кодек-строка: ни позитив, ни негатив 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 not RE_SCENE.match(line) and rng.random() > 0.15: continue # простые негативы прореживаем, трудные (с намёками, scene-имена релизов) берём все ex.append((*make_input(prev, line, nxt), y)) if not want_test: negs = [l for l, y in zip(lines, labels) if y == 0] 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 and not RE_SCENE.match(l)] # ---------------------------------------------------------------- обучение 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()