DenozordecandClaude Sonnet 5.5
8cbd800ec7
- src/fetch.py: загрузка по URL (allowlist трекеров, защита от SSRF, cp1251)
- src/release.py: агрегат релиза, студии из заголовка раздачи (ПМ (LostFilm))
- predict --url, POST /predict {"url"}; ответ расширен, tracks сохранён
- export_onnx --quantize: tiny2 int8 29 МБ, parity 99.2-100%
- тесты (52 ok), docs/plan.md, docs; target/ убран из git
Co-Authored-By: Claude Sonnet 5.5 <[email protected]>
204 lines
8.7 KiB
Python
204 lines
8.7 KiB
Python
"""Экспорт обученной модели в ONNX + проверка эквивалентности PyTorch.
|
|
|
|
Запуск:
|
|
python -m src.export_onnx [--model models/best]
|
|
|
|
Результат:
|
|
models/best-onnx/model.onnx — энкодер с обеими головами
|
|
models/best-onnx/tokenizer.json … — файлы токенизатора (HF fast)
|
|
models/best-onnx/label_maps.json — классы типа + BIO
|
|
models/best-onnx/aliases.json — словарь алиасов студий для клиента
|
|
reports/onnx_parity.md — сравнение ONNX vs PyTorch на test
|
|
|
|
Выходы ONNX: type_logits [B, n_types], bio_logits [B, seq, 3],
|
|
age_logits [B, 2] (дорожка 18+).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
from transformers import AutoTokenizer
|
|
|
|
from .model import load_model
|
|
from .studio_norm import STUDIO_ALIASES, STUDIO_PLACEHOLDERS
|
|
from .train import MAX_LEN, Rows, decode_span, make_collate
|
|
from torch.utils.data import DataLoader
|
|
|
|
|
|
class _ExportWrapper(torch.nn.Module):
|
|
"""Фиксированная сигнатура для экспорта: три входа, три выхода."""
|
|
|
|
def __init__(self, model):
|
|
super().__init__()
|
|
self.m = model
|
|
|
|
def forward(self, input_ids, attention_mask, token_type_ids):
|
|
_, type_logits, bio_logits, age_logits = self.m(
|
|
input_ids, attention_mask, token_type_ids)
|
|
return type_logits, bio_logits, age_logits
|
|
|
|
|
|
def export(model_dir: str, out_dir: str) -> Path:
|
|
d = Path(model_dir)
|
|
out = Path(out_dir)
|
|
out.mkdir(parents=True, exist_ok=True)
|
|
model = load_model(d)
|
|
model.eval()
|
|
wrapper = _ExportWrapper(model).eval()
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(d)
|
|
sample = tokenizer(["Профессиональный (многоголосый закадровый) HDRezka Studio"],
|
|
return_tensors="pt", truncation=True, max_length=MAX_LEN)
|
|
ids = sample["input_ids"]
|
|
mask = sample["attention_mask"]
|
|
tti = sample["token_type_ids"]
|
|
|
|
f = out / "model.onnx"
|
|
try:
|
|
# legacy-экспортёр: проверенные dynamic_axes для batch/seq
|
|
torch.onnx.export(
|
|
wrapper, (ids, mask, tti), f, dynamo=False,
|
|
input_names=["input_ids", "attention_mask", "token_type_ids"],
|
|
output_names=["type_logits", "bio_logits", "age_logits"],
|
|
dynamic_axes={"input_ids": {0: "batch", 1: "seq"},
|
|
"attention_mask": {0: "batch", 1: "seq"},
|
|
"token_type_ids": {0: "batch", 1: "seq"},
|
|
"type_logits": {0: "batch"},
|
|
"bio_logits": {0: "batch", 1: "seq"},
|
|
"age_logits": {0: "batch"}},
|
|
opset_version=17,
|
|
)
|
|
except Exception as e: # noqa: BLE001 — fallback на dynamo-экспортёр
|
|
print(f"legacy-экспорт не удался ({e}); пробуем dynamo")
|
|
torch.onnx.export(
|
|
wrapper, (ids, mask, tti), f, dynamo=True,
|
|
dynamic_shapes={
|
|
"input_ids": {0: "batch", 1: "seq"},
|
|
"attention_mask": {0: "batch", 1: "seq"},
|
|
"token_type_ids": {0: "batch", 1: "seq"},
|
|
},
|
|
external_data=False,
|
|
)
|
|
|
|
tokenizer.save_pretrained(out) # tokenizer.json и др.
|
|
shutil.copy(d / "label_maps.json", out / "label_maps.json")
|
|
(out / "aliases.json").write_text(json.dumps(
|
|
{"aliases": STUDIO_ALIASES, "placeholders": sorted(STUDIO_PLACEHOLDERS)},
|
|
ensure_ascii=False, indent=1), encoding="utf-8")
|
|
return f
|
|
|
|
|
|
def quantize_int8(onnx_dir: str) -> Path:
|
|
"""Динамическая int8-квантизация весов -> model.int8.onnx (в 3-4 раза меньше)."""
|
|
from onnxruntime.quantization import QuantType, quantize_dynamic
|
|
|
|
src = Path(onnx_dir) / "model.onnx"
|
|
dst = Path(onnx_dir) / "model.int8.onnx"
|
|
quantize_dynamic(str(src), str(dst), weight_type=QuantType.QInt8)
|
|
return dst
|
|
|
|
|
|
def parity(model_dir: str, onnx_dir: str, data_path: str, limit: int = 2000,
|
|
onnx_name: str = "model.onnx") -> dict:
|
|
"""Сравнение ONNX vs PyTorch: совпадение типа и текста спана."""
|
|
import onnxruntime as ort
|
|
|
|
rows = [r for r in json.loads(Path(data_path).read_text(encoding="utf-8"))
|
|
if r["split"] == "test"][:limit]
|
|
label_maps = json.loads((Path(model_dir) / "label_maps.json")
|
|
.read_text(encoding="utf-8"))
|
|
types = label_maps["types"]
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_dir)
|
|
model = load_model(model_dir)
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
model = model.to(device)
|
|
|
|
sess = ort.InferenceSession(str(Path(onnx_dir) / onnx_name),
|
|
providers=["CPUExecutionProvider"])
|
|
input_names = {i.name for i in sess.get_inputs()}
|
|
|
|
collate = make_collate(tokenizer, {t: i for i, t in enumerate(types)})
|
|
loader = DataLoader(Rows(rows), batch_size=64, collate_fn=collate)
|
|
|
|
n = same_type = same_span = same_age = 0
|
|
max_logit_diff = 0.0
|
|
with torch.no_grad():
|
|
for enc, _, _, ages, offsets in loader:
|
|
enc_dev = {k: v.to(device) for k, v in enc.items()}
|
|
t_torch, b_torch, a_torch = model.predict(**enc_dev)
|
|
feed = {"input_ids": enc["input_ids"].numpy(),
|
|
"attention_mask": enc["attention_mask"].numpy()}
|
|
if "token_type_ids" in input_names:
|
|
feed["token_type_ids"] = enc["token_type_ids"].numpy()
|
|
t_onnx, b_onnx, a_onnx = sess.run(None, feed)
|
|
|
|
tpt = t_torch.argmax(-1).cpu().numpy()
|
|
pot = t_onnx.argmax(-1)
|
|
apt = a_torch.argmax(-1).cpu().numpy()
|
|
apo = a_onnx.argmax(-1)
|
|
bpt = b_torch.argmax(-1).cpu().numpy()
|
|
bot = b_onnx.argmax(-1)
|
|
max_logit_diff = max(max_logit_diff, float(np.abs(
|
|
t_torch.cpu().numpy() - t_onnx).max()))
|
|
for i in range(len(pot)):
|
|
n += 1
|
|
same_type += tpt[i] == pot[i]
|
|
same_age += apt[i] == apo[i]
|
|
attn = enc["attention_mask"][i].tolist()
|
|
st = decode_span(bpt[i], offsets[i], attn)
|
|
so = decode_span(bot[i], offsets[i], attn)
|
|
row = loader.dataset.rows[n - 1]
|
|
st_txt = row["text"][st[0]:st[1]] if st else None
|
|
so_txt = row["text"][so[0]:so[1]] if so else None
|
|
same_span += (st_txt or "") == (so_txt or "")
|
|
return {"n": n, "type_match": same_type / max(n, 1),
|
|
"span_match": same_span / max(n, 1),
|
|
"age_match": same_age / max(n, 1),
|
|
"max_logit_diff": max_logit_diff}
|
|
|
|
|
|
def main() -> None:
|
|
ap = argparse.ArgumentParser(
|
|
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
ap.add_argument("--model", default="models/best")
|
|
ap.add_argument("--out", default=None, help="по умолчанию <model>-onnx")
|
|
ap.add_argument("--data", default="data/processed/dataset.json")
|
|
ap.add_argument("--quantize", action="store_true",
|
|
help="дополнительно model.int8.onnx; parity считается по нему")
|
|
ap.add_argument("--report", default="reports/onnx_parity.md")
|
|
args = ap.parse_args()
|
|
|
|
out = args.out or f"{args.model}-onnx"
|
|
f = export(args.model, out)
|
|
print(f"exported: {f} ({f.stat().st_size / 1e6:.1f} MB)")
|
|
onnx_name = "model.onnx"
|
|
if args.quantize:
|
|
q = quantize_int8(out)
|
|
print(f"quantized: {q} ({q.stat().st_size / 1e6:.1f} MB)")
|
|
onnx_name = q.name
|
|
res = parity(args.model, out, args.data, onnx_name=onnx_name)
|
|
ok = (res["type_match"] >= 0.99 and res["span_match"] >= 0.99
|
|
and res["age_match"] >= 0.99)
|
|
md = ["# ONNX parity (test split)", "",
|
|
f"- модель: {args.model} / {onnx_name}"
|
|
f" ({(Path(out) / onnx_name).stat().st_size / 1e6:.1f} MB)",
|
|
f"- строк: {res['n']}",
|
|
f"- совпадение класса типа: {res['type_match']:.4f}",
|
|
f"- совпадение текста спана: {res['span_match']:.4f}",
|
|
f"- совпадение флага 18+: {res['age_match']:.4f}",
|
|
f"- макс. расхождение type_logits: {res['max_logit_diff']:.2e}",
|
|
"", f"Вердикт: {'**OK** (≥99% все критерии)' if ok else '**FAIL**'}"]
|
|
Path(args.report).write_text("\n".join(md), encoding="utf-8")
|
|
print("\n".join(md[1:]))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|