DenozordecandClaude Sonnet 5.5
05537ac9e9
docker / build (push) Successful in 1m40s
- harvest_core: дубляж — тип, а не часть имени студии; split_multi_rows; audio_only_rows; kinozal сопоставляет все аудиодорожки - prepare_data: --max-dup, отсев остатков многотиповых строк - train: веса классов (корень из обратной частоты, потолок 4), --no-class-weights - исправлены регулярки с \b, записанные как символ BS (predict.py, audio18.py) - PARSE_VERSION 3: кэш разбора сбрасывается - docs/dataset.md, docs/api-reference.md обновлены Co-Authored-By: Claude Sonnet 5.5 <[email protected]>
135 lines
7.0 KiB
Python
135 lines
7.0 KiB
Python
"""Multi-task модель: тип озвучки + BIO-спан автора + признак 18+.
|
|
|
|
Одна общая кодирующая часть (любой BERT-подобный энкодер с HuggingFace) и три
|
|
головы:
|
|
* type_head — [CLS] -> класс translation_type;
|
|
* bio_head — по токенам -> O / B / I (имя автора в тексте);
|
|
* age_head — [CLS] -> 0/1 (дорожка 18+, «ненормативная лексика»-варианты);
|
|
* lang_head — [CLS] -> язык дорожки (необязательная: n_langs=0 — как раньше;
|
|
метка -100 = язык неизвестен, в лосс не входит).
|
|
|
|
Функция потерь: CE(тип) + span_weight*CE(BIO, ignore_index=-100)
|
|
+ age_weight*CE(18+).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from transformers import AutoConfig, AutoModel
|
|
|
|
|
|
class DubTagger(nn.Module):
|
|
def __init__(self, backbone: str, n_types: int, n_bio: int = 3,
|
|
dropout: float = 0.1, span_weight: float = 1.0,
|
|
age_weight: float = 0.5, n_langs: int = 0, lang_weight: float = 0.3,
|
|
pretrained: bool = True, config=None):
|
|
super().__init__()
|
|
self.backbone_name = backbone
|
|
if pretrained: # обучение: веса основы с Hugging Face
|
|
self.encoder = AutoModel.from_pretrained(backbone)
|
|
else: # загрузка чекпоинта: веса придут из model.pt
|
|
self.encoder = AutoModel.from_config(config or AutoConfig.from_pretrained(backbone))
|
|
hidden = self.encoder.config.hidden_size
|
|
self.dropout = nn.Dropout(dropout)
|
|
self.type_head = nn.Linear(hidden, n_types)
|
|
self.bio_head = nn.Linear(hidden, n_bio)
|
|
self.age_head = nn.Linear(hidden, 2)
|
|
self.n_langs = n_langs
|
|
self.lang_weight = lang_weight
|
|
self.lang_head = nn.Linear(hidden, n_langs) if n_langs else None
|
|
self.n_types = n_types
|
|
self.n_bio = n_bio
|
|
self.span_weight = span_weight
|
|
self.age_weight = age_weight
|
|
self.register_buffer("type_weights", None, persistent=False) # только для обучения; в model.pt не пишется
|
|
|
|
def set_type_weights(self, weights) -> None:
|
|
"""Веса классов для потерь типа (редкие классы весят больше); None — без весов."""
|
|
self.type_weights = None if weights is None else weights.to(self.type_head.weight.device)
|
|
|
|
def forward(self, input_ids, attention_mask, token_type_ids=None,
|
|
type_labels=None, bio_labels=None, age_labels=None,
|
|
lang_labels=None, return_lang: bool = False):
|
|
enc_kwargs = {"input_ids": input_ids, "attention_mask": attention_mask}
|
|
if token_type_ids is not None:
|
|
enc_kwargs["token_type_ids"] = token_type_ids
|
|
out = self.encoder(**enc_kwargs)
|
|
seq = out.last_hidden_state
|
|
|
|
cls = self.dropout(seq[:, 0])
|
|
type_logits = self.type_head(cls)
|
|
age_logits = self.age_head(cls)
|
|
bio_logits = self.bio_head(self.dropout(seq))
|
|
|
|
loss = None
|
|
if type_labels is not None:
|
|
loss = F.cross_entropy(type_logits, type_labels, weight=self.type_weights)
|
|
if bio_labels is not None:
|
|
loss = loss + self.span_weight * F.cross_entropy(
|
|
bio_logits.reshape(-1, self.n_bio),
|
|
bio_labels.reshape(-1), ignore_index=-100)
|
|
if age_labels is not None:
|
|
loss = loss + self.age_weight * F.cross_entropy(
|
|
age_logits, age_labels)
|
|
if lang_labels is not None and self.lang_head is not None \
|
|
and bool((lang_labels != -100).any()):
|
|
loss = loss + self.lang_weight * F.cross_entropy(
|
|
self.lang_head(cls), lang_labels, ignore_index=-100)
|
|
if return_lang:
|
|
lang_logits = self.lang_head(cls) if self.lang_head is not None else None
|
|
return loss, type_logits, bio_logits, age_logits, lang_logits
|
|
return loss, type_logits, bio_logits, age_logits
|
|
|
|
@torch.no_grad()
|
|
def predict(self, input_ids, attention_mask, token_type_ids=None):
|
|
"""Логиты (type, bio, age) без градиентов; модель должна быть в eval()."""
|
|
_, type_logits, bio_logits, age_logits = self(
|
|
input_ids, attention_mask, token_type_ids)
|
|
return type_logits, bio_logits, age_logits
|
|
|
|
|
|
@torch.no_grad()
|
|
def predict_with_lang(model: "DubTagger", **enc):
|
|
"""(type, bio, age, lang) логиты; lang=None, если у модели нет языковой головы."""
|
|
_, t, b, a, lang = model(**enc, return_lang=True)
|
|
return t, b, a, lang
|
|
|
|
|
|
def save_model(model: DubTagger, tokenizer, model_dir: str | Path) -> None:
|
|
"""Сохраняет чекпоинт + токенизатор + конфиг в каталог модели."""
|
|
d = Path(model_dir)
|
|
d.mkdir(parents=True, exist_ok=True)
|
|
torch.save(model.state_dict(), d / "model.pt")
|
|
model.encoder.config.save_pretrained(d) # config.json: сервису не нужен интернет и Hugging Face
|
|
tokenizer.save_pretrained(d)
|
|
cfg = {"backbone": model.backbone_name, "n_types": model.n_types,
|
|
"n_bio": model.n_bio, "dropout": model.dropout.p,
|
|
"span_weight": model.span_weight, "age_weight": model.age_weight,
|
|
"n_langs": model.n_langs, "lang_weight": model.lang_weight}
|
|
(d / "dub_config.json").write_text(
|
|
json.dumps(cfg, ensure_ascii=False, indent=1), encoding="utf-8")
|
|
|
|
|
|
def load_model(model_dir: str | Path, device: str = "cpu") -> DubTagger:
|
|
"""Загружает чекпоинт, сохранённый save_model.
|
|
|
|
Архитектура берётся из config.json рядом с весами (сервис работает офлайн); у старых чекпоинтов
|
|
без него конфиг запрашивается у Hugging Face (нужен кэш или сеть), сами веса основы не качаются.
|
|
"""
|
|
d = Path(model_dir)
|
|
cfg = json.loads((d / "dub_config.json").read_text(encoding="utf-8"))
|
|
enc_cfg = AutoConfig.from_pretrained(d) if (d / "config.json").exists() else None
|
|
model = DubTagger(cfg["backbone"], cfg["n_types"], cfg["n_bio"], pretrained=False, config=enc_cfg,
|
|
dropout=cfg.get("dropout", 0.1),
|
|
span_weight=cfg.get("span_weight", 1.0),
|
|
age_weight=cfg.get("age_weight", 0.5),
|
|
n_langs=cfg.get("n_langs", 0), lang_weight=cfg.get("lang_weight", 0.3))
|
|
state = torch.load(d / "model.pt", map_location="cpu", weights_only=True)
|
|
model.load_state_dict(state) # чекпоинт без lang_head грузится: n_langs=0
|
|
return model.to(device).eval()
|