DenozordecandClaude Sonnet 5.5
729eeac045
docker / build (push) Successful in 1m49s
Kinozal 2160145: после переобучения модель выбирала первую строку вкладки «Релиз» («Memoirs.of.a.Geisha.2005.2160p.UHD.Blu-ray.Remux…-CiNEPHiLES», оценка 0.96), и имя релиза попадало в ответ как «студия». Прежние веса строку не выбирали. - html_extract.RE_SCENE: без пробелов и ≥ 4 точек — дорожкой не бывает; LineSegmenter.select_lines отбрасывает такие строки при любой оценке модели. - build_examples: scene-имена — трудные негативы, берутся все (раньше прореживались до 15 %); эффект — при следующем переобучении. - На реальных данных страницы (вкладки Техданные/Интересно/Награды/Релиз): 1 дорожка «проф. дубляж», без студии. Тесты: 288 ok. Co-Authored-By: Claude Sonnet 5.5 <[email protected]>
299 lines
14 KiB
Python
299 lines
14 KiB
Python
"""Обучаемый сегментатор: какие строки страницы описывают дорожки озвучки.
|
||
|
||
Главный путь выделения дорожек (правила из 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()
|