File size: 8,097 Bytes
dd18e84 f8d2531 dd18e84 f8d2531 dd18e84 f8d2531 dd18e84 f8d2531 dd18e84 f8d2531 dd18e84 f8d2531 dd18e84 f8d2531 dd18e84 f8d2531 dd18e84 f8d2531 dd18e84 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | #!/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")
|