- 引入 MemoryRouterXL 与 v5/v6 流式多线程训练/编码管线 - 修复 prepare_memory_router_dataset 候选池重建缺陷(mega 家族 3568x 加速,输出逐字节相同) - 修复 v5 被破坏的拒答与多跳标签(train 未知样本 319 -> 16319,multi_hop 平均正例 1.00 -> 2.00) - 同存储预算下 V2-128 v6 逐轴 22/22 通过:Top-1 41.12% -> 94.62%,未知拒答 0.00% -> 100.00% - 记录三条被实测推翻的显然优化(logits_to_keep=1 反而慢 55%、XL 容量未带来收益) - 记忆手术跨架构可移植性 14/14,读写关闭时与原生模型逐位相同
258 lines
12 KiB
Python
258 lines
12 KiB
Python
"""Can "the bank does not contain this attribute" be detected at all?
|
||
|
||
The runtime answers 75% of unanswerable paraphrased questions with an invented code. A
|
||
score threshold was already measured to be useless: the retrieval top-score distributions
|
||
of answerable and unanswerable episodes overlap almost entirely. Before building any
|
||
classifier head, this script asks the prior question properly and cheaply -- **without
|
||
generating text** -- by scoring every candidate of every episode with the same scorers the
|
||
runtime uses and testing how well answerability can be predicted from the resulting score
|
||
geometry (top-1, margin, top-k mean, spread...).
|
||
|
||
Method, so the number is trustworthy:
|
||
|
||
* features come from the frozen feature bank, i.e. exactly the 2560-dim keys the runtime
|
||
scores;
|
||
* a logistic head is **fitted on the train split and evaluated on the eval split**, so the
|
||
reported AUC is not an in-sample artefact;
|
||
* the reported ceiling is the AUC of the single best feature and of the fitted head, plus
|
||
the achievable operating points (unknown-flagging rate vs answerable rejection rate).
|
||
|
||
If the AUC is near 0.5 the conclusion is firm: abstention cannot be recovered from score
|
||
geometry on this feature set, and the fix has to be structural.
|
||
|
||
Usage::
|
||
|
||
python -m V2_dpskw.analyze_abstention_separability --with-text-retriever
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import hashlib
|
||
import json
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn.functional as F
|
||
|
||
from .eval_router_scorecard import load_router_any
|
||
|
||
|
||
def text_key(text: str) -> str:
|
||
return hashlib.sha1(str(text).encode("utf-8", "replace")).hexdigest()
|
||
|
||
|
||
def load_episodes(path: Path) -> list[dict]:
|
||
rows = []
|
||
with path.open("r", encoding="utf-8") as handle:
|
||
for line in handle:
|
||
line = line.strip()
|
||
if line:
|
||
rows.append(json.loads(line))
|
||
return rows
|
||
|
||
|
||
def episode_vectors(row: dict, lookup: dict, bank: np.ndarray):
|
||
query = row["query"]
|
||
candidates = [c["text"] for c in row["candidates"]]
|
||
missing = [t for t in [query, *candidates] if text_key(t) not in lookup]
|
||
if missing:
|
||
raise SystemExit(f"{row['id']}: {len(missing)} texts missing from the bank")
|
||
return bank[lookup[text_key(query)]], bank[[lookup[text_key(t)] for t in candidates]]
|
||
|
||
|
||
def features_from_scores(scores: np.ndarray) -> dict:
|
||
order = np.sort(scores)[::-1]
|
||
top1 = float(order[0])
|
||
top2 = float(order[1]) if order.size > 1 else 0.0
|
||
top3 = float(order[2]) if order.size > 2 else top2
|
||
spread = float(np.std(scores))
|
||
median = float(np.median(scores))
|
||
return {
|
||
"top1": top1,
|
||
"top2": top2,
|
||
"margin": top1 - top2,
|
||
"top3_mean": float(np.mean(order[:3])),
|
||
"top1_minus_median": top1 - median,
|
||
"std": spread,
|
||
"z_top1": (top1 - float(np.mean(scores))) / (spread + 1e-6),
|
||
# how many candidates sit close to the best one: a truly held fact should be
|
||
# clearly ahead of 23 unrelated same-shape facts
|
||
"n_within_10pct": float(np.sum(scores >= top1 - 0.10 * max(abs(top1), 1e-6))),
|
||
"entropy": float(-np.sum(np.exp(scores - top1) / np.sum(np.exp(scores - top1))
|
||
* (scores - top1))),
|
||
}
|
||
|
||
|
||
FEATURE_NAMES = ["top1", "top2", "margin", "top3_mean", "top1_minus_median",
|
||
"std", "z_top1", "n_within_10pct", "entropy"]
|
||
|
||
|
||
def auc(labels: np.ndarray, values: np.ndarray) -> float:
|
||
"""Rank-based AUC (Mann-Whitney), ties averaged."""
|
||
positives = values[labels == 1]
|
||
negatives = values[labels == 0]
|
||
if positives.size == 0 or negatives.size == 0:
|
||
return float("nan")
|
||
order = np.argsort(np.concatenate([positives, negatives]))
|
||
ranks = np.empty(order.size, dtype=float)
|
||
ranks[order] = np.arange(1, order.size + 1)
|
||
# average ties
|
||
combined = np.concatenate([positives, negatives])
|
||
_, inverse, counts = np.unique(combined, return_inverse=True, return_counts=True)
|
||
for index, count in enumerate(counts):
|
||
if count > 1:
|
||
mask = inverse == index
|
||
ranks[mask] = ranks[mask].mean()
|
||
rank_sum = ranks[: positives.size].sum()
|
||
return float((rank_sum - positives.size * (positives.size + 1) / 2)
|
||
/ (positives.size * negatives.size))
|
||
|
||
|
||
def collect(rows, lookup, bank, scorers) -> tuple[dict[str, list[dict]], np.ndarray]:
|
||
per_scorer: dict[str, list[dict]] = {name: [] for name in scorers}
|
||
labels = []
|
||
for row in rows:
|
||
query_vector, candidate_vectors = episode_vectors(row, lookup, bank)
|
||
query = torch.from_numpy(np.ascontiguousarray(query_vector)).float().reshape(1, -1)
|
||
candidates = torch.from_numpy(
|
||
np.ascontiguousarray(candidate_vectors)).float().reshape(1, -1, bank.shape[1])
|
||
labels.append(1 if row.get("positive_indices") else 0)
|
||
with torch.no_grad():
|
||
for name, fn in scorers.items():
|
||
scores = fn(query, candidates).reshape(-1).numpy()
|
||
per_scorer[name].append(features_from_scores(scores))
|
||
return per_scorer, np.array(labels)
|
||
|
||
|
||
def main() -> int:
|
||
parser = argparse.ArgumentParser(description=__doc__)
|
||
parser.add_argument("--data-dir", default="data/zero_overlap")
|
||
parser.add_argument("--feature-cache", default=r"H:\Memory\nm_cache\nm_zero_overlap\feature_cache")
|
||
parser.add_argument("--router", default="checkpoints/router_replay_v7_v2_128/memory_router_v2.pt")
|
||
parser.add_argument("--with-text-retriever", action="store_true")
|
||
parser.add_argument("--package", default="qwen3_5_4b_natural_memory_v2")
|
||
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
|
||
parser.add_argument("--output", default="abstention_separability.json")
|
||
parser.add_argument("--markdown", default="abstention_separability.md")
|
||
args = parser.parse_args()
|
||
|
||
cache = Path(args.feature_cache)
|
||
bank = np.load(cache / "features.f16.npy", mmap_mode="r")
|
||
lookup = json.loads((cache / "index.json").read_text(encoding="utf-8"))
|
||
device = torch.device(args.device)
|
||
|
||
router, _, _ = load_router_any(Path(args.router))
|
||
router = router.to(device).eval()
|
||
|
||
def cosine(query, candidates):
|
||
return F.cosine_similarity(query, candidates.reshape(-1, candidates.shape[-1]), dim=-1)
|
||
|
||
def router_score(query, candidates):
|
||
projected = router.encode_key(candidates.reshape(-1, candidates.shape[-1]).to(device))
|
||
scores, _ = router.projected_scores(
|
||
query.to(device), projected.reshape(1, -1, projected.shape[-1]))
|
||
return torch.sigmoid(scores.reshape(-1)).cpu()
|
||
|
||
scorers = {"cosine": cosine, "router": router_score}
|
||
|
||
retriever = None
|
||
if args.with_text_retriever:
|
||
from .qwen_integration import load_qwen_dynamic
|
||
|
||
model = load_qwen_dynamic(args.package)
|
||
model.eval()
|
||
if model.text_retriever is None or not getattr(model, "_text_retriever_ready", False):
|
||
raise SystemExit("packaged text_retriever is not ready")
|
||
retriever = model.text_retriever.to(device).eval()
|
||
|
||
def retriever_score(query, candidates):
|
||
with torch.no_grad():
|
||
out = retriever(query.to(device),
|
||
candidates.reshape(-1, candidates.shape[-1]).to(device))
|
||
return torch.sigmoid(out.reshape(-1)).cpu()
|
||
|
||
scorers["text_retriever"] = retriever_score
|
||
scorers["blend_50_50"] = lambda q, c: 0.5 * retriever_score(q, c) + 0.5 * router_score(q, c)
|
||
|
||
train_rows = load_episodes(Path(args.data_dir) / "train.jsonl")
|
||
eval_rows = load_episodes(Path(args.data_dir) / "eval.jsonl")
|
||
print(json.dumps({"train_episodes": len(train_rows), "eval_episodes": len(eval_rows),
|
||
"train_unknown": sum(1 for r in train_rows if not r.get("positive_indices")),
|
||
"eval_unknown": sum(1 for r in eval_rows if not r.get("positive_indices")),
|
||
"scorers": list(scorers)}), flush=True)
|
||
|
||
train_feats, train_labels = collect(train_rows, lookup, bank, scorers)
|
||
eval_feats, eval_labels = collect(eval_rows, lookup, bank, scorers)
|
||
|
||
report: dict = {"scorers": {}, "train_episodes": len(train_rows), "eval_episodes": len(eval_rows)}
|
||
|
||
for name in scorers:
|
||
entry: dict = {"single_feature_auc_on_eval": {}}
|
||
for feature in FEATURE_NAMES:
|
||
values = np.array([row[feature] for row in eval_feats[name]], dtype=float)
|
||
entry["single_feature_auc_on_eval"][feature] = round(auc(eval_labels, values), 4)
|
||
|
||
# Logistic head, fitted on train only.
|
||
X_train = np.array([[row[f] for f in FEATURE_NAMES] for row in train_feats[name]], dtype=np.float64)
|
||
X_eval = np.array([[row[f] for f in FEATURE_NAMES] for row in eval_feats[name]], dtype=np.float64)
|
||
mean, std = X_train.mean(axis=0), X_train.std(axis=0) + 1e-9
|
||
Xtr = (X_train - mean) / std
|
||
Xev = (X_eval - mean) / std
|
||
weights = np.zeros(Xtr.shape[1])
|
||
bias = 0.0
|
||
# plain gradient descent with L2; no sklearn dependency
|
||
for _ in range(4000):
|
||
logits = Xtr @ weights + bias
|
||
probs = 1.0 / (1.0 + np.exp(-logits))
|
||
grad_w = Xtr.T @ (probs - train_labels) / len(train_labels) + 1e-3 * weights
|
||
grad_b = float(np.mean(probs - train_labels))
|
||
weights -= 0.5 * grad_w
|
||
bias -= 0.5 * grad_b
|
||
eval_probs = 1.0 / (1.0 + np.exp(-(Xev @ weights + bias)))
|
||
entry["head_auc_on_eval"] = round(auc(eval_labels, eval_probs), 4)
|
||
entry["head_weights"] = {f: round(float(w), 4) for f, w in zip(FEATURE_NAMES, weights)}
|
||
|
||
# Operating points: keep answerable episodes (maximise) while flagging unknowns.
|
||
curve = []
|
||
for threshold in (0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9):
|
||
flagged_unknown = float(np.mean(eval_probs[eval_labels == 0] < threshold))
|
||
rejected_answerable = float(np.mean(eval_probs[eval_labels == 1] < threshold))
|
||
curve.append({
|
||
"threshold": threshold,
|
||
"unknown_flagged_pct": round(100 * flagged_unknown, 2),
|
||
"answerable_rejected_pct": round(100 * rejected_answerable, 2),
|
||
})
|
||
entry["operating_curve"] = curve
|
||
report["scorers"][name] = entry
|
||
|
||
Path(args.output).write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
|
||
lines = ["# 未知问题可分性分析(能否从分数几何判断\"库里没有这个属性\")", "",
|
||
f"训练集 {len(train_rows)} 条(未知 {report['train_episodes'] and sum(1 for r in train_rows if not r.get('positive_indices'))}),"
|
||
f"评测集 {len(eval_rows)} 条(未知 {sum(1 for r in eval_rows if not r.get('positive_indices'))})。",
|
||
"逻辑回归头**只在训练集上拟合**,AUC 在评测集上计算。AUC 0.5 = 完全不可分。", ""]
|
||
for name, entry in report["scorers"].items():
|
||
lines.append(f"## 打分器 `{name}`")
|
||
lines.append("")
|
||
lines.append(f"**单特征最佳 AUC**:" + ", ".join(
|
||
f"{k} {v:.4f}" for k, v in sorted(entry["single_feature_auc_on_eval"].items(),
|
||
key=lambda kv: -kv[1])[:4]))
|
||
lines.append("")
|
||
lines.append(f"**拟合头 AUC(评测集)**:{entry['head_auc_on_eval']:.4f}")
|
||
lines.append("")
|
||
lines.append("| 判定阈值 | 标出未知的比例 | 误拒可回答的比例 |")
|
||
lines.append("|---:|---:|---:|")
|
||
for point in entry["operating_curve"]:
|
||
lines.append("| {threshold:.2f} | {unknown_flagged_pct:.2f}% | {answerable_rejected_pct:.2f}% |".format(**point))
|
||
lines.append("")
|
||
text = "\n".join(lines) + "\n"
|
||
Path(args.markdown).write_text(text, encoding="utf-8")
|
||
print(text)
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|