Files
natural-memory-nm21/compact_mega_router_source.py
T
WpyQwq 643e22ecb9 Natural Memory NM2.1: 记忆路由器分叉、数据集缺陷修复与全轴评测证据
- 引入 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,读写关闭时与原生模型逐位相同
2026-09-19 11:11:31 +08:00

79 lines
2.8 KiB
Python

"""Reduce mega memory records to the evidence needed by a router episode.
The original stress file keeps very large long-context fact lists. The
router only needs the gold evidence plus a small local context; global hard
negatives are supplied later by prepare_memory_router_dataset.py.
"""
from __future__ import annotations
import argparse
import json
from collections import Counter
from pathlib import Path
from typing import Any
def clean(value: Any) -> str:
return " ".join(str(value or "").replace("\x00", " ").split()).strip()
def compact(row: dict[str, Any], local_limit: int) -> dict[str, Any]:
facts = [fact for fact in row.get("facts", []) if isinstance(fact, dict) and clean(fact.get("text"))]
acceptable = [clean(value) for value in row.get("acceptable", []) if clean(value)]
positive = [fact for fact in facts if any(value.lower() in clean(fact.get("text")).lower() for value in acceptable)]
category = clean(row.get("category")) or "unknown"
selected: list[dict[str, Any]] = []
seen: set[str] = set()
def add(fact: dict[str, Any]) -> None:
key = clean(fact.get("text"))
if key and key not in seen and len(selected) < local_limit:
selected.append(fact)
seen.add(key)
for fact in positive:
add(fact)
if category == "unknown_abstention":
for fact in facts[:local_limit]:
add(fact)
else:
for fact in facts[:local_limit]:
add(fact)
for fact in reversed(facts[-local_limit:]):
add(fact)
output = dict(row)
output["facts"] = selected
output["metadata"] = {**(row.get("metadata") if isinstance(row.get("metadata"), dict) else {}), "compact_router_source": True}
return output
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--input", required=True)
parser.add_argument("--output", required=True)
parser.add_argument("--local-limit", type=int, default=8)
args = parser.parse_args()
if args.local_limit < 1:
raise SystemExit("local-limit must be positive")
counts: Counter[str] = Counter()
source = Path(args.input)
output = Path(args.output)
output.parent.mkdir(parents=True, exist_ok=True)
with source.open("r", encoding="utf-8") as src, output.open("w", encoding="utf-8") as dst:
for line in src:
line = line.strip()
if not line:
continue
row = compact(json.loads(line), args.local_limit)
counts[clean(row.get("category")) or "unknown"] += 1
dst.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")
print(json.dumps({"input": str(source), "output": str(output), "rows": sum(counts.values()), "categories": dict(counts)}, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()