DenozordecandClaude Sonnet 5.5
05537ac9e9
docker / build (push) Successful in 1m40s
- harvest_core: дубляж — тип, а не часть имени студии; split_multi_rows; audio_only_rows; kinozal сопоставляет все аудиодорожки - prepare_data: --max-dup, отсев остатков многотиповых строк - train: веса классов (корень из обратной частоты, потолок 4), --no-class-weights - исправлены регулярки с \b, записанные как символ BS (predict.py, audio18.py) - PARSE_VERSION 3: кэш разбора сбрасывается - docs/dataset.md, docs/api-reference.md обновлены Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
248 lines
12 KiB
Python
248 lines
12 KiB
Python
"""Инференс модели: CLI и Python API.
|
|
|
|
CLI:
|
|
python -m src.predict --text "Перевод: Профессиональный (многоголосый закадровый) | HDRezka Studio 18+"
|
|
python -m src.predict --html page.html
|
|
python -m src.predict --file page.txt [--out result.json] [--model models/best]
|
|
|
|
Python API:
|
|
from src.predict import DubPredictor
|
|
model = DubPredictor() # models/best по умолчанию
|
|
model.predict_text("Авторский (одноголосый закадровый) - Ю. Сербин")
|
|
model.predict_html(open("page.html", encoding="utf-8").read())
|
|
|
|
Формат результата (JSON):
|
|
{"segment": "...", "translation_type": "...", "type_confidence": 0.99,
|
|
"studio_span": "..." | null, "studio": "..." | null, "studio_confidence": 0.97,
|
|
"is_18plus": false, "age_confidence": 0.99}
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import re
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
from transformers import AutoTokenizer
|
|
|
|
from .html_extract import (RE_LINE_PREFIX, extract_segments, html_to_text,
|
|
read_input)
|
|
from .model import load_model, predict_with_lang
|
|
from .studio_norm import canonical
|
|
from .spans import MAX_LEN, decode_span
|
|
from .trackinfo import lang_of_text, model_text, role_of
|
|
|
|
|
|
# мусор, который модель иногда принимает за автора: раскладка каналов, кодеки
|
|
RE_NOT_AUTHOR = re.compile(
|
|
r"(?i)^[\s,/:+.\d]*(?:[LRC]|Ls|Rs|LFE|l|r|Front|Side|ch)?(?:[\s,/:+.\d]*(?:[LRC]|Ls|Rs|LFE|l|r|Front|Side|ch))*[\s,/:+.\d]*$"
|
|
r"|dolby|digital|audio|kbps|khz|\bAC-?3\b|\bDTS\b|\bAAC\b|\bOpus\b")
|
|
|
|
|
|
def refine_span(text: str, span: tuple[int, int] | None) -> tuple[int, int] | None:
|
|
"""Постобработка спана автора: до границ слов; мусор и 1 символ -> None."""
|
|
if span is None:
|
|
return None
|
|
a, z = span
|
|
while a > 0 and text[a - 1].isalnum() and text[a].isalnum():
|
|
a -= 1
|
|
while z < len(text) and text[z].isalnum() and text[z - 1].isalnum():
|
|
z += 1
|
|
piece = text[a:z].strip()
|
|
if len(piece) < 2 or RE_NOT_AUTHOR.search(piece):
|
|
return None
|
|
return a, z
|
|
|
|
|
|
|
|
_RE_VOICES = re.compile(r"(?i)(много|двух|одно)голос")
|
|
_RE_DUB_WORD = re.compile(r"(?i)дубл")
|
|
_PREFIXES = (("проф.", re.compile(r"(?i)профессиональн")), ("люб.", re.compile(r"(?i)любительск")),
|
|
("авт.", re.compile(r"(?i)авторск")))
|
|
|
|
|
|
def guard_type(ttype: str, text: str, span_text: str | None, types: list[str]) -> str:
|
|
"""Дубляж из-за имени студии («Дубляжная», «Продубляж»): если вне имени о дубляже ничего нет,
|
|
а число голосов названо прямо («Любительский (многоголосый закадровый) …») — берём класс по тексту."""
|
|
if "дубляж" not in ttype:
|
|
return ttype
|
|
rest = text.replace(span_text, " ") if span_text else text
|
|
voices = _RE_VOICES.search(rest)
|
|
if _RE_DUB_WORD.search(rest) or not voices:
|
|
return ttype
|
|
prefix = next((p for p, rx in _PREFIXES if rx.search(rest)), "")
|
|
cand = f"{prefix} {voices.group(1).lower()}голосый".strip()
|
|
return cand if cand in types else ttype
|
|
|
|
|
|
class DubPredictor:
|
|
"""Загруженная модель для предсказания типа озвучки и автора."""
|
|
|
|
def __init__(self, model_dir: str | Path = "models/best", device: str | None = None,
|
|
segmenter_dir: str | Path = "models/segmenter"):
|
|
self.model_dir = Path(model_dir)
|
|
if device is None:
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
self.device = device
|
|
maps = json.loads((self.model_dir / "label_maps.json").read_text(encoding="utf-8"))
|
|
self.types = maps["types"]
|
|
self.langs = maps.get("langs") or [] # пусто: у модели нет языковой головы
|
|
# модели, обученные до появления строки «Аудио» во входе, её не ждут
|
|
self.use_audio_context = bool(maps.get("audio_context"))
|
|
self.tokenizer = AutoTokenizer.from_pretrained(self.model_dir)
|
|
self.model = load_model(self.model_dir, device)
|
|
# обученный сегментатор страниц (главный путь); нет модели -> правила
|
|
self.segmenter = None
|
|
seg_dir = Path(segmenter_dir)
|
|
if seg_dir.exists():
|
|
from .segmenter import LineSegmenter
|
|
self.segmenter = LineSegmenter(seg_dir, device)
|
|
|
|
# ------------------------------------------------------------------
|
|
@torch.no_grad()
|
|
def _predict_batch(self, texts: list[str], infos: list[dict | None] | None = None) -> list[dict]:
|
|
"""texts — тексты дорожек; infos — разобранные строки «Аудио» (или None) того же порядка.
|
|
|
|
Вход модели = текст + «[RUS AC3 6ch 448kbps 18+]», как при обучении (prepare_data).
|
|
"""
|
|
infos = infos or [None] * len(texts)
|
|
inputs = [model_text(t, i if self.use_audio_context else None) for t, i in zip(texts, infos)]
|
|
enc = self.tokenizer(inputs, truncation=True, max_length=MAX_LEN,
|
|
padding=True, return_tensors="pt",
|
|
return_offsets_mapping=True)
|
|
offsets = enc.pop("offset_mapping").tolist()
|
|
enc = {k: v.to(self.device) for k, v in enc.items()}
|
|
lang_prob = None
|
|
if self.langs and self.model.lang_head is not None:
|
|
type_logits, bio_logits, age_logits, lang_logits = predict_with_lang(self.model, **enc)
|
|
lang_prob = torch.softmax(lang_logits, dim=-1).cpu().numpy()
|
|
else:
|
|
type_logits, bio_logits, age_logits = self.model.predict(**enc)
|
|
type_prob = torch.softmax(type_logits, dim=-1).cpu().numpy()
|
|
age_prob = torch.softmax(age_logits, dim=-1).cpu().numpy()
|
|
bio_prob = torch.softmax(bio_logits, dim=-1).cpu().numpy()
|
|
bio_pred = bio_prob.argmax(-1)
|
|
|
|
results = []
|
|
for i, text in enumerate(inputs):
|
|
info = infos[i]
|
|
seg = texts[i]
|
|
tid = int(type_prob[i].argmax())
|
|
aid = int(age_prob[i].argmax())
|
|
span = refine_span(text, decode_span(
|
|
bio_pred[i], offsets[i], enc["attention_mask"][i].tolist()))
|
|
span_text = text[span[0]:span[1]].strip() if span else None
|
|
ttype = guard_type(self.types[tid], text, span_text, self.types)
|
|
span_conf = None
|
|
if span is not None:
|
|
toks = [t for t, (a, z) in enumerate(offsets[i])
|
|
if z > a and a < span[1] and z > span[0]]
|
|
confs = [float(bio_prob[i, t, bio_pred[i, t]]) for t in toks]
|
|
span_conf = float(np.mean(confs)) if confs else None
|
|
lang, lang_conf, lang_src = "", None, None
|
|
if info and info["lang"]: # язык, прямо названный в строке «Аудио», надёжнее нейросети
|
|
lang, lang_src = info["lang"], "audio"
|
|
elif lang_prob is not None:
|
|
j = int(lang_prob[i].argmax())
|
|
if lang_prob[i, j] >= 0.5:
|
|
lang, lang_conf, lang_src = self.langs[j], round(float(lang_prob[i, j]), 4), "model"
|
|
if not lang:
|
|
lang = lang_of_text(seg)
|
|
lang_src = "text" if lang else None
|
|
results.append({
|
|
"segment": seg,
|
|
"translation_type": ttype,
|
|
"type_confidence": round(float(type_prob[i, tid]), 4),
|
|
"studio_span": span_text,
|
|
"studio": canonical(span_text) if span_text else None,
|
|
"studio_confidence": round(span_conf, 4) if span_conf else None,
|
|
"is_18plus": bool(aid == 1) or bool(info and info["adult"]), # 18+ в строке «Аудио» — прямое свидетельство
|
|
"age_confidence": round(float(age_prob[i, aid]), 4),
|
|
"lang": lang or None,
|
|
"lang_confidence": lang_conf,
|
|
"lang_source": lang_src,
|
|
"role": role_of(info, seg),
|
|
"official": bool(info and info["official"]),
|
|
"audio": ({"codec": info["codec"] or None, "channels": int(info["channels"]) if info["channels"] else None,
|
|
"bitrate_kbps": int(info["kbps"]) if info["kbps"] else None}
|
|
if info else None),
|
|
"_audio18": info["adult"] if info else None,
|
|
})
|
|
return results
|
|
|
|
# ------------------------------------------------------------------
|
|
def predict_text(self, text: str) -> dict:
|
|
"""Текст одной дорожки (можно с префиксом «Перевод:») -> предсказание."""
|
|
clean = re.sub(r"\s+", " ",
|
|
RE_LINE_PREFIX.sub("", (text or "").strip())).strip()
|
|
if not clean:
|
|
raise ValueError("пустой текст дорожки")
|
|
return self._predict_batch([clean])[0]
|
|
|
|
def predict_page(self, page_text: str) -> list[dict]:
|
|
"""Текст страницы релиза -> предсказания для всех найденных дорожек."""
|
|
segments = extract_segments(page_text)
|
|
return self._predict_batch(segments) if segments else []
|
|
|
|
def predict_html(self, html: str) -> list[dict]:
|
|
"""HTML-страница релиза -> предсказания для всех найденных дорожек."""
|
|
return self.predict_page(html_to_text(html))
|
|
|
|
def predict_file(self, path: str | Path) -> list[dict]:
|
|
"""Файл (HTML или текстовый дамп страницы) -> все дорожки."""
|
|
content = read_input(path)
|
|
if looks_like_html(content):
|
|
return self.predict_html(content)
|
|
return self.predict_page(content)
|
|
|
|
|
|
def looks_like_html(content: str) -> bool:
|
|
return bool(re.search(r"<(?i:html|body|div|br|table)", content[:2000]))
|
|
|
|
|
|
def main() -> None:
|
|
ap = argparse.ArgumentParser(
|
|
description="Предсказание типа озвучки и автора (см. docstring модуля)",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
ap.add_argument("--model", default="models/best", help="каталог модели")
|
|
src = ap.add_mutually_exclusive_group(required=True)
|
|
src.add_argument("--text", help="текст одной дорожки")
|
|
src.add_argument("--html", help="путь к HTML-файлу страницы")
|
|
src.add_argument("--file", help="путь к текстовому дампу страницы")
|
|
src.add_argument("--url", help="ссылка на раздачу (rutracker / kinozal / nnm)")
|
|
ap.add_argument("--cookie", help="Cookie для --url (нужен kinozal)")
|
|
ap.add_argument("--flaresolverr", help="адрес FlareSolverr (обход Cloudflare), напр. http://127.0.0.1:8191")
|
|
ap.add_argument("--out", help="записать JSON в файл вместо stdout")
|
|
args = ap.parse_args()
|
|
|
|
predictor = DubPredictor(args.model)
|
|
if args.text:
|
|
payload = {"tracks": [predictor.predict_text(args.text)]}
|
|
elif args.url:
|
|
from .fetch import fetch_html
|
|
from .release import predict_release
|
|
html, tracker = fetch_html(args.url, cookie=args.cookie,
|
|
flaresolverr=args.flaresolverr)
|
|
payload = {**predict_release(predictor, html=html), "tracker": tracker}
|
|
else:
|
|
from .release import predict_release
|
|
content = read_input(args.html or args.file)
|
|
if looks_like_html(content):
|
|
payload = predict_release(predictor, html=content)
|
|
else:
|
|
payload = predict_release(predictor, text=content)
|
|
|
|
text_json = json.dumps(payload, ensure_ascii=False, indent=2)
|
|
if args.out:
|
|
Path(args.out).write_text(text_json, encoding="utf-8")
|
|
print(f"written: {args.out}")
|
|
else:
|
|
print(text_json)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|