laya-ad / code /laya-server /backend /server /ad_classifier.py
x6nux's picture
laya-ad v2.2: sync laya-server code (default version v2.2, build_push.sh, .env.example)
0cd599e verified
Raw History Blame Contribute Delete
22.8 kB
"""广告分类 Predictor:把 laya-ad 微调的 mmBERT 二分类模型接入 Jev 兼容 API。
包含此前工作流验证的 CPU 优化:
- 动态微批:并发请求在 15ms 窗口内合批前向(单 worker 多线程下 100 并发整批 ≤1s SLO 的关键)
- CPU 线程数:LAYA_CPU_THREADS(默认物理核数 ≈ 逻辑核一半,Zen5 实测最优档)
- bf16 推理:Zen5 AVX-512 BF16 有效
支持的题目(其余标记 unsupported 并返回无信息中性答案,绝不冒充判断):
- choice 且标签恰为 {ad, clean}(大小写不敏感)→ 分类器直接输出
- noul 且 instructions 提及广告(ad/advert/广告)→ "noul" 字段输出 P(is_ad)
usage.input_tokens 为真实 token 数;output_tokens = 题目数(每题一个决策位)。
"""
from __future__ import annotations
import json
import os
import queue
import re
import threading
import time
from pathlib import Path
from typing import Any, Callable, Protocol
import torch
from .quant_loader import detect_precision, load_ad_model
_AD_RE = re.compile(r"ad|advert|广告", re.IGNORECASE)
_AD_LABELS = {"ad", "clean"}
class Predictor(Protocol):
"""Jev 推理 Predictor 协议(与上游 LayaAdapter 相同的形状契约)。"""
def predict(self, state: Any, questions: dict[str, Any], model: str | None = None,
*, max_len: int | None = None) -> dict[str, Any]: ...
def model_identity(model_dir: Path) -> tuple[str, str]:
"""读取模型目录身份:(model_id, precision)。缺省文件时回退通用名。"""
model_dir = Path(model_dir)
precision = "unknown"
p_file = model_dir / "PRECISION.txt"
if p_file.is_file():
precision = p_file.read_text(encoding="utf-8").strip() or "unknown"
mid_file = model_dir / "MODEL_ID.txt"
model_id = mid_file.read_text(encoding="utf-8").strip() if mid_file.is_file() else ""
if not model_id:
model_id = "laya" if precision == "fp8" else f"laya-{precision}"
return model_id, precision
def _version_of(model_id: str) -> str | None:
"""'laya-v1.0-fp8' → 'v1.0';不符合 laya-<版本>-<精度> 模式则 None。"""
if not model_id.startswith("laya-"):
return None
rest = model_id[len("laya-"):]
return rest.rsplit("-", 1)[0] if "-" in rest else None
def _accepted_ids(model_id: str, precision: str) -> list[str]:
"""单模型镜像可接受的 model id 集合(规范序:auto, laya, laya-<版本>, 完整 id)。
laya/laya-<版本> 别名仅 fp8 默认镜像提供。"""
ids = ["auto"]
if precision == "fp8":
version = _version_of(model_id)
if version:
ids += ["laya", f"laya-{version}"]
ids.append(model_id)
return ids
def _extract_text(state: Any) -> str:
"""把 Jev state 序列化为分类文本:消息正文优先,随后 sender/bio/bio_links、
matched_rules、recent_context、known_ad_patterns/false_positives 也进入输入。"""
if isinstance(state, str):
return state
if isinstance(state, list):
parts = []
for turn in state:
if isinstance(turn, dict) and isinstance(turn.get("text"), str):
parts.append(turn["text"])
return "\n".join(parts) if parts else json.dumps(state, ensure_ascii=False)
if not isinstance(state, dict):
return str(state)
msg = state.get("message")
if isinstance(msg, dict):
text = msg.get("text") or ""
elif isinstance(msg, str):
text = msg
else:
text = ""
if not isinstance(text, str):
text = str(text)
if not text:
for key in ("text", "body", "content"):
if isinstance(state.get(key), str):
text = state[key]
break
parts: list[str] = []
if text:
parts.append("[消息]" + text)
chat = state.get("chat")
if isinstance(chat, dict) and chat.get("title"):
parts.append("[群]" + str(chat["title"]))
sender = state.get("sender")
if isinstance(sender, dict):
bits = []
name = " ".join(str(sender[k]) for k in ("first_name", "last_name")
if sender.get(k)).strip()
if name:
bits.append("昵称: " + name)
if sender.get("username"):
bits.append("用户名: @" + str(sender["username"]).lstrip("@"))
if sender.get("bio"):
bits.append("简介: " + str(sender["bio"]))
if sender.get("age_hours") is not None:
bits.append(f"入群时长(h): {sender['age_hours']}")
if sender.get("msgs_in_group") is not None:
bits.append(f"群内发言数: {sender['msgs_in_group']}")
if sender.get("prior_ad_hits") is not None:
bits.append(f"历史广告命中: {sender['prior_ad_hits']}")
if bits:
parts.append("[发送者]" + " | ".join(bits))
links = sender.get("bio_links")
if isinstance(links, list) and links:
lb = []
for x in links[:4]:
if not isinstance(x, dict):
continue
s = " / ".join(str(x[k]) for k in ("kind", "title", "handle", "about")
if x.get(k))
if s:
lb.append(s)
if lb:
parts.append("[资料链接]" + " ; ".join(lb))
rules = state.get("matched_rules")
if isinstance(rules, list) and rules:
rb = []
for x in rules[:4]:
if not isinstance(x, dict):
continue
rb.append(" / ".join(str(x[k]) for k in ("name", "category", "note")
if x.get(k)))
if rb:
parts.append("[命中规则]" + " ; ".join(rb))
ctx = state.get("recent_context")
if isinstance(ctx, list) and ctx:
cb = []
for x in ctx[-5:]:
if isinstance(x, dict):
cb.append(f"{x.get('user', '')}: {x.get('text', '')}")
elif isinstance(x, str):
cb.append(x)
if cb:
parts.append("[近期消息]" + " / ".join(cb))
for key, label in (("known_ad_patterns", "[已知广告模式]"),
("known_false_positives", "[已知误报]")):
val = state.get(key)
if val:
parts.append(label + str(val)[:400])
if isinstance(msg, dict):
flags = []
for key, label in (("has_link", "含链接"), ("has_mention", "含@"),
("has_media", "含媒体"), ("is_forwarded", "转发")):
if msg.get(key):
flags.append(label)
if flags:
parts.append("[消息特征]" + " ".join(flags))
return "\n".join(parts) if parts else json.dumps(state, ensure_ascii=False)
def _labels(criteria: Any) -> list[str]:
if isinstance(criteria, dict):
return [str(k) for k in criteria]
if isinstance(criteria, list):
return [str(x) for x in criteria]
return []
class _MicroBatcher:
"""请求合批器:把并发单条请求在窗口内合并成一次 padding 前向。"""
def __init__(self, forward_fn: Callable[[list[str], int], list[float]],
max_batch: int, window_ms: int):
self.forward_fn = forward_fn
self.max_batch = max(1, max_batch)
self.window_s = max(0, window_ms) / 1000.0
self._queue: queue.Queue = queue.Queue()
self._worker = threading.Thread(target=self._loop, daemon=True,
name="ad-clf-microbatch")
self._worker.start()
def submit(self, text: str, seq_len: int) -> tuple[float, float]:
item = {"text": text, "seq_len": seq_len, "event": threading.Event()}
self._queue.put(item)
item["event"].wait()
if item.get("error") is not None:
raise item["error"]
return item["probs"]
def _loop(self) -> None:
while True:
batch = [self._queue.get()]
deadline = time.monotonic() + self.window_s
while len(batch) < self.max_batch:
remaining = deadline - time.monotonic()
if remaining <= 0:
break
try:
batch.append(self._queue.get(timeout=remaining))
except queue.Empty:
break
try:
seq_len = min(512, max(b["seq_len"] for b in batch))
probs = self.forward_fn([b["text"] for b in batch], seq_len)
for b, p in zip(batch, probs):
b["probs"] = p
b["event"].set()
except Exception as exc: # 单批失败只影响本批请求
for b in batch:
b["error"] = exc
b["event"].set()
class AdClassifierPredictor:
"""实现与 LayaAdapter 相同的 Predictor 协议,后端为 laya-ad 二分类模型。"""
def __init__(self, model_dir: Path, device: str | None):
threads = int(os.environ.get("LAYA_CPU_THREADS",
str(max(1, (os.cpu_count() or 2) // 2))))
if torch.cuda.is_available() and (device or "").startswith("cuda"):
self.device = "cuda"
else:
self.device = device or "cpu"
torch.set_num_threads(threads if self.device == "cpu" else max(4, threads // 2))
self.model_dir = Path(model_dir)
self.model, self.tokenizer = load_ad_model(self.model_dir, self.device)
self.model_id, self.precision = model_identity(self.model_dir)
self.accepted_ids = _accepted_ids(self.model_id, self.precision)
id2label = {int(k): str(v).casefold() for k, v in self.model.config.id2label.items()}
self.ad_id = next((i for i, name in id2label.items() if name == "ad"), None)
self.clean_id = next((i for i, name in id2label.items() if name == "normal"), None)
if self.ad_id is None or self.clean_id is None:
raise RuntimeError(f"模型 id2label 缺少 ad/normal: {id2label}")
self.kind_labels = [str(x) for x in getattr(self.model, "kind_labels", [])]
self.kind_index = {label.casefold(): i for i, label in enumerate(self.kind_labels)}
self.scope_labels = [str(x) for x in getattr(self.model, "scope_labels", [])]
self.scope_index = {label.casefold(): i for i, label in enumerate(self.scope_labels)}
self.severity_labels = [str(x) for x in getattr(self.model, "severity_labels", [])]
self.severity_index = {label.casefold(): i for i, label in enumerate(self.severity_labels)}
self._batcher = _MicroBatcher(self._forward_batch,
int(os.environ.get("LAYA_AD_BATCH", "32")),
int(os.environ.get("LAYA_AD_BATCH_WINDOW_MS", "15")))
def served_ids(self) -> list[str]:
return self.accepted_ids
def _forward_batch(self, texts: list[str], seq_len: int):
enc = self.tokenizer(texts, truncation=True, max_length=seq_len,
padding=True, return_tensors="pt")
enc = {k: v.to(self.device) for k, v in enc.items()}
with torch.inference_mode():
out = self.model(**enc)
logits = out.logits.float()
probs = torch.softmax(logits, dim=-1)
kind_logits = getattr(out, "kind_logits", None)
kind_probs = (torch.softmax(kind_logits.float(), dim=-1)
if kind_logits is not None else None)
scope_logits = getattr(out, "scope_logits", None)
scope_probs = (torch.softmax(scope_logits.float(), dim=-1)
if scope_logits is not None else None)
sev_logits = getattr(out, "severity_logits", None)
sev_probs = (torch.softmax(sev_logits.float(), dim=-1)
if sev_logits is not None else None)
return [(float(p[self.ad_id]), float(p[self.clean_id]),
kind_probs[i].cpu().tolist() if kind_probs is not None else None,
scope_probs[i].cpu().tolist() if scope_probs is not None else None,
sev_probs[i].cpu().tolist() if sev_probs is not None else None)
for i, p in enumerate(probs)]
def predict(self, state: Any, questions: dict[str, Any], model: str | None = None,
*, max_len: int | None = None) -> dict[str, Any]:
text = _extract_text(state)
seq_len = min(int(max_len) if max_len else 512, 512)
p_ad, p_clean, kind_probs, scope_probs, sev_probs = self._batcher.submit(text, seq_len)
pred_label = "ad" if p_ad >= p_clean else "clean"
answers: dict[str, Any] = {}
for qid, q in questions.items():
qtype = q.get("type")
labels = _labels(q.get("criteria"))
if qtype == "choice" and labels and {l.casefold() for l in labels} == _AD_LABELS:
prob_map = {}
for label in labels:
prob_map[label] = round(p_ad if label.casefold() == "ad" else p_clean, 4)
answers[qid] = {
"type": "choice",
"choice": pred_label if pred_label in [l.casefold() for l in labels] else labels[0],
"probabilities": prob_map,
"confidence": round(max(p_ad, p_clean), 4),
"answer_confidence": round(max(p_ad, p_clean), 4),
"unsupported": False,
}
elif (qtype == "choice" and labels and kind_probs is not None
and all(l.casefold() in self.kind_index or l.casefold() == "none"
for l in labels)):
raw = {}
for label in labels:
key = label.casefold()
if key == "none":
raw[label] = p_clean
elif key in self.kind_index:
raw[label] = p_ad * kind_probs[self.kind_index[key]]
else:
raw = None
break
if raw is None:
answers[qid] = self._unsupported(q)
else:
total = sum(raw.values()) or 1.0
prob_map = {label: round(v / total, 4) for label, v in raw.items()}
best = max(raw.items(), key=lambda x: x[1])[0]
answers[qid] = {
"type": "choice",
"choice": best,
"probabilities": prob_map,
"confidence": round(max(raw.values()) / total, 4),
"answer_confidence": round(max(raw.values()) / total, 4),
"unsupported": False,
}
elif (qtype == "choice" and labels and scope_probs is not None
and all(l.casefold() in self.scope_index for l in labels)):
if p_ad >= p_clean:
vals = [scope_probs[self.scope_index[l.casefold()]] for l in labels]
else: # 非广告按 teacher 口径默认 message
vals = [1.0 if l.casefold() == "message" else 0.0 for l in labels]
total = sum(vals) or 1.0
prob_map = {label: round(v / total, 4) for label, v in zip(labels, vals)}
best = max(zip(labels, vals), key=lambda x: x[1])[0]
answers[qid] = {
"type": "choice",
"choice": best,
"probabilities": prob_map,
"confidence": round(max(vals) / total, 4),
"answer_confidence": round(max(vals) / total, 4),
"unsupported": False,
}
elif (qtype == "score" and sev_probs is not None and len(labels) == 4
and all(l.casefold() in self.severity_index for l in labels)):
if p_ad >= p_clean:
vals = [sev_probs[self.severity_index[l.casefold()]] for l in labels]
else:
vals = [1.0, 0.0, 0.0, 0.0]
total = sum(vals) or 1.0
probs = {str(i): round(v / total, 4) for i, v in enumerate(vals)}
score = sum(i * (v / total) for i, v in enumerate(vals))
answers[qid] = {
"type": "score",
"score": round(score, 4),
"legend": {str(i): str(label) for i, label in enumerate(labels)},
"probabilities": probs,
"confidence": round(max(vals) / total, 4),
"answer_confidence": round(max(vals) / total, 4),
"unsupported": False,
}
elif qtype == "noul" and _AD_RE.search(str(q.get("instructions", ""))):
# noul 约定:true = 是广告(仅当题目语义与广告相关时才有意义)
answers[qid] = {
"type": "noul",
"noul": round(p_ad, 4),
"confidence": round(max(p_ad, 1.0 - p_ad), 4),
"answer_confidence": round(max(p_ad, 1.0 - p_ad), 4),
"unsupported": False,
}
else:
answers[qid] = self._unsupported(q)
usage = {"input_tokens": self._count_tokens(text, seq_len) * len(questions),
"output_tokens": len(questions)}
return {"answers": answers, "model": self.model_id, "usage": usage}
def _count_tokens(self, text: str, seq_len: int) -> int:
enc = self.tokenizer(text, truncation=True, max_length=seq_len)
return int(len(enc["input_ids"]))
@staticmethod
def _unsupported(q: dict[str, Any]) -> dict[str, Any]:
"""无信息中性答案:均匀概率 + 低置信 + unsupported 标记,绝不冒充判断。"""
qtype = q.get("type")
if qtype == "choice":
labels = _labels(q.get("criteria")) or ["unknown"]
n = len(labels)
return {"type": "choice", "choice": labels[0],
"probabilities": {l: round(1.0 / n, 4) for l in labels},
"confidence": round(1.0 / n, 4), "answer_confidence": round(1.0 / n, 4),
"unsupported": True}
if qtype == "score":
n = len(q.get("criteria") or [])
uniform = 1.0 / max(1, n)
return {"type": "score", "score": round(sum(i * uniform for i in range(n)), 4),
"legend": {str(i): str(c) for i, c in enumerate(q.get("criteria") or [])},
"probabilities": {str(i): round(uniform, 4) for i in range(n)},
"confidence": round(uniform, 4), "answer_confidence": round(uniform, 4),
"unsupported": True}
# noul(与广告无关的是非题):0.5 = 无信息
return {"type": "noul", "noul": 0.5, "confidence": 0.5, "answer_confidence": 0.5,
"unsupported": True}
class MultiPrecisionPredictor:
"""多精度镜像的 Predictor:LAYA_MODEL_DIR 下每个子目录一个精度变体
(各自含 MODEL_ID.txt / PRECISION.txt / model.safetensors),
按请求的 model id 懒加载切换;auto/laya/laya-<版本> 指向 fp8 变体。"""
def __init__(self, model_dir: Path, device: str | None):
self.root = Path(model_dir)
self.variants: dict[str, tuple[Path, str]] = {}
self.default_id: str | None = None
for d in sorted(self.root.iterdir()):
if not d.is_dir() or not (d / "model.safetensors").is_file():
continue
mid, prec = model_identity(d)
self.variants[mid] = (d, prec)
if prec == "fp8":
self.default_id = mid
if not self.variants:
raise RuntimeError(f"多模型目录为空: {self.root}")
if self.default_id is None:
self.default_id = next(iter(self.variants))
self.version = _version_of(self.default_id)
self.accepted_ids = ["auto", "laya"] \
+ ([f"laya-{self.version}"] if self.version else []) \
+ sorted(self.variants)
self.device = device or "cpu"
self._cache: dict[str, AdClassifierPredictor] = {}
def resolve(self, model: str | None) -> "AdClassifierPredictor":
if model in (None, "auto", "laya") or (self.version and model == f"laya-{self.version}"):
mid = self.default_id
elif model in self.variants:
mid = model
else:
raise ValueError(f"unknown model id: {model}")
if mid not in self._cache:
d, _ = self.variants[mid]
self._cache[mid] = AdClassifierPredictor(d, self.device)
return self._cache[mid]
def refresh(self) -> None:
"""重新扫描模型目录(新下载的变体上线)。"""
for d in sorted(self.root.iterdir()):
if d.is_dir() and (d / "MODEL_ID.txt").is_file() and (d / "model.safetensors").is_file():
mid, prec = model_identity(d)
if mid not in self.variants:
self.variants[mid] = (d, prec)
if prec == "fp8" and self.default_id is None:
self.default_id = mid
self.version = _version_of(self.default_id) or self.version
self.accepted_ids = ["auto", "laya"] \
+ ([f"laya-{self.version}"] if self.version else []) \
+ sorted(self.variants)
def served_ids(self) -> list[str]:
self.refresh()
return self.accepted_ids
def predict(self, state: Any, questions: dict[str, Any], model: str | None = None,
*, max_len: int | None = None) -> dict[str, Any]:
predictor = self.resolve(model)
return predictor.predict(state, questions, max_len=max_len)
def served_model_ids(model_dir: Path) -> list[str]:
"""create_app 启动时确定对外可用的 model id 列表(单模型 / 多模型目录皆可)。"""
model_dir = Path(model_dir)
if (model_dir / "MODEL_ID.txt").is_file():
mid, prec = model_identity(model_dir)
return _accepted_ids(mid, prec)
ids = ["auto", "laya"]
default_version = None
seen: list[str] = []
for d in sorted(model_dir.iterdir()):
if d.is_dir() and (d / "MODEL_ID.txt").is_file():
mid, prec = model_identity(d)
if mid not in seen:
seen.append(mid)
if prec == "fp8":
v = _version_of(mid)
if v:
default_version = v
ids.append(f"laya-{v}")
return ids + seen