DUB-AI: multi-task модель (тип озвучки + автор-спан + 18+), обучение/инференс/ONNX/сервис, тесты и документация
This commit is contained in:
DUB-AI committed
commit
8dbca3f4a3
2873 files changed
+44479
No files matched your search
@@ -0,0 +1,167 @@
|
||||
"""Юнит-тесты 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)
|
||||
Reference in new issue
Block a user