| |
| """单编码器双头模型：共享 ModernBERT encoder，两个任务头各自保留自己的 prediction head。 |
| |
| - binary branch: head + classifier 来自 v1.1（768 → 2） |
| - kind branch: head + classifier 来自已有 7 类模型（768 → 7） |
| - 一次前向共享 encoder，两个 head 只多约 0.6M 参数；总参数 ≈ 一个 322M 编码器， |
| 而不是两个完整模型（约 644M）。 |
| """ |
| import copy |
| import json |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn as nn |
| from safetensors import safe_open |
| from safetensors.torch import load_file |
| from transformers import ModernBertConfig, ModernBertForSequenceClassification |
|
|
| INT4_GROUP = 64 |
|
|
|
|
| class DualHeadModernBert(nn.Module): |
| def __init__(self, base_bin, base_kind, kind_num_labels=7, use_kind_adapter=False, |
| kind_adapter_scope="pooled", use_scope_head=False, use_severity_head=False): |
| super().__init__() |
| self.config = base_bin.config |
| self.model = base_bin.model |
| self.drop = base_bin.drop |
| self.binary_head = base_bin.head |
| self.binary_classifier = base_bin.classifier |
| self.kind_head = base_kind.head |
| self.kind_classifier = base_kind.classifier |
| assert self.kind_classifier.out_features == kind_num_labels |
| self.kind_adapter_scope = kind_adapter_scope |
| h = self.config.hidden_size |
| if use_kind_adapter: |
| self.kind_adapter = nn.Sequential( |
| nn.Linear(h, h), nn.GELU(), nn.Dropout(0.1), nn.Linear(h, h)) |
| nn.init.zeros_(self.kind_adapter[-1].weight) |
| nn.init.zeros_(self.kind_adapter[-1].bias) |
| else: |
| self.kind_adapter = None |
| if use_scope_head: |
| self.scope_head = copy.deepcopy(base_kind.head) |
| self.scope_classifier = nn.Linear(h, 2) |
| else: |
| self.scope_head = None |
| self.scope_classifier = None |
| if use_severity_head: |
| self.severity_head = copy.deepcopy(base_kind.head) |
| self.severity_classifier = nn.Linear(h, 4) |
| else: |
| self.severity_head = None |
| self.severity_classifier = None |
|
|
| def _pool(self, h, attention_mask): |
| if self.config.classifier_pooling == "cls": |
| return h[:, 0] |
| if attention_mask is None: |
| attention_mask = torch.ones(h.shape[:2], device=h.device, dtype=torch.bool) |
| return (h * attention_mask.unsqueeze(-1)).sum(1) / attention_mask.sum(1, keepdim=True) |
|
|
| def forward_all(self, input_ids, attention_mask=None, position_ids=None): |
| outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, |
| position_ids=position_ids) |
| tokens = outputs[0] |
| h = self._pool(tokens, attention_mask) |
| bin_logits = self.binary_classifier(self.drop(self.binary_head(h))) |
| if self.kind_adapter is None: |
| hk = h |
| elif self.kind_adapter_scope == "token": |
| hk = self._pool(tokens + self.kind_adapter(tokens), attention_mask) |
| else: |
| hk = h + self.kind_adapter(h) |
| kind_logits = self.kind_classifier(self.drop(self.kind_head(hk))) |
| scope_logits = (self.scope_classifier(self.drop(self.scope_head(h))) |
| if self.scope_head is not None else None) |
| severity_logits = (self.severity_classifier(self.drop(self.severity_head(h))) |
| if self.severity_head is not None else None) |
| return bin_logits, kind_logits, scope_logits, severity_logits, h |
|
|
| def forward(self, input_ids, attention_mask=None, position_ids=None): |
| bin_logits, kind_logits, _, _, _ = self.forward_all( |
| input_ids, attention_mask=attention_mask, position_ids=position_ids) |
| return bin_logits, kind_logits |
|
|
|
|
| def load_dual_head(model_dir, device="cpu", dtype=torch.bfloat16): |
| model_dir = Path(model_dir) |
| dc = json.loads((model_dir / "dual_config.json").read_text(encoding="utf-8")) |
| heads = dc.get("heads") or [] |
| config = ModernBertConfig.from_pretrained(str(model_dir)) |
| cfg_kind = copy.deepcopy(config) |
| cfg_kind.num_labels = len(dc["kind_labels"]) |
| base_bin = ModernBertForSequenceClassification(config) |
| base_kind = ModernBertForSequenceClassification(cfg_kind) |
| model = DualHeadModernBert(base_bin, base_kind, len(dc["kind_labels"]), |
| use_kind_adapter=bool(dc.get("use_kind_adapter", False)), |
| kind_adapter_scope=dc.get("kind_adapter_scope", "pooled"), |
| use_scope_head=("scope" in heads) or bool(dc.get("use_scope_head", False)), |
| use_severity_head=("severity" in heads) or bool(dc.get("use_severity_head", False))) |
| state = load_file(str(model_dir / "model.safetensors")) |
| model.load_state_dict(state) |
| return model.to(device=device, dtype=dtype).eval() |
|
|
|
|
| def _dual_skeleton(model_dir): |
| model_dir = Path(model_dir) |
| dc = json.loads((model_dir / "dual_config.json").read_text(encoding="utf-8")) |
| heads = dc.get("heads") or [] |
| config = ModernBertConfig.from_pretrained(str(model_dir)) |
| cfg_kind = copy.deepcopy(config) |
| cfg_kind.num_labels = len(dc["kind_labels"]) |
| base_bin = ModernBertForSequenceClassification(config) |
| base_kind = ModernBertForSequenceClassification(cfg_kind) |
| model = DualHeadModernBert(base_bin, base_kind, len(dc["kind_labels"]), |
| use_kind_adapter=bool(dc.get("use_kind_adapter", False)), |
| kind_adapter_scope=dc.get("kind_adapter_scope", "pooled"), |
| use_scope_head=("scope" in heads) or bool(dc.get("use_scope_head", False)), |
| use_severity_head=("severity" in heads) or bool(dc.get("use_severity_head", False))) |
| return model, dc |
|
|
|
|
| def _dtypes(path): |
| with safe_open(str(path), framework="pt") as f: |
| return {f.get_slice(k).get_dtype() for k in f.keys()} |
|
|
|
|
| def load_dual_head_quantized(model_dir, device="cpu", dtype=torch.bfloat16): |
| """加载 fp8 / int8 / int4 变体：反量化成双头模型。""" |
| model_dir = Path(model_dir) |
| model, _ = _dual_skeleton(model_dir) |
| param_shapes = {k: v.shape for k, v in model.state_dict().items()} |
| raw = load_file(str(model_dir / "model.safetensors")) |
| state = {} |
| for k, v in raw.items(): |
| if k.endswith(".__scale__") or k.endswith(".__pad__"): |
| continue |
| if v.dtype == torch.float8_e4m3fn: |
| state[k] = v.to(dtype) |
| elif v.dtype == torch.int8 and k + ".__scale__" in raw: |
| scale = raw[k + ".__scale__"].float().view(-1, *([1] * (v.dim() - 1))) |
| state[k] = (v.float() * scale).to(dtype) |
| elif v.dtype == torch.uint8 and k + ".__scale__" in raw: |
| hi = (v >> 4).to(torch.int16) - 8 |
| lo = (v & 0xF).to(torch.int16) - 8 |
| q = torch.stack([hi, lo], dim=-1).reshape(-1).float() |
| scale = raw[k + ".__scale__"].float() |
| s_exp = scale.repeat_interleave(INT4_GROUP)[:q.numel()] |
| w = q * s_exp |
| pad = int(raw[k + ".__pad__"].item()) if k + ".__pad__" in raw else 0 |
| if pad: |
| w = w[:-pad] |
| state[k] = w.reshape(param_shapes[k]).to(dtype) |
| else: |
| state[k] = v.to(dtype) |
| model.load_state_dict(state, strict=True) |
| return model.to(device=device, dtype=dtype).eval() |
|
|
|
|
| def load_dual_head_any(model_dir, device="cpu", dtype=torch.bfloat16): |
| """按 safetensors 存储类型自动选择普通/量化加载器。""" |
| dts = _dtypes(Path(model_dir) / "model.safetensors") |
| if {"I8", "U8", "F8_E4M3"} & dts: |
| return load_dual_head_quantized(model_dir, device=device, dtype=dtype) |
| return load_dual_head(model_dir, device=device, dtype=dtype) |
|
|
|
|
| if __name__ == "__main__": |
| import sys |
| m = load_dual_head_any(sys.argv[1] if len(sys.argv) > 1 else "outputs/ad-dual/best") |
| print("加载成功 | 参数", sum(p.numel() for p in m.parameters()) / 1e6, "M") |
|
|
