"""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()