"""Инференс модели: 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()