File size: 5,647 Bytes
71dabc7 0cd599e 71dabc7 0cd599e 71dabc7 0cd599e 71dabc7 0cd599e 71dabc7 | 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 | """量化变体加载器（laya-ad 发布格式）。
支持 laya-ad Release 的五种 safetensors 变体并统一加载为可推理模型：
- fp16 / bf16：transformers 原生加载
- fp8（float8_e4m3fn 存储）：cast 回 bf16
- int8（per-channel 对称 weight-only，`__scale__` 伴生 fp16）
- int4（group=64 对称、双 nibble 打包 uint8、`__scale__`/`__pad__` 伴生）
由 scripts/quantize_variants.py 产出；格式与 laya-ad 仓库 scripts/load_quantized.py 一致。
"""
import json
from pathlib import Path
from types import SimpleNamespace
import torch
from safetensors import safe_open
INT4_GROUP = 64
class _DualHeadBinary(torch.nn.Module):
"""把多任务模型包装成服务端可用的模型：binary/kind/scope/severity 四组 logits。"""
def __init__(self, model, kind_labels: list[str] | None = None,
scope_labels: list[str] | None = None,
severity_labels: list[str] | None = None):
super().__init__()
self.model = model
self.config = model.config
self.kind_labels = list(kind_labels or [])
self.scope_labels = list(scope_labels or [])
self.severity_labels = list(severity_labels or [])
self.has_scope = bool(self.scope_labels) and getattr(model, "scope_head", None) is not None
self.has_severity = (bool(self.severity_labels)
and getattr(model, "severity_head", None) is not None)
def forward(self, input_ids=None, attention_mask=None, **kwargs):
binary_logits, kind_logits, scope_logits, severity_logits, _ = self.model.forward_all(
input_ids, attention_mask)
return SimpleNamespace(logits=binary_logits, kind_logits=kind_logits,
scope_logits=scope_logits, severity_logits=severity_logits)
def detect_precision(model_dir: Path) -> str:
"""优先读 PRECISION.txt（CI 注入），否则从权重 dtype 推断。"""
marker = Path(model_dir) / "PRECISION.txt"
if marker.is_file():
return marker.read_text(encoding="utf-8").strip() or "unknown"
sd = Path(model_dir) / "model.safetensors"
dts = set()
with safe_open(str(sd), framework="pt") as f:
for k in f.keys():
dts.add(f.get_slice(k).get_dtype())
if "F8_E4M3" in dts:
return "fp8"
if "U8" in dts:
return "int4"
if "I8" in dts:
return "int8"
if "BF16" in dts:
return "bf16"
if "F16" in dts:
return "fp16"
return "fp32"
def load_ad_model(model_dir, device: str = "cpu", dtype: torch.dtype = torch.bfloat16):
"""加载广告分类模型（自动识别量化格式），返回 (model, tokenizer)，eval 模式。"""
from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer
model_dir = Path(model_dir)
if (model_dir / "dual_config.json").is_file():
# v2 双头模型：共享 encoder + binary/kind 两个头
from transformers import AutoTokenizer
from .dual_head_model import load_dual_head_any
dual_cfg = json.loads((model_dir / "dual_config.json").read_text(encoding="utf-8"))
kind_labels = [str(x) for x in dual_cfg.get("kind_labels") or []]
scope_labels = [str(x) for x in dual_cfg.get("scope_labels") or []]
severity_labels = [str(x) for x in dual_cfg.get("severity_labels") or []]
dual = load_dual_head_any(str(model_dir), device=device, dtype=dtype)
return (_DualHeadBinary(dual, kind_labels, scope_labels, severity_labels).to(device).eval(),
AutoTokenizer.from_pretrained(str(model_dir)))
sd_path = model_dir / "model.safetensors"
dts = set()
with safe_open(str(sd_path), framework="pt") as f:
for k in f.keys():
dts.add(f.get_slice(k).get_dtype())
if not ({"I8", "U8", "F8_E4M3"} & dts):
kw = {"dtype": dtype}
model = AutoModelForSequenceClassification.from_pretrained(
str(model_dir), attn_implementation="sdpa", **kw)
tok = AutoTokenizer.from_pretrained(str(model_dir))
return model.to(device).eval(), tok
model = AutoModelForSequenceClassification.from_config(
AutoConfig.from_pretrained(str(model_dir)))
param_shapes = {k: v.shape for k, v in model.state_dict().items()}
with safe_open(str(sd_path), framework="pt") as f:
raw = {k: f.get_tensor(k) for k in f.keys()}
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: # 未量化透传（1D bias/norm 等保留精度）
state[k] = v.to(dtype)
model.load_state_dict(state, strict=True)
tok = AutoTokenizer.from_pretrained(str(model_dir))
return model.to(device).eval(), tok
|