| """量化变体加载器（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(): |
| |
| 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: |
| 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 |
|
|
