WpyQwq
643e22ecb9
- 引入 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，读写关闭时与原生模型逐位相同
879 lines
34 KiB
Python
879 lines
34 KiB
Python
"""Run large-scale natural-language and real-repository memory benchmarks.
|
||
|
||
The benchmark keeps the memory bank in process RAM and caps the model process
|
||
at a declared HBM budget. It measures two separate workloads:
|
||
|
||
* naturalistic multi-turn memory episodes from the local native-memory corpus;
|
||
* real source and documentation from this repository, expanded to a large
|
||
page library and queried through actual file/symbol questions.
|
||
|
||
The primary outputs are peak VRAM and correctness. Retrieval is evaluated on
|
||
every query, while generation is evaluated on a substantial holdout subset so
|
||
that a high retrieval score cannot be mistaken for end-to-end chat quality.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import gc
|
||
import json
|
||
import random
|
||
import re
|
||
import time
|
||
from collections import Counter
|
||
from pathlib import Path
|
||
from statistics import mean, median
|
||
from typing import Any, Iterable
|
||
|
||
import torch
|
||
import torch.nn.functional as F
|
||
|
||
from .qwen_integration import format_memory_evidence, load_qwen_dynamic, load_tokenizer
|
||
from .stream_chat_qwen_memory import _chat_tensor, _memory_system_prefix
|
||
|
||
|
||
PROJECT_ROOT = Path(__file__).resolve().parent
|
||
DEFAULT_MEMORY_MODEL = PROJECT_ROOT / "qwen3_5_4b_natural_memory_v2"
|
||
REFUSAL_MARKERS = ("不知道", "没有记录", "无相关", "未找到", "不清楚", "无法确认")
|
||
|
||
|
||
def _path(value: str | Path) -> Path:
|
||
path = Path(value)
|
||
return path if path.is_absolute() or path.exists() else PROJECT_ROOT / path
|
||
|
||
|
||
def _max_memory(gpu_memory_gb: float) -> dict[Any, str] | None:
|
||
if gpu_memory_gb <= 0.0 or not torch.cuda.is_available():
|
||
return None
|
||
return {0: f"{gpu_memory_gb:.1f}GiB", "cpu": "64GiB"}
|
||
|
||
|
||
def _set_cuda_process_cap(gpu_memory_gb: float) -> None:
|
||
if gpu_memory_gb <= 0.0 or not torch.cuda.is_available():
|
||
return
|
||
total = torch.cuda.get_device_properties(0).total_memory
|
||
fraction = min(0.95, max(0.05, gpu_memory_gb * 1024**3 / total))
|
||
torch.cuda.set_per_process_memory_fraction(fraction, device=0)
|
||
|
||
|
||
def _sync(device: torch.device) -> None:
|
||
if device.type == "cuda":
|
||
torch.cuda.synchronize(device)
|
||
|
||
|
||
def _encode_texts(
|
||
model: Any,
|
||
tokenizer: Any,
|
||
texts: list[str],
|
||
device: torch.device,
|
||
*,
|
||
batch_size: int = 16,
|
||
) -> torch.Tensor:
|
||
outputs: list[torch.Tensor] = []
|
||
for start in range(0, len(texts), max(1, batch_size)):
|
||
batch = texts[start : start + max(1, batch_size)]
|
||
encoded = tokenizer(
|
||
batch,
|
||
add_special_tokens=False,
|
||
padding=True,
|
||
return_tensors="pt",
|
||
)
|
||
input_ids = encoded["input_ids"].to(device)
|
||
mask = encoded["attention_mask"].to(device)
|
||
outputs.append(model._encode_model_key(input_ids, mask).detach().cpu())
|
||
if not outputs:
|
||
return torch.empty((0, model.memory.hidden_size), dtype=torch.float32)
|
||
return torch.cat(outputs, dim=0)
|
||
|
||
|
||
def _query_token_ids(tokenizer: Any, text: str) -> torch.Tensor:
|
||
return torch.tensor(
|
||
tokenizer(text, add_special_tokens=False)["input_ids"],
|
||
dtype=torch.long,
|
||
)
|
||
|
||
|
||
def _normal(text: Any) -> str:
|
||
return re.sub(r"\s+", "", str(text)).lower()
|
||
|
||
|
||
def _contains_answer(response: str, expected: str) -> bool:
|
||
expected_normal = _normal(expected)
|
||
return bool(expected_normal) and expected_normal in _normal(response)
|
||
|
||
|
||
def _is_refusal(response: str) -> bool:
|
||
"""Recognize concise and natural-language abstentions."""
|
||
|
||
normalized = _normal(response)
|
||
if any(marker in normalized for marker in REFUSAL_MARKERS):
|
||
return True
|
||
return bool(
|
||
re.search(
|
||
r"(没有|无|未|不包含|无法).{0,80}(记录|资料|信息|数据|找到|知道|访问|交互)",
|
||
normalized,
|
||
)
|
||
)
|
||
|
||
|
||
def _quantiles(values: Iterable[float]) -> dict[str, float | None]:
|
||
items = sorted(float(value) for value in values)
|
||
if not items:
|
||
return {"mean": None, "median": None, "p95": None, "max": None}
|
||
index = min(len(items) - 1, max(0, int(round(0.95 * (len(items) - 1)))))
|
||
return {
|
||
"mean": mean(items),
|
||
"median": median(items),
|
||
"p95": items[index],
|
||
"max": items[-1],
|
||
}
|
||
|
||
|
||
def _load_jsonl(path: Path, limit: int | None = None) -> list[dict[str, Any]]:
|
||
rows: list[dict[str, Any]] = []
|
||
with path.open("r", encoding="utf-8") as handle:
|
||
for line in handle:
|
||
if not line.strip():
|
||
continue
|
||
rows.append(json.loads(line))
|
||
if limit is not None and len(rows) >= limit:
|
||
break
|
||
return rows
|
||
|
||
|
||
def _native_cases(data_root: Path, max_cases: int) -> list[dict[str, Any]]:
|
||
"""Load train/eval native-memory episodes without loading large files."""
|
||
|
||
rows: list[dict[str, Any]] = []
|
||
sources = (
|
||
("train", data_root / "native_memory" / "train.jsonl"),
|
||
("eval", data_root / "native_memory" / "eval.jsonl"),
|
||
)
|
||
remaining = max(1, max_cases)
|
||
for split, path in sources:
|
||
if not path.exists() or remaining <= 0:
|
||
continue
|
||
loaded = _load_jsonl(path, remaining)
|
||
for row in loaded:
|
||
rows.append({"split": split, **row})
|
||
remaining -= len(loaded)
|
||
if not rows:
|
||
raise FileNotFoundError(f"native memory corpus not found under {data_root}")
|
||
return rows
|
||
|
||
|
||
def _fact_text(chunk: dict[str, Any]) -> str:
|
||
messages = chunk.get("messages") or []
|
||
for message in messages:
|
||
if isinstance(message, dict) and message.get("role") == "user":
|
||
return str(message.get("content", ""))
|
||
return str(chunk.get("text", ""))
|
||
|
||
|
||
def _build_general_records(
|
||
cases: list[dict[str, Any]],
|
||
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
|
||
"""Convert local conversation episodes into versioned memory records."""
|
||
|
||
records: list[dict[str, Any]] = []
|
||
query_rows: list[dict[str, Any]] = []
|
||
for case in cases:
|
||
subject = str(case.get("subject", ""))
|
||
attribute = str(case.get("attribute", ""))
|
||
chunks = case.get("memory_chunks") or []
|
||
for chunk in chunks:
|
||
if not isinstance(chunk, dict) or float(chunk.get("write_label", 0.0)) < 0.5:
|
||
continue
|
||
text = _fact_text(chunk).strip()
|
||
value = str(chunk.get("value", ""))
|
||
if not text or not value:
|
||
continue
|
||
records.append(
|
||
{
|
||
"text": text,
|
||
"entity": subject,
|
||
"attribute": attribute,
|
||
"value": value,
|
||
"memory_type": "personal_fact",
|
||
"importance": 0.9,
|
||
"confidence": 0.95,
|
||
"source": "native_memory_corpus",
|
||
"trusted": True,
|
||
"force": True,
|
||
}
|
||
)
|
||
query_messages = case.get("query") or []
|
||
query = ""
|
||
for message in query_messages:
|
||
if isinstance(message, dict) and message.get("role") == "user":
|
||
query = str(message.get("content", ""))
|
||
break
|
||
if not query:
|
||
query = f"请查询 {subject} 的 {attribute}。"
|
||
answer = str(case.get("answer", ""))
|
||
answerable = bool(case.get("answerable", False)) and answer not in REFUSAL_MARKERS
|
||
query_rows.append(
|
||
{
|
||
"id": str(case.get("id", "")),
|
||
"split": str(case.get("split", "unknown")),
|
||
"query": query,
|
||
"subject": subject,
|
||
"attribute": attribute,
|
||
"expected": answer,
|
||
"answerable": answerable,
|
||
}
|
||
)
|
||
return records, query_rows
|
||
|
||
|
||
def _source_files() -> list[Path]:
|
||
allowed = {".py", ".md", ".json"}
|
||
files: list[Path] = []
|
||
for path in PROJECT_ROOT.rglob("*"):
|
||
if not path.is_file() or path.suffix.lower() not in allowed:
|
||
continue
|
||
if "__pycache__" in path.parts or "checkpoints" in path.parts:
|
||
continue
|
||
if path.name.endswith(".safetensors"):
|
||
continue
|
||
files.append(path)
|
||
return sorted(files)
|
||
|
||
|
||
def _source_chunks(files: list[Path], *, chars_per_chunk: int = 1800) -> list[str]:
|
||
chunks: list[str] = []
|
||
for path in files:
|
||
text = path.read_text(encoding="utf-8", errors="replace")
|
||
if not text:
|
||
continue
|
||
for start in range(0, len(text), chars_per_chunk):
|
||
piece = text[start : start + chars_per_chunk].strip()
|
||
if piece:
|
||
chunks.append(piece)
|
||
if not chunks:
|
||
raise RuntimeError("repository source corpus is empty")
|
||
return chunks
|
||
|
||
|
||
def _project_targets(files: list[Path], limit: int) -> list[dict[str, Any]]:
|
||
pattern = re.compile(r"^\s*(class|async\s+def|def)\s+([A-Za-z_][A-Za-z0-9_]*)")
|
||
candidates: list[dict[str, Any]] = []
|
||
for path in files:
|
||
if path.suffix.lower() != ".py":
|
||
continue
|
||
relative = path.relative_to(PROJECT_ROOT).as_posix()
|
||
for line_number, line in enumerate(path.read_text(encoding="utf-8", errors="replace").splitlines(), 1):
|
||
match = pattern.match(line)
|
||
if match is None:
|
||
continue
|
||
kind = match.group(1).replace("async ", "")
|
||
name = match.group(2)
|
||
text = (
|
||
f"真实代码库事实：文件 {relative} 的第 {line_number} 行定义了 "
|
||
f"{kind} {name}。原始代码行：{line.strip()}"
|
||
)
|
||
candidates.append(
|
||
{
|
||
"file": relative,
|
||
"line": line_number,
|
||
"kind": kind,
|
||
"name": name,
|
||
"text": text,
|
||
"expected": relative,
|
||
}
|
||
)
|
||
if not candidates:
|
||
raise RuntimeError("no Python symbols found in repository corpus")
|
||
if len(candidates) <= limit:
|
||
return candidates
|
||
# Evenly sample the repository instead of measuring only the first file.
|
||
indices = [int(index * len(candidates) / limit) for index in range(limit)]
|
||
return [candidates[index] for index in indices]
|
||
|
||
|
||
def _build_project_records(
|
||
files: list[Path],
|
||
*,
|
||
record_count: int,
|
||
target_count: int,
|
||
chunk_tokens: int,
|
||
) -> tuple[list[dict[str, Any]], list[dict[str, Any]], dict[str, Any]]:
|
||
chunks = _source_chunks(files)
|
||
targets = _project_targets(files, target_count)
|
||
rng = random.Random(20260905)
|
||
positions = rng.sample(range(record_count), min(len(targets), record_count))
|
||
position_to_target = dict(zip(positions, targets))
|
||
records: list[dict[str, Any]] = []
|
||
for index in range(record_count):
|
||
target = position_to_target.get(index)
|
||
if target is None:
|
||
text = (
|
||
f"项目源码快照分片 {index}：\n"
|
||
f"{chunks[index % len(chunks)]}"
|
||
)
|
||
records.append(
|
||
{
|
||
"text": text,
|
||
"key_kind": "random_filler",
|
||
"memory_type": "repository_chunk",
|
||
"importance": 0.4,
|
||
"confidence": 0.8,
|
||
"source": "real_repository_snapshot",
|
||
"trusted": True,
|
||
"force": True,
|
||
}
|
||
)
|
||
continue
|
||
records.append(
|
||
{
|
||
"text": target["text"],
|
||
"entity": target["file"],
|
||
"attribute": f"symbol:{target['name']}",
|
||
"value": target["file"],
|
||
"memory_type": "repository_symbol",
|
||
"importance": 0.9,
|
||
"confidence": 0.95,
|
||
"source": "real_repository_snapshot",
|
||
"target": target,
|
||
"trusted": True,
|
||
"force": True,
|
||
}
|
||
)
|
||
queries: list[dict[str, Any]] = []
|
||
for target in targets:
|
||
queries.extend(
|
||
[
|
||
{
|
||
"query": f"在真实代码库中，文件 {target['file']} 里的 {target['name']} 定义在哪个文件？",
|
||
"expected": target["expected"],
|
||
"target": target,
|
||
},
|
||
{
|
||
"query": f"请从项目记忆查找：{target['file']} 的 {target['kind']} {target['name']} 位于哪里？",
|
||
"expected": target["expected"],
|
||
"target": target,
|
||
},
|
||
]
|
||
)
|
||
return records, queries, {
|
||
"source_file_count": len(files),
|
||
"source_characters": sum(path.stat().st_size for path in files),
|
||
"source_chunk_count": len(chunks),
|
||
"capacity_tokens": int(record_count * chunk_tokens),
|
||
"chunk_tokens": int(chunk_tokens),
|
||
"target_count": len(targets),
|
||
}
|
||
|
||
|
||
def _prepare_records(
|
||
model: Any,
|
||
tokenizer: Any,
|
||
records: list[dict[str, Any]],
|
||
device: torch.device,
|
||
*,
|
||
batch_size: int,
|
||
) -> list[dict[str, Any]]:
|
||
semantic_indices = [
|
||
index for index, record in enumerate(records)
|
||
if record.get("key_kind") != "random_filler"
|
||
]
|
||
semantic_keys = _encode_texts(
|
||
model,
|
||
tokenizer,
|
||
[str(records[index]["text"]) for index in semantic_indices],
|
||
device,
|
||
batch_size=batch_size,
|
||
)
|
||
key_by_index = {
|
||
index: semantic_keys[position]
|
||
for position, index in enumerate(semantic_indices)
|
||
}
|
||
prepared: list[dict[str, Any]] = []
|
||
for index, record in enumerate(records):
|
||
item = {key: value for key, value in record.items() if key not in {"target", "key_kind"}}
|
||
if record.get("key_kind") == "random_filler":
|
||
generator = torch.Generator(device="cpu").manual_seed(900000 + index)
|
||
key = F.normalize(torch.randn(model.memory.hidden_size, generator=generator), dim=0)
|
||
semantic_key = None
|
||
else:
|
||
key = F.normalize(key_by_index[index].float(), dim=0)
|
||
semantic_key = key
|
||
item["key"] = key
|
||
item["summary"] = key
|
||
# Background corpus chunks are intentionally not semantic candidates
|
||
# in this bounded stress run. Leaving their semantic key unset lets
|
||
# the bank use the learned reranker for real evidence while retaining
|
||
# the chunks as corpus noise rather than pretending their random
|
||
# addresses were learned representations.
|
||
item["semantic_key"] = semantic_key
|
||
if record.get("key_kind") != "random_filler":
|
||
# Match the production writer: retrieved memories enter Qwen as
|
||
# an internal system-message prefix, while record.text remains
|
||
# the raw fact used by the address/routing layer.
|
||
evidence_text = format_memory_evidence(
|
||
str(record["text"]),
|
||
entity=str(record.get("entity", "")),
|
||
attribute=str(record.get("attribute", "")),
|
||
value=str(record.get("value", "")),
|
||
)
|
||
storage = _memory_system_prefix(
|
||
tokenizer,
|
||
"以下是与当前用户相关的已保存长期记忆。仅在问题相关时使用，"
|
||
"只能依据明确证据；先核对实体、属性和已确认值；冲突优先最新可靠来源，"
|
||
"不要拼接不确定候选，证据不足就明确说不知道；涉及名称、路径、token、"
|
||
"参数或结论时，原样复述证据中的关键短语：\n"
|
||
+ evidence_text,
|
||
)
|
||
token_ids = storage["input_ids"].reshape(-1).to(dtype=torch.long)
|
||
token_mask = storage["attention_mask"].reshape(-1).to(dtype=torch.bool)
|
||
item["token_ids"] = token_ids
|
||
item["token_mask"] = token_mask
|
||
prepared.append(item)
|
||
return prepared
|
||
|
||
|
||
def _clear_memory_bank(model: Any, device: torch.device) -> None:
|
||
model.memory_os_v2 = model._new_memory_os_v2(model.memory.hidden_size)
|
||
model.reset_memory(batch_size=1, device=device)
|
||
|
||
|
||
def _direct_read(
|
||
model: Any,
|
||
tokenizer: Any,
|
||
query_key: torch.Tensor,
|
||
row: dict[str, Any],
|
||
device: torch.device,
|
||
) -> dict[str, Any]:
|
||
query = str(row["query"])
|
||
token_ids = _query_token_ids(tokenizer, query)
|
||
_sync(device)
|
||
started = time.perf_counter()
|
||
records, decision = model.memory_os_v2.read(
|
||
query_key=query_key,
|
||
query_text=query,
|
||
query_token_ids=token_ids,
|
||
top_k_pages=4,
|
||
top_k_records=2,
|
||
max_hops=3,
|
||
)
|
||
_sync(device)
|
||
elapsed_ms = (time.perf_counter() - started) * 1000.0
|
||
returned_values = [record.value for record in records]
|
||
if row.get("answerable", True):
|
||
hit = any(_contains_answer(value, str(row.get("expected", ""))) for value in returned_values)
|
||
else:
|
||
hit = len(records) == 0
|
||
return {
|
||
"id": row.get("id", ""),
|
||
"query": query,
|
||
"expected": row.get("expected", ""),
|
||
"answerable": bool(row.get("answerable", True)),
|
||
"hit": bool(hit),
|
||
"returned_values": returned_values,
|
||
"returned_ids": [record.record_id for record in records],
|
||
"stop_reason": decision.stop_reason,
|
||
"coarse_candidates": model.memory_os_v2.bank._last_coarse_candidates,
|
||
"elapsed_ms": elapsed_ms,
|
||
}
|
||
|
||
|
||
def _record_audit_view(record: Any) -> dict[str, Any]:
|
||
"""Keep generation diagnostics readable without serializing embeddings."""
|
||
|
||
return {
|
||
"record_id": str(getattr(record, "record_id", "")),
|
||
"entity": str(getattr(record, "entity", "")),
|
||
"attribute": str(getattr(record, "attribute", "")),
|
||
"value": str(getattr(record, "value", "")),
|
||
"status": str(getattr(record, "status", "")),
|
||
"version": int(getattr(record, "version", 0)),
|
||
"source": str(getattr(record, "source", "")),
|
||
"text_preview": str(getattr(record, "text", ""))[:300],
|
||
}
|
||
|
||
|
||
def _chat_generate(
|
||
model: Any,
|
||
tokenizer: Any,
|
||
query: str,
|
||
device: torch.device,
|
||
max_new_tokens: int,
|
||
) -> dict[str, Any]:
|
||
encoded = _chat_tensor(
|
||
tokenizer,
|
||
query,
|
||
)
|
||
encoded = {key: value.to(device) for key, value in encoded.items()}
|
||
raw = tokenizer(query, add_special_tokens=False, return_tensors="pt")
|
||
query_ids = raw["input_ids"].to(device)
|
||
query_mask = raw.get("attention_mask")
|
||
if query_mask is None:
|
||
query_mask = torch.ones_like(query_ids)
|
||
_sync(device)
|
||
generation_started = time.perf_counter()
|
||
output = model.generate(
|
||
**encoded,
|
||
max_new_tokens=max_new_tokens,
|
||
do_sample=False,
|
||
update_memory=False,
|
||
memory_query_input_ids=query_ids,
|
||
memory_query_attention_mask=query_mask.to(device),
|
||
memory_query_text=query,
|
||
use_cache=True,
|
||
pad_token_id=tokenizer.pad_token_id,
|
||
)
|
||
_sync(device)
|
||
generation_seconds = time.perf_counter() - generation_started
|
||
response_ids = output[0, encoded["input_ids"].shape[1] :]
|
||
response = tokenizer.decode(response_ids.detach().cpu().tolist(), skip_special_tokens=True).strip()
|
||
|
||
runtime = getattr(model, "runtime", None)
|
||
raw_decisions = list(getattr(runtime, "v2_last_decisions", []) or [])
|
||
decisions: list[dict[str, Any]] = []
|
||
selected_record_ids: list[str] = []
|
||
for raw_decision in raw_decisions:
|
||
decision = dict(raw_decision)
|
||
record_ids = [str(item) for item in decision.get("record_ids", [])]
|
||
selected_record_ids.extend(record_ids)
|
||
decisions.append(
|
||
{
|
||
"need_memory": bool(decision.get("need_memory", False)),
|
||
"record_ids": record_ids,
|
||
"page_ids": [str(item) for item in decision.get("page_ids", [])],
|
||
"hop_count": int(decision.get("hop_count", 0)),
|
||
"confidence": float(decision.get("confidence", 0.0)),
|
||
"top_score": float(decision.get("top_score", 0.0)),
|
||
"score_margin": float(decision.get("score_margin", 0.0)),
|
||
"evidence_score": float(decision.get("evidence_score", 0.0)),
|
||
"stop_reason": str(decision.get("stop_reason", "")),
|
||
}
|
||
)
|
||
selected_record_ids = list(dict.fromkeys(selected_record_ids))
|
||
records_by_id = getattr(getattr(model, "memory_os_v2", None), "bank", None)
|
||
records_by_id = getattr(records_by_id, "records", {})
|
||
selected_records = [
|
||
_record_audit_view(records_by_id[record_id])
|
||
for record_id in selected_record_ids
|
||
if record_id in records_by_id
|
||
]
|
||
return {
|
||
"response": response,
|
||
"public_prompt_tokens": int(encoded["input_ids"].shape[1]),
|
||
"generated_tokens": int(response_ids.numel()),
|
||
"generation_seconds": float(generation_seconds),
|
||
"prefix_used": bool(getattr(runtime, "text_prefix_used", False)),
|
||
"guard_used": bool(getattr(runtime, "text_guard_used", False)),
|
||
"prefix_tokens": int(getattr(runtime, "text_prefix_tokens", 0)),
|
||
"memory_read_seconds": float(getattr(runtime, "text_read_seconds", 0.0)),
|
||
"selected_record_ids": selected_record_ids,
|
||
"selected_records": selected_records,
|
||
"selected_values": [record["value"] for record in selected_records],
|
||
"decisions": decisions,
|
||
"stop_reasons": [decision["stop_reason"] for decision in decisions],
|
||
}
|
||
|
||
|
||
def _generation_eval(
|
||
model: Any,
|
||
tokenizer: Any,
|
||
rows: list[dict[str, Any]],
|
||
device: torch.device,
|
||
*,
|
||
max_new_tokens: int,
|
||
reference_retrieval_rows: dict[str, dict[str, Any]] | None = None,
|
||
) -> dict[str, Any]:
|
||
output_rows: list[dict[str, Any]] = []
|
||
for row in rows:
|
||
generated = _chat_generate(model, tokenizer, str(row["query"]), device, max_new_tokens)
|
||
response = str(generated["response"])
|
||
answerable = bool(row.get("answerable", True))
|
||
expected = str(row.get("expected", ""))
|
||
retrieved_target = answerable and any(
|
||
_contains_answer(str(record.get("value", "")), expected)
|
||
or _contains_answer(str(record.get("text_preview", "")), expected)
|
||
for record in generated["selected_records"]
|
||
)
|
||
reference = (reference_retrieval_rows or {}).get(str(row.get("id", "")))
|
||
reference_target = bool(reference and reference.get("hit", False))
|
||
passed = (
|
||
_contains_answer(response, expected)
|
||
if answerable
|
||
else _is_refusal(response)
|
||
)
|
||
if answerable and passed:
|
||
error_class = "correct"
|
||
elif not answerable and passed:
|
||
error_class = "correct_refusal"
|
||
elif not answerable:
|
||
error_class = "refusal_failure"
|
||
elif reference_target and not retrieved_target:
|
||
error_class = "generation_retrieval_divergence"
|
||
elif not generated["selected_record_ids"]:
|
||
error_class = "no_memory_prefix"
|
||
elif not retrieved_target:
|
||
error_class = "retrieval_miss_or_wrong_prefix"
|
||
else:
|
||
error_class = "correct_evidence_ignored_or_overridden"
|
||
output_rows.append(
|
||
{
|
||
"id": row.get("id", ""),
|
||
"query": row["query"],
|
||
"expected": expected,
|
||
"response": response,
|
||
"answerable": answerable,
|
||
"correct": bool(passed),
|
||
"retrieval_target_found": bool(retrieved_target),
|
||
"reference_retrieval_hit": reference_target,
|
||
"reference_retrieved_values": list(reference.get("returned_values", [])) if reference else [],
|
||
"retrieved_values": list(generated["selected_values"]),
|
||
"retrieved_ids": list(generated["selected_record_ids"]),
|
||
"selected_records": list(generated["selected_records"]),
|
||
"prefix_used": bool(generated["prefix_used"]),
|
||
"prefix_tokens": int(generated["prefix_tokens"]),
|
||
"memory_read_seconds": float(generated["memory_read_seconds"]),
|
||
"decisions": list(generated["decisions"]),
|
||
"stop_reasons": list(generated["stop_reasons"]),
|
||
"error_class": error_class,
|
||
}
|
||
)
|
||
answerable_rows = [row for row in output_rows if row["answerable"]]
|
||
unknown_rows = [row for row in output_rows if not row["answerable"]]
|
||
error_counts = Counter(row["error_class"] for row in output_rows)
|
||
failed_answerable = [row for row in answerable_rows if not row["correct"]]
|
||
return {
|
||
"cases": len(output_rows),
|
||
"answerable_cases": len(answerable_rows),
|
||
"answerable_correct": sum(int(row["correct"]) for row in answerable_rows),
|
||
"answerable_accuracy": sum(int(row["correct"]) for row in answerable_rows) / max(1, len(answerable_rows)),
|
||
"unknown_cases": len(unknown_rows),
|
||
"unknown_correct": sum(int(row["correct"]) for row in unknown_rows),
|
||
"unknown_refusal_accuracy": sum(int(row["correct"]) for row in unknown_rows) / max(1, len(unknown_rows)),
|
||
"error_class_counts": dict(sorted(error_counts.items())),
|
||
"failed_answerable_cases": len(failed_answerable),
|
||
"failed_answerable_with_correct_evidence": sum(
|
||
int(row["retrieval_target_found"]) for row in failed_answerable
|
||
),
|
||
"failed_answerable_with_reference_retrieval_hit": sum(
|
||
int(row["reference_retrieval_hit"]) for row in failed_answerable
|
||
),
|
||
"failed_answerable_with_retrieval_miss": sum(
|
||
int(not row["retrieval_target_found"]) for row in failed_answerable
|
||
),
|
||
"rows": output_rows,
|
||
}
|
||
|
||
|
||
def _retrieval_summary(rows: list[dict[str, Any]]) -> dict[str, Any]:
|
||
answerable = [row for row in rows if row["answerable"]]
|
||
unknown = [row for row in rows if not row["answerable"]]
|
||
return {
|
||
"queries": len(rows),
|
||
"answerable_queries": len(answerable),
|
||
"answerable_hits": sum(int(row["hit"]) for row in answerable),
|
||
"answerable_recall": sum(int(row["hit"]) for row in answerable) / max(1, len(answerable)),
|
||
"unknown_queries": len(unknown),
|
||
"unknown_correct_abstentions": sum(int(row["hit"]) for row in unknown),
|
||
"unknown_abstention_accuracy": sum(int(row["hit"]) for row in unknown) / max(1, len(unknown)),
|
||
"read_latency_ms": _quantiles(row["elapsed_ms"] for row in rows),
|
||
"mean_coarse_candidates": mean(row["coarse_candidates"] for row in rows) if rows else 0.0,
|
||
"max_coarse_candidates": max((row["coarse_candidates"] for row in rows), default=0),
|
||
}
|
||
|
||
|
||
def _vram_snapshot(device: torch.device) -> dict[str, float | None]:
|
||
if device.type != "cuda":
|
||
return {"allocated_gb": None, "reserved_gb": None, "peak_allocated_gb": None, "peak_reserved_gb": None}
|
||
return {
|
||
"allocated_gb": torch.cuda.memory_allocated(device) / 1024**3,
|
||
"reserved_gb": torch.cuda.memory_reserved(device) / 1024**3,
|
||
"peak_allocated_gb": torch.cuda.max_memory_allocated(device) / 1024**3,
|
||
"peak_reserved_gb": torch.cuda.max_memory_reserved(device) / 1024**3,
|
||
}
|
||
|
||
|
||
@torch.inference_mode()
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser(description=__doc__)
|
||
parser.add_argument("--base-model", default=r"H:\Memory")
|
||
parser.add_argument("--memory-model", default=str(DEFAULT_MEMORY_MODEL))
|
||
parser.add_argument("--data-root", default=str(PROJECT_ROOT / "data"))
|
||
parser.add_argument("--output", default=str(PROJECT_ROOT / "real_scale_memory_4b.json"))
|
||
parser.add_argument("--general-cases", type=int, default=640)
|
||
parser.add_argument("--project-records", type=int, default=8192)
|
||
parser.add_argument("--project-targets", type=int, default=256)
|
||
parser.add_argument("--project-chunk-tokens", type=int, default=512)
|
||
parser.add_argument("--generation-general", type=int, default=128)
|
||
parser.add_argument("--generation-project", type=int, default=64)
|
||
parser.add_argument("--max-new-tokens", type=int, default=32)
|
||
parser.add_argument("--encode-batch-size", type=int, default=16)
|
||
parser.add_argument("--gpu-memory-gb", type=float, default=10.0)
|
||
parser.add_argument("--no-4bit", action="store_true")
|
||
args = parser.parse_args()
|
||
_set_cuda_process_cap(args.gpu_memory_gb)
|
||
|
||
tokenizer = load_tokenizer(_path(args.base_model))
|
||
use_4bit = not args.no_4bit
|
||
print("loading Natural Memory 4B")
|
||
model = load_qwen_dynamic(
|
||
_path(args.memory_model),
|
||
load_in_4bit=use_4bit,
|
||
max_memory=_max_memory(args.gpu_memory_gb),
|
||
)
|
||
model.configure_memory_grounding_guard(tokenizer)
|
||
model.eval()
|
||
model.memory_config.memory_top_k_records = 2
|
||
model.memory_os_v2.bank.top_k_records = 2
|
||
device = model._find_layer_device()
|
||
model.reset_memory(batch_size=1, device=device)
|
||
_sync(device)
|
||
if device.type == "cuda":
|
||
torch.cuda.reset_peak_memory_stats(device)
|
||
baseline_vram = _vram_snapshot(device)
|
||
|
||
data_root = _path(args.data_root)
|
||
general_cases = _native_cases(data_root, max(1, int(args.general_cases)))
|
||
general_raw_records, general_queries = _build_general_records(general_cases)
|
||
general_records = _prepare_records(
|
||
model,
|
||
tokenizer,
|
||
general_raw_records,
|
||
device,
|
||
batch_size=max(1, int(args.encode_batch_size)),
|
||
)
|
||
_sync(device)
|
||
general_write_start = time.perf_counter()
|
||
model.memory_os_v2.write_batch(general_records)
|
||
_sync(device)
|
||
general_write_ms = (time.perf_counter() - general_write_start) * 1000.0
|
||
general_query_keys = _encode_texts(
|
||
model,
|
||
tokenizer,
|
||
[str(row["query"]) for row in general_queries],
|
||
device,
|
||
batch_size=max(1, int(args.encode_batch_size)),
|
||
)
|
||
general_retrieval_rows = [
|
||
_direct_read(model, tokenizer, general_query_keys[index], row, device)
|
||
for index, row in enumerate(general_queries)
|
||
]
|
||
general_retrieval = _retrieval_summary(general_retrieval_rows)
|
||
general_generation_rows = [
|
||
row
|
||
for row in general_queries
|
||
if row["split"] == "eval"
|
||
][: max(1, int(args.generation_general))]
|
||
general_generation = _generation_eval(
|
||
model,
|
||
tokenizer,
|
||
general_generation_rows,
|
||
device,
|
||
max_new_tokens=max(1, int(args.max_new_tokens)),
|
||
reference_retrieval_rows={str(row["id"]): row for row in general_retrieval_rows},
|
||
)
|
||
general_vram = _vram_snapshot(device)
|
||
|
||
_clear_memory_bank(model, device)
|
||
project_files = _source_files()
|
||
project_raw_records, project_queries, project_meta = _build_project_records(
|
||
project_files,
|
||
record_count=max(128, int(args.project_records)),
|
||
target_count=max(1, int(args.project_targets)),
|
||
chunk_tokens=max(1, int(args.project_chunk_tokens)),
|
||
)
|
||
project_records = _prepare_records(
|
||
model,
|
||
tokenizer,
|
||
project_raw_records,
|
||
device,
|
||
batch_size=max(1, int(args.encode_batch_size)),
|
||
)
|
||
_sync(device)
|
||
project_write_start = time.perf_counter()
|
||
model.memory_os_v2.write_batch(project_records)
|
||
_sync(device)
|
||
project_write_ms = (time.perf_counter() - project_write_start) * 1000.0
|
||
project_query_keys = _encode_texts(
|
||
model,
|
||
tokenizer,
|
||
[str(row["query"]) for row in project_queries],
|
||
device,
|
||
batch_size=max(1, int(args.encode_batch_size)),
|
||
)
|
||
project_retrieval_rows = []
|
||
for index, row in enumerate(project_queries):
|
||
direct_row = dict(row)
|
||
direct_row["answerable"] = True
|
||
project_retrieval_rows.append(
|
||
_direct_read(model, tokenizer, project_query_keys[index], direct_row, device)
|
||
)
|
||
project_retrieval = _retrieval_summary(project_retrieval_rows)
|
||
project_generation_rows = [
|
||
{"id": f"project-{index}", "query": row["query"], "expected": row["expected"], "answerable": True}
|
||
for index, row in enumerate(project_queries[: max(1, int(args.generation_project))])
|
||
]
|
||
project_generation = _generation_eval(
|
||
model,
|
||
tokenizer,
|
||
project_generation_rows,
|
||
device,
|
||
max_new_tokens=max(1, int(args.max_new_tokens)),
|
||
reference_retrieval_rows={str(row["id"]): row for row in project_retrieval_rows},
|
||
)
|
||
project_vram = _vram_snapshot(device)
|
||
|
||
report = {
|
||
"benchmark": "real_scale_memory_4b",
|
||
"model": str(_path(args.memory_model)),
|
||
"base_model": str(_path(args.base_model)),
|
||
"quantization": "4bit_nf4" if use_4bit else "none",
|
||
"gpu_memory_cap_gb": float(args.gpu_memory_gb),
|
||
"device": str(device),
|
||
"priority_metrics": ["peak_vram", "correctness"],
|
||
"baseline_vram_after_load": baseline_vram,
|
||
"general_chat_memory": {
|
||
"source": "local_native_memory_train_plus_eval",
|
||
"cases": len(general_cases),
|
||
"records_written": len(general_records),
|
||
"bank_records": len(model.memory_os_v2.bank.records),
|
||
"write_ms": general_write_ms,
|
||
"retrieval": general_retrieval,
|
||
"generation": general_generation,
|
||
"vram_after_general": general_vram,
|
||
},
|
||
"project_repository_memory": {
|
||
**project_meta,
|
||
"records_written": len(project_records),
|
||
"bank_records": len(model.memory_os_v2.bank.records),
|
||
"page_count": len(model.memory_os_v2.bank.pages),
|
||
"write_ms": project_write_ms,
|
||
"retrieval": project_retrieval,
|
||
"generation": project_generation,
|
||
"vram_after_project": project_vram,
|
||
},
|
||
"limitations": [
|
||
"本地 native_memory 语料是工程内置的自然语言基准，不等同于真实用户导出数据。",
|
||
"项目库目标来自当前仓库的真实源码和文档；大规模背景页用于压力测试，主要考察有界路由和显存。",
|
||
"生成质量与检索召回分别报告，不能用检索正确率替代端到端聊天正确率。",
|
||
"没有启用 SQLite 或磁盘分页；记忆主体保持在进程 RAM，只有命中的记录进入有界显存缓存。",
|
||
],
|
||
}
|
||
output = _path(args.output)
|
||
output.parent.mkdir(parents=True, exist_ok=True)
|
||
output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
print(json.dumps({
|
||
"general_retrieval": general_retrieval,
|
||
"general_generation": {key: value for key, value in general_generation.items() if key != "rows"},
|
||
"project_retrieval": project_retrieval,
|
||
"project_generation": {key: value for key, value in project_generation.items() if key != "rows"},
|
||
"vram": {"baseline": baseline_vram, "general": general_vram, "project": project_vram},
|
||
}, ensure_ascii=False, indent=2))
|
||
print(f"saved={output}")
|
||
del model
|
||
gc.collect()
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|