370953729b
- perf_stats aggregator lives in eval/, not model/: the import failed silently and EVERY perf column was empty (not just ttft). Now warns on stderr instead of swallowing. - repeats > 1 get their own checkpoint key (:rep2, :rep3, ...): repeat 2 previously restored repeat 1's predictions and finished instantly with identical scores. rep1 keeps the legacy key (existing checkpoints still resume). - repeats summary: report the MEAN score and aggregate time/tokens over ALL runs (was: last run only). - README: six-benchmark command as the primary example. Co-Authored-By: Claude <noreply@anthropic.com>
384 lines
17 KiB
Python
384 lines
17 KiB
Python
#!/usr/bin/env python3
|
||
"""FP-Fusion engine: 双池并行调度 + 归一化 + JSD/split-half 统计.
|
||
|
||
并行结构:
|
||
pool D : D 层 420 条单 token 采样 (semaphore=d_concurrency, temperature=1.0)
|
||
pool TXT : 基线 T₀ (20 条极短请求) → I/K/C/S 文本层 36 条 (semaphore=text_concurrency,
|
||
max_tokens=TEXT_MAX_TOKENS 截断)
|
||
两池互不依赖, asyncio.gather 同时跑。
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import math
|
||
import random
|
||
import re
|
||
import time
|
||
|
||
import httpx
|
||
|
||
from .battery import (ALL_CELL_DEFS, BASELINE_MAX_TOKENS, BASELINE_PROMPT,
|
||
BASELINE_SAMPLES, D_SAMPLES_PER_CELL, D_TEMPERATURE,
|
||
REFUSAL_STARTERS, TEXT_MAX_TOKENS, TEXT_TEMPERATURE)
|
||
|
||
# ---------------------------------------------------------------- 归一化 ----
|
||
|
||
_CN_DIGITS = {'零': 0, '一': 1, '二': 2, '两': 2, '三': 3, '四': 4, '五': 5,
|
||
'六': 6, '七': 7, '八': 8, '九': 9, '十': 10}
|
||
_COLOR_EN = {'red': 'red', 'blue': 'blue', 'green': 'green', 'yellow': 'yellow',
|
||
'black': 'black', 'white': 'white', 'purple': 'purple', 'violet': 'purple',
|
||
'orange': 'orange', 'pink': 'pink', 'brown': 'brown', 'gray': 'gray',
|
||
'grey': 'gray'}
|
||
_COIN_MAP = {'heads': 'heads', 'tails': 'tails', '正面': 'heads', '反面': 'tails',
|
||
'head': 'heads', 'tail': 'tails'}
|
||
_WEEKDAY_MAP = {'monday': 'monday', 'tuesday': 'tuesday', 'wednesday': 'wednesday',
|
||
'thursday': 'thursday', 'friday': 'friday', 'saturday': 'saturday',
|
||
'sunday': 'sunday', '周一': 'monday', '星期一': 'monday', '礼拜一': 'monday',
|
||
'周二': 'tuesday', '星期二': 'tuesday', '礼拜二': 'tuesday',
|
||
'周三': 'wednesday', '星期三': 'wednesday', '礼拜三': 'wednesday',
|
||
'周四': 'thursday', '星期四': 'thursday', '礼拜四': 'thursday',
|
||
'周五': 'friday', '星期五': 'friday', '礼拜五': 'friday',
|
||
'周六': 'saturday', '星期六': 'saturday', '礼拜六': 'saturday',
|
||
'周日': 'sunday', '星期日': 'sunday', '星期天': 'sunday',
|
||
'礼拜日': 'sunday', '礼拜天': 'sunday', '周末': 'sunday'}
|
||
|
||
_PUNCT_RE = re.compile(r'[\W_]+', re.UNICODE)
|
||
|
||
|
||
def _first_token(text):
|
||
return text.split()[0] if text.split() else ''
|
||
|
||
|
||
def normalize_answer(raw, domain):
|
||
"""移植 detector normalizer 的主干规则, 返回 (canonical, category)."""
|
||
if raw is None:
|
||
return None, 'error'
|
||
text = raw.strip()
|
||
if not text:
|
||
return None, 'empty'
|
||
low = text.lower()
|
||
if any(s in low for s in REFUSAL_STARTERS):
|
||
return None, 'refusal'
|
||
# NFC + 去标点/emoji + 全角数字转半角
|
||
cleaned = re.sub(r'[\uFF01-\uFF5E]',
|
||
lambda m: chr(ord(m.group(0)) - 0xFEE0), text)
|
||
cleaned = _PUNCT_RE.sub(' ', cleaned).strip().lower()
|
||
if not cleaned:
|
||
return None, 'empty'
|
||
tok = _first_token(cleaned)
|
||
|
||
kind = domain[0]
|
||
if kind == 'int':
|
||
digits = ''.join(ch if ch.isdigit() else str(_CN_DIGITS.get(ch, '')) for ch in tok)
|
||
digits = re.sub(r'\s+', '', digits)
|
||
if digits.isdigit():
|
||
v = int(digits)
|
||
if domain[1] <= v <= domain[2]:
|
||
return str(v), 'valid'
|
||
return tok, 'invalid'
|
||
if kind == 'letter':
|
||
if len(tok) == 1 and tok.isalpha():
|
||
return tok, 'valid'
|
||
m = re.fullmatch(r'[a-z]', tok) or re.match(r'^([a-z])', cleaned)
|
||
if m:
|
||
return m.group(1), 'valid'
|
||
return tok, 'invalid'
|
||
if kind == 'color':
|
||
# 对齐 detector 参考约定: en 归一为英文标准色; zh 保留中文、去掉尾部'色'
|
||
# (参考库实证: random-color:zh keys = {"蓝","蓝紫"}, en = {"blue",...})
|
||
tok = _first_token(cleaned)
|
||
if not tok:
|
||
return None, 'empty'
|
||
if all(ord(ch) < 128 for ch in tok):
|
||
return _COLOR_EN.get(tok, tok), 'valid'
|
||
if len(tok) > 1 and tok.endswith('色'):
|
||
tok = tok[:-1]
|
||
return tok, 'valid'
|
||
if kind == 'coin':
|
||
for k, v in _COIN_MAP.items():
|
||
if k in cleaned:
|
||
return v, 'valid'
|
||
return tok, 'invalid'
|
||
if kind == 'enum':
|
||
for v in domain[1]:
|
||
if v in cleaned:
|
||
return v, 'valid'
|
||
for k, v in _WEEKDAY_MAP.items():
|
||
if k in cleaned:
|
||
return v, 'valid'
|
||
return tok, 'invalid'
|
||
return tok, 'valid' # word 域: 任意词有效
|
||
|
||
|
||
# ---------------------------------------------------------------- 统计 ----
|
||
|
||
def jsd_bits(counts_p, counts_q):
|
||
"""Jensen-Shannon divergence, base 2, 范围 [0,1]."""
|
||
tp, tq = sum(counts_p.values()), sum(counts_q.values())
|
||
if tp <= 0 or tq <= 0:
|
||
return 0.0
|
||
support = set(counts_p) | set(counts_q)
|
||
hm = hp = hq = 0.0
|
||
for k in support:
|
||
p = counts_p.get(k, 0) / tp
|
||
q = counts_q.get(k, 0) / tq
|
||
m = (p + q) / 2
|
||
if m > 0:
|
||
hm -= m * math.log2(m)
|
||
if p > 0:
|
||
hp -= p * math.log2(p)
|
||
if q > 0:
|
||
hq -= q * math.log2(q)
|
||
return min(1.0, max(0.0, hm - (hp + hq) / 2))
|
||
|
||
|
||
def distributions_by_cell(samples):
|
||
"""samples: [{'cell':, 'norm':, 'cat':, 'arrival':}] → {cell: Counter}"""
|
||
out = {}
|
||
for s in samples:
|
||
if s['cat'] == 'valid' and s['norm'] is not None:
|
||
out.setdefault(s['cell'], {})
|
||
out[s['cell']][s['norm']] = out[s['cell']].get(s['norm'], 0) + 1
|
||
return out
|
||
|
||
|
||
def compare_cells(dist_a, dist_b, min_valid=10):
|
||
"""逐 cell JSD(双方 ≥min_valid 才可比), 返回按 JSD 降序的 entries + meanJsd."""
|
||
entries = []
|
||
for cell in sorted(set(dist_a) & set(dist_b)):
|
||
if sum(dist_a[cell].values()) < min_valid or sum(dist_b[cell].values()) < min_valid:
|
||
continue
|
||
entries.append({'cell': cell, 'jsd': jsd_bits(dist_a[cell], dist_b[cell]),
|
||
'valid_a': sum(dist_a[cell].values()),
|
||
'valid_b': sum(dist_b[cell].values())})
|
||
entries.sort(key=lambda e: e['jsd'], reverse=True)
|
||
mean = (sum(e['jsd'] for e in entries) / len(entries)) if entries else None
|
||
return entries, mean
|
||
|
||
|
||
def split_half_jsd(samples, min_per_half=5):
|
||
"""按到达顺序奇偶对半分, 逐 cell JSD 后取平均."""
|
||
halves = {}
|
||
for s in samples:
|
||
if s['cat'] != 'valid' or s['norm'] is None:
|
||
continue
|
||
key = (s['cell'], s['arrival'] % 2)
|
||
halves.setdefault(key, {})
|
||
halves[key][s['norm']] = halves[key].get(s['norm'], 0) + 1
|
||
jsds = []
|
||
for cell in {k[0] for k in halves}:
|
||
even, odd = halves.get((cell, 0)), halves.get((cell, 1))
|
||
if even and odd and sum(even.values()) >= min_per_half and sum(odd.values()) >= min_per_half:
|
||
jsds.append(jsd_bits(even, odd))
|
||
return (sum(jsds) / len(jsds)) if jsds else None
|
||
|
||
|
||
# ---------------------------------------------------------------- 引擎 ----
|
||
|
||
class FusionEngine:
|
||
def __init__(self, api_url, model, timeout=120, d_samples=D_SAMPLES_PER_CELL,
|
||
baseline_samples=BASELINE_SAMPLES, text_limit=0,
|
||
d_concurrency=4, text_concurrency=3,
|
||
text_max_tokens=TEXT_MAX_TOKENS, thinking=False,
|
||
extra_body=None, system_prompt_override=None, logprobs=False,
|
||
temperature_sweep=None, prompt_variants=0, d_cells=None,
|
||
api_key=None):
|
||
self.base = api_url.rstrip('/')
|
||
self.model = model
|
||
self.timeout = timeout
|
||
self.d_samples = d_samples
|
||
self.baseline_samples = baseline_samples
|
||
self.text_limit = text_limit # >0 时只跑前 N 条文本探针(冒烟用)
|
||
self.d_sem = asyncio.Semaphore(d_concurrency)
|
||
self.text_sem = asyncio.Semaphore(text_concurrency)
|
||
self.text_max_tokens = text_max_tokens
|
||
self.extra = {'chat_template_kwargs': {'thinking': thinking}}
|
||
# 维度A/C 扩展：额外请求体字段、全局 system prompt 覆盖（对抗伪装）、
|
||
# logprobs 采集（灰盒预留，一期只采集不评分）
|
||
self.extra_body = dict(extra_body or {})
|
||
self.system_prompt_override = system_prompt_override
|
||
self.logprobs = logprobs
|
||
# 维度D：温度扫描 + D 层 paraphrase 变体（鲁棒性正交）
|
||
self.temperature_sweep = temperature_sweep or []
|
||
self.prompt_variants = max(0, int(prompt_variants or 0))
|
||
# 剪枝落地：D cell 白名单（None=全 26 cell，保持旧行为）
|
||
self.d_cells = set(d_cells) if d_cells else None
|
||
# API 鉴权（vectron 等需 Bearer；None=本地无鉴权端点，同旧行为）
|
||
self.api_key = (api_key or '').strip() or None
|
||
self.records = [] # 所有请求的原始记录
|
||
self.baseline_p50 = None
|
||
self.tokens_in = self.tokens_out = 0
|
||
|
||
# -- 单次请求 ---------------------------------------------------------
|
||
async def _chat(self, sem, user_prompt, max_tokens, temperature, system_prompt,
|
||
tag, rec):
|
||
# 对抗模式：全局覆盖 system prompt（伪装角色）；否则用探针自带
|
||
sys_text = (self.system_prompt_override or '').strip() or system_prompt
|
||
body = {'model': self.model, 'temperature': temperature,
|
||
'max_tokens': max_tokens, 'stream': False,
|
||
'messages': [{'role': 'system', 'content': sys_text},
|
||
{'role': 'user', 'content': user_prompt}],
|
||
**self.extra, **self.extra_body}
|
||
if self.logprobs:
|
||
body['logprobs'] = True
|
||
body.setdefault('top_logprobs', 5)
|
||
headers = {'Content-Type': 'application/json'}
|
||
if self.api_key:
|
||
headers['Authorization'] = f'Bearer {self.api_key}'
|
||
t0 = time.perf_counter()
|
||
err = None
|
||
content, usage = None, None
|
||
top_logprobs = None
|
||
# 慢端点排队不可控: 所有请求统一放宽下限 300s(含 D 层单 token 请求)
|
||
eff_timeout = max(self.timeout, 300)
|
||
try:
|
||
async with sem:
|
||
async with httpx.AsyncClient(timeout=eff_timeout) as client:
|
||
r = await client.post(f'{self.base}/chat/completions',
|
||
json=body, headers=headers)
|
||
r.raise_for_status()
|
||
data = r.json()
|
||
content = data['choices'][0]['message'].get('content') or ''
|
||
usage = data.get('usage') or {}
|
||
# 灰盒预留：记录首个完成 token 的 top-logprobs（供 B 期概率校准）
|
||
if self.logprobs:
|
||
lp = (data.get('choices') or [{}])[0].get('logprobs') or {}
|
||
tokens = lp.get('content') or []
|
||
top_logprobs = [
|
||
{'token': t.get('token'), 'logprob': t.get('logprob'),
|
||
'top': t.get('top_logprobs')}
|
||
for t in tokens[:16]
|
||
] if isinstance(tokens, list) else None
|
||
except Exception as e:
|
||
# httpx 超时异常的 str() 可能为空串 → 永远保证 error 非空(fail-closed 依赖它)
|
||
err = str(e) or type(e).__name__
|
||
latency = (time.perf_counter() - t0) * 1000
|
||
if usage:
|
||
self.tokens_in += usage.get('prompt_tokens') or 0
|
||
self.tokens_out += usage.get('completion_tokens') or 0
|
||
rec.update({'latency_ms': round(latency, 1), 'response': content,
|
||
'error': err, 'tag': tag,
|
||
'completion_tokens': (usage or {}).get('completion_tokens'),
|
||
'top_logprobs': top_logprobs})
|
||
self.records.append(rec)
|
||
return rec
|
||
|
||
SYS_ONE_WORD = {'en': 'Answer with exactly one word. No punctuation, no explanation.',
|
||
'zh': '只回答一个词，不要标点，不要解释。'}
|
||
SYS_DEFAULT = 'You are a helpful assistant.'
|
||
|
||
# -- 基线 --------------------------------------------------------------
|
||
async def _run_baseline(self):
|
||
sem = asyncio.Semaphore(3)
|
||
tasks = [self._chat(sem, BASELINE_PROMPT, BASELINE_MAX_TOKENS, 0.0,
|
||
self.SYS_DEFAULT, 'baseline',
|
||
{'id': f'baseline_{i}', 'layer': 'BASE'})
|
||
for i in range(self.baseline_samples)]
|
||
done = [r for r in await asyncio.gather(*tasks) if not r['error']]
|
||
lats = sorted(r['latency_ms'] for r in done)
|
||
self.baseline_p50 = lats[len(lats) // 2] if lats else None
|
||
|
||
# -- D 层 ---------------------------------------------------------------
|
||
def _d_jobs(self):
|
||
jobs = []
|
||
variants = self.prompt_variants if self.prompt_variants > 0 else None
|
||
for c in ALL_CELL_DEFS:
|
||
for lang in ('en', 'zh'):
|
||
cell_id = f"{c['id']}:{lang}"
|
||
if self.d_cells is not None and cell_id not in self.d_cells:
|
||
continue
|
||
pool = c['par'][lang]
|
||
for i in range(self.d_samples):
|
||
# 维度D 改写轴：轮换使用池中前 N 个 paraphrase；否则随机
|
||
if variants:
|
||
prompt = pool[i % min(variants, len(pool))]
|
||
else:
|
||
prompt = random.choice(pool)
|
||
jobs.append((cell_id, c['domain'], prompt, lang, prompt))
|
||
random.shuffle(jobs) # 防单 cell 突发(触发缓存/限流偏差)
|
||
# 统一为 5 元组：cell, domain, prompt, lang, prompt_var(原样副本)
|
||
return [(j[0], j[1], j[2], j[3], j[4]) for j in jobs]
|
||
|
||
async def _run_d_layer(self):
|
||
sem = self.d_sem
|
||
tasks = []
|
||
for idx, (cell_id, domain, prompt, lang, prompt_var) in enumerate(self._d_jobs()):
|
||
rec = {'id': f'd_{cell_id}_{idx}', 'layer': 'D', 'cell': cell_id,
|
||
'lang': lang, 'prompt': prompt, 'prompt_var': prompt_var,
|
||
'arrival': idx}
|
||
tasks.append(self._chat(sem, prompt, 16, D_TEMPERATURE,
|
||
self.SYS_ONE_WORD[lang], 'dist', rec))
|
||
await asyncio.gather(*tasks)
|
||
|
||
# -- 文本层 -------------------------------------------------------------
|
||
async def _run_text_layer(self, probes):
|
||
sem = self.text_sem
|
||
tasks = []
|
||
temps = self.temperature_sweep or [TEXT_TEMPERATURE]
|
||
for p in probes:
|
||
# 探针自带 temperature（如 V 层确定性探针用 0.0）优先；
|
||
# 否则 sweep 多温度；再否则默认 TEXT_TEMPERATURE。
|
||
p_meta = p.get('meta') or {}
|
||
probe_temp = p_meta.get('temperature')
|
||
per_temps = [probe_temp] if probe_temp is not None else temps
|
||
for t in per_temps:
|
||
rec = {'id': p['id'], 'layer': p['layer'], 'prompt': p['text'],
|
||
'temperature': t,
|
||
'meta': {k: v for k, v in p.items()
|
||
if k in ('pair', 'lang', 'metacog', 'refusal_grad',
|
||
'len_ctrl', 'role', 'expect_family')}}
|
||
tasks.append(self._chat(sem, p['text'], self.text_max_tokens,
|
||
t, self.SYS_DEFAULT, 'text', rec))
|
||
await asyncio.gather(*tasks)
|
||
|
||
# -- 主入口 --------------------------------------------------------------
|
||
async def run(self, text_probes):
|
||
if self.text_limit > 0:
|
||
text_probes = text_probes[:self.text_limit]
|
||
# 基线必须独占测量: 若与 D 池并发, CPU 端点的基线会被排队延迟污染
|
||
await self._run_baseline()
|
||
await asyncio.gather(self._run_d_layer(), self._run_text_layer(text_probes))
|
||
return self.records
|
||
|
||
# -- 结果整理 -------------------------------------------------------------
|
||
def d_samples_normalized(self):
|
||
"""返回 [{cell, norm, cat, arrival}]"""
|
||
out = []
|
||
for r in self.records:
|
||
if r.get('tag', r.get('layer')) != 'dist':
|
||
continue
|
||
# _chat 不知 domain; 由调用方(cell)反查 —— 在 run_fp_fusion 里完成
|
||
out.append(r)
|
||
return out
|
||
|
||
|
||
def build_d_normalized(records):
|
||
"""records(D层) → 归一化样本列表"""
|
||
from .battery import ALL_CELL_DEFS
|
||
dom = {}
|
||
for c in ALL_CELL_DEFS:
|
||
dom[f'{c["id"]}:en'] = c['domain']
|
||
dom[f'{c["id"]}:zh'] = c['domain']
|
||
out = []
|
||
for r in records:
|
||
if r['layer'] != 'D':
|
||
continue
|
||
norm, cat = (None, 'error') if r['error'] else normalize_answer(
|
||
r.get('response'), dom[r['cell']])
|
||
out.append({'cell': r['cell'], 'norm': norm, 'cat': cat,
|
||
'arrival': r['arrival'], 'error': r['error']})
|
||
return out
|
||
|
||
|
||
def load_reference(path):
|
||
"""加载 detector schema 的单模型参考指纹 → {cellId: counts}"""
|
||
with open(path, encoding='utf-8') as f:
|
||
ref = json.load(f)
|
||
if ref.get('formatVersion') != 1 or not isinstance(ref.get('cells'), dict):
|
||
raise ValueError(f'unsupported reference format: {path}')
|
||
cells = {}
|
||
for cid, c in ref['cells'].items():
|
||
counts = c.get('counts', {})
|
||
cells[cid] = {str(k): v for k, v in counts.items()}
|
||
return {'model': ref.get('model'), 'cells': cells}
|