190 lines
8.6 KiB
Python
190 lines
8.6 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
|
|
from .studio_norm import canonical
|
|
from .train import MAX_LEN, decode_span
|
|
|
|
|
|
# мусор, который модель иногда принимает за автора: раскладка каналов, кодеки
|
|
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|AC-?3|DTS|AAC|Opus")
|
|
|
|
|
|
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
|
|
|
|
|
|
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.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]) -> list[dict]:
|
|
enc = self.tokenizer(texts, 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()}
|
|
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(texts):
|
|
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
|
|
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
|
|
results.append({
|
|
"segment": text,
|
|
"translation_type": self.types[tid],
|
|
"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),
|
|
"age_confidence": round(float(age_prob[i, aid]), 4),
|
|
})
|
|
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()
|