#!/usr/bin/env python3 """单编码器双头模型：共享 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")