168 lines
9.4 KiB
Python
168 lines
9.4 KiB
Python
"""Юнит-тесты DUB-AI (запуск: .venv/Scripts/python tests/run_tests.py).
|
||
|
||
Проверяют ключевые части без обращения к GPU:
|
||
* поиск спана автора (точно / алиас / инициалы / кавычки / разделители);
|
||
* канонизация имён студий;
|
||
* снятие префикса строки дорожки;
|
||
* сегментацию страниц на дорожки;
|
||
* HTML -> текст;
|
||
* BIO-выравнивание спана по offsets токенизатора.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||
|
||
from src.html_extract import (RE_LINE_PREFIX, _split_entries, # noqa: E402
|
||
extract_segments, html_to_text)
|
||
from src.studio_norm import canonical, find_studio_span # noqa: E402
|
||
|
||
PASS = 0
|
||
FAIL = 0
|
||
|
||
|
||
def check(name: str, cond: bool, extra: str = "") -> None:
|
||
global PASS, FAIL
|
||
if cond:
|
||
PASS += 1
|
||
print(f" ok {name}")
|
||
else:
|
||
FAIL += 1
|
||
print(f"FAIL {name} {extra}")
|
||
|
||
|
||
# ---------------------------------------------------------------- спаны
|
||
print("== find_studio_span ==")
|
||
txt = "Профессиональный (многоголосый закадровый) | HDRezka Studio 18+"
|
||
s, e = find_studio_span(txt, "HDrezka")
|
||
check("HDrezka найден", s is not None and "HDRezka" in txt[s:e], f"span={(s, e)}")
|
||
|
||
s, e = find_studio_span("Авторский (одноголосый закадровый) - Ю. Сербин", "Юрий Сербин")
|
||
check("инициалы Ю. Сербин", "Сербин" in "Авторский (одноголосый закадровый) - Ю. Сербин"[s:e] if s is not None else False)
|
||
|
||
s, e = find_studio_span("Авторский (одноголосый закадровый) Гаврилов", "Андрей Гаврилов")
|
||
check("фамилия Гаврилов", s is not None and e is not None and s > 0)
|
||
|
||
s, e = find_studio_span('Профессиональный (многоголосый закадровый) Студия «Кравец», LostFilm',
|
||
"Студия Кравец, LostFilm")
|
||
check("кавычки в тексте", s is not None)
|
||
|
||
s, e = find_studio_span("Профессиональный (многоголосый закадровый) - SDI Media | TV 1000",
|
||
"SDI Media TV 1000")
|
||
check("разделитель |", s is not None)
|
||
|
||
check("нет автора -> None", find_studio_span("Профессиональный (дублированный)", "") is None)
|
||
check("нет вхождения -> None", find_studio_span("Профессиональный (дублированный)", "HDrezka") is None)
|
||
|
||
# ---------------------------------------------------------------- канонизация
|
||
print("== canonical ==")
|
||
check("HDRezka Studio -> HDrezka", canonical("HDRezka Studio") == "HDrezka")
|
||
check("hdrezka -> HDrezka", canonical("hdrezka") == "HDrezka")
|
||
check("Read Head Sound -> Red Head Sound", canonical("Read Head Sound") == "Red Head Sound")
|
||
check("Ultradox Studio -> Ultradox", canonical("Ultradox Studio") == "Ultradox")
|
||
check("Jaskier ненормативная лексика -> Jaskier",
|
||
canonical("Jaskier ненормативная лексика") == "Jaskier")
|
||
check("официальный -> ''", canonical("официальный") == "")
|
||
|
||
# ---------------------------------------------------------------- префиксы
|
||
print("== RE_LINE_PREFIX ==")
|
||
for line, body in [
|
||
("Перевод: Профессиональный (дублированный)", "Профессиональный (дублированный)"),
|
||
("Перевод 2: Любительский (одноголосый закадровый)", "Любительский (одноголосый закадровый)"),
|
||
("Перевод #1: Дублированный", "Дублированный"),
|
||
("Озвучивание: Дублированное | Official", "Дублированное | Official"),
|
||
("Аудио 1: Opus, 6 ch, 221 Kbps | Русский AVO (Дмитрий Пучков)",
|
||
"Opus, 6 ch, 221 Kbps | Русский AVO (Дмитрий Пучков)"),
|
||
("Перевод:Профессиональный (дублированный)", "Профессиональный (дублированный)"),
|
||
]:
|
||
check(f"префикс «{line[:20]}…»", RE_LINE_PREFIX.sub("", line) == body)
|
||
|
||
# ---------------------------------------------------------------- сегментация
|
||
print("== extract_segments ==")
|
||
segs = extract_segments(
|
||
"Описание релиза\n"
|
||
"Перевод 1: Профессиональный (дублированный)\n"
|
||
"Перевод 2: Профессиональный (двухголосый закадровый) tycoon\n"
|
||
"Субтитры: русские, английские\n")
|
||
check("две дорожки", len(segs) == 2, f"got {segs}")
|
||
|
||
segs = extract_segments("Перевод 1: Профессиональный (дублированный) 2. Авторский (одноголосый закадровый) Гоблин")
|
||
check("нумерация 1./2.", len(segs) == 2, f"got {segs}")
|
||
|
||
segs = extract_segments("Перевод: Одноголосый, многоголосый закадровый")
|
||
check("запятая-перечисление", len(segs) == 2, f"got {segs}")
|
||
|
||
segs = extract_segments("Перевод: Профессиональный (многоголосый, закадровый)")
|
||
check("скобки-уточнение не режутся", len(segs) == 1, f"got {segs}")
|
||
|
||
segs = extract_segments("Перевод:\nПрофессиональный (дублированный) Невафильм")
|
||
check("значение на следующей строке",
|
||
segs == ["Профессиональный (дублированный) Невафильм"], f"got {segs}")
|
||
|
||
segs = extract_segments("Перевод: Профессиональный (двухголосый закадровый) + Любительский (одноголосый закадровый)")
|
||
check("разделитель +", len(segs) == 2, f"got {segs}")
|
||
|
||
segs = extract_segments("Аудио 1: Opus, 6 ch, 221 Kbps | Русский AVO (Дмитрий Пучков)\n"
|
||
"Аудио 2: Opus, 6 ch, 199 Kbps | Русский DVO (Кубик в Кубе)")
|
||
check("nnm аудио-дорожки", len(segs) == 2, f"got {segs}")
|
||
|
||
segs = extract_segments("Аудио: Русский (AC3, 6 ch, 640 Кбит/с), Английский (AC3)")
|
||
check("codec-строка без типа пропущена", len(segs) == 0, f"got {segs}")
|
||
|
||
check("_split_entries одиночная", _split_entries("Дублированный") == ["Дублированный"])
|
||
|
||
# ---------------------------------------------------------------- html
|
||
print("== html_to_text ==")
|
||
html = ("<html><body><div>меню форума</div>"
|
||
'<div class="post_body">Перевод: Профессиональный (дублированный) Невафильм</div>'
|
||
"</body></html>")
|
||
text = html_to_text(html)
|
||
check("селектор .post_body", "Невафильм" in text and "меню форума" not in text)
|
||
|
||
html2 = "<html><body><p>hello</p></body></html>"
|
||
check("fallback на весь текст", "hello" in html_to_text(html2))
|
||
|
||
print("== read_input / mojibake ==")
|
||
import tempfile
|
||
|
||
from src.html_extract import read_input
|
||
|
||
with tempfile.TemporaryDirectory() as td:
|
||
p = Path(td) / "mojibake.html"
|
||
p.write_bytes("Перевод: Дублированный".encode("utf-8").decode("cp1251",
|
||
errors="ignore").encode("utf-8"))
|
||
check("mojibake восстановлен", "Дублированный" in read_input(p))
|
||
p2 = Path(td) / "cp1251.txt"
|
||
p2.write_bytes("Перевод: Дублированный".encode("cp1251"))
|
||
check("cp1251 прочитан", "Дублированный" in read_input(p2))
|
||
|
||
# ---------------------------------------------------------------- BIO
|
||
print("== BIO-выравнивание (требует токенизатор) ==")
|
||
try:
|
||
from transformers import AutoTokenizer
|
||
|
||
from src.train import make_collate
|
||
|
||
tok = AutoTokenizer.from_pretrained("models/best")
|
||
collate = make_collate(tok, {"проф. дубляж": 0})
|
||
text = "Авторский (одноголосый закадровый) - Сергей Визгунов"
|
||
s0 = text.index("Сергей")
|
||
enc, types, bio, ages, offsets = collate(
|
||
[{"text": text, "translation_type": "проф. дубляж",
|
||
"span": [s0, s0 + len("Сергей Визгунов")], "is_18plus": False}])
|
||
marked = {i for i, lab in enumerate(bio[0].tolist()) if lab >= 1}
|
||
span_toks = [tok.convert_ids_to_tokens(int(enc["input_ids"][0, i])) for i in sorted(marked)]
|
||
joined = "".join(t.lstrip("#") for t in span_toks).lower()
|
||
# у ruBert-base токенизатор strips accents: «й» может стать «и»
|
||
check("токены спана = ФИ", "визгунов" in joined and joined.startswith("сер"),
|
||
f"tokens={span_toks}")
|
||
check("первый токен B", bio[0, min(marked)].item() == 1 if marked else False)
|
||
except Exception as e: # noqa: BLE001
|
||
print(f" пропущено (нет модели models/best): {e}")
|
||
|
||
print(f"\nитог: {PASS} ok, {FAIL} failed")
|
||
sys.exit(1 if FAIL else 0)
|