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，读写关闭时与原生模型逐位相同
This commit is contained in:
WpyQwq
@@ -0,0 +1,231 @@
|
||||
"""Evaluate the production-style local dense+rereank baseline on dirty data.
|
||||
|
||||
This is the companion baseline for ``benchmark_dirty_real_corpus_4b``. It
|
||||
uses the same real repository corpus, the same user-dialogue slice, the same
|
||||
queries, tokenizer, greedy decoding and GPU cap. Retrieval is CPU-only:
|
||||
local BERT embeddings provide recall and a bounded transparent reranker uses
|
||||
lexical/entity/attribute signals. It is stronger than lexical Chunk RAG,
|
||||
but is explicitly not presented as a trained public cross-encoder.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import time
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
from statistics import mean
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from .benchmark_compare_baselines_4b import (
|
||||
_base_generate,
|
||||
_rag_prompt,
|
||||
_vram_snapshot,
|
||||
)
|
||||
from .benchmark_dirty_real_corpus_4b import (
|
||||
_build_all_records,
|
||||
_dirty_source_files,
|
||||
_is_refusal_strict,
|
||||
_path,
|
||||
_quality_pass,
|
||||
_retrieval_hit,
|
||||
_summarize,
|
||||
)
|
||||
from .benchmark_real_scale_memory_4b import _max_memory, _set_cuda_process_cap
|
||||
from .strong_rag_baseline import LocalEmbeddingReranker, _resolve_local_encoder
|
||||
from .qwen_integration import load_qwen_base, load_tokenizer
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
def _run_dirty(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
rows: list[dict[str, Any]],
|
||||
index: LocalEmbeddingReranker,
|
||||
device: torch.device,
|
||||
*,
|
||||
top_k: int,
|
||||
max_new_tokens: int,
|
||||
) -> dict[str, Any]:
|
||||
output_rows: list[dict[str, Any]] = []
|
||||
for row_index, row in enumerate(rows, 1):
|
||||
if row_index == 1 or row_index % 32 == 0 or row_index == len(rows):
|
||||
print(f"Strong RAG generation: {row_index}/{len(rows)}", flush=True)
|
||||
query = str(row["query"])
|
||||
retrieved, retrieve_ms = index.retrieve(query, top_k=top_k)
|
||||
prompt = _rag_prompt(query, retrieved)
|
||||
generated = _base_generate(
|
||||
model,
|
||||
tokenizer,
|
||||
prompt,
|
||||
device,
|
||||
max_new_tokens=max_new_tokens,
|
||||
)
|
||||
selected = [
|
||||
{
|
||||
"value": item.get("value", ""),
|
||||
"text_preview": item.get("text", ""),
|
||||
"entity": item.get("entity", ""),
|
||||
"attribute": item.get("attribute", ""),
|
||||
}
|
||||
for item in retrieved
|
||||
]
|
||||
response = str(generated.get("response", ""))
|
||||
output_rows.append(
|
||||
{
|
||||
**row,
|
||||
"response": response,
|
||||
"correct": _quality_pass(response, row),
|
||||
"retrieval_target_found": _retrieval_hit(selected, row),
|
||||
"retrieved_values": [item.get("value", "") for item in retrieved],
|
||||
"retrieved_ids": [],
|
||||
"selected_records": selected,
|
||||
"prefix_used": bool(retrieved),
|
||||
"prefix_tokens": int(generated.get("prompt_tokens", 0)),
|
||||
"reader_ms": float(retrieve_ms),
|
||||
"retriever_ms": float(retrieve_ms),
|
||||
"prompt_tokens": int(generated.get("prompt_tokens", 0)),
|
||||
"generated_tokens": int(generated.get("generated_tokens", 0)),
|
||||
"total_latency_s": float(generated.get("total_latency_s", 0.0)),
|
||||
"decode_tok_s": float(generated.get("decode_tok_s", 0.0)),
|
||||
"allocated_gb": generated.get("allocated_gb"),
|
||||
"reserved_gb": generated.get("reserved_gb"),
|
||||
}
|
||||
)
|
||||
return _summarize(output_rows, include_reader=False)
|
||||
|
||||
|
||||
def _release(model: Any) -> None:
|
||||
del model
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def _compact(summary: dict[str, Any]) -> dict[str, Any]:
|
||||
return {key: value for key, value in summary.items() if key != "rows"}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--base-model", default=r"H:\Memory")
|
||||
parser.add_argument("--output", default=str(PROJECT_ROOT / "dirty_strong_rag_compare_4b.json"))
|
||||
parser.add_argument("--project-records", type=int, default=8192)
|
||||
parser.add_argument("--project-targets", type=int, default=96)
|
||||
parser.add_argument("--top-k", type=int, default=4)
|
||||
parser.add_argument("--candidate-k", type=int, default=64)
|
||||
parser.add_argument("--embedding-batch-size", type=int, default=16)
|
||||
parser.add_argument("--embedding-max-length", type=int, default=384)
|
||||
parser.add_argument("--embedding-model", default=None)
|
||||
parser.add_argument("--max-new-tokens", type=int, default=64)
|
||||
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))
|
||||
files = _dirty_source_files()
|
||||
user_records, user_rows, code_records, code_rows, corpus_meta = _build_all_records(
|
||||
files,
|
||||
project_records=max(128, int(args.project_records)),
|
||||
project_targets=max(1, int(args.project_targets)),
|
||||
)
|
||||
all_records = user_records + code_records
|
||||
all_rows = user_rows + code_rows
|
||||
print(
|
||||
f"dirty corpus: files={len(files)} real_chunks={corpus_meta['stored_real_chunk_records']} "
|
||||
f"semantic_targets={corpus_meta['semantic_target_records']} user_records={len(user_records)} "
|
||||
f"queries={len(all_rows)}"
|
||||
)
|
||||
|
||||
encoder_path = _resolve_local_encoder(args.embedding_model)
|
||||
print(f"loading local embedding model on CPU: {encoder_path}")
|
||||
embedder = LocalEmbeddingReranker(
|
||||
encoder_path,
|
||||
max_length=args.embedding_max_length,
|
||||
batch_size=args.embedding_batch_size,
|
||||
candidate_k=args.candidate_k,
|
||||
)
|
||||
started = time.perf_counter()
|
||||
embedder.add(all_records)
|
||||
index_seconds = time.perf_counter() - started
|
||||
print(f"dense_index_seconds={index_seconds:.3f}")
|
||||
|
||||
print("loading original Qwen3.5 4B for strong-RAG generation")
|
||||
base = load_qwen_base(
|
||||
_path(args.base_model),
|
||||
load_in_4bit=not args.no_4bit,
|
||||
max_memory=_max_memory(args.gpu_memory_gb),
|
||||
)
|
||||
base.eval()
|
||||
device = base.get_input_embeddings().weight.device
|
||||
load_vram = _vram_snapshot(device)
|
||||
result = _run_dirty(
|
||||
base,
|
||||
tokenizer,
|
||||
all_rows,
|
||||
embedder,
|
||||
device,
|
||||
top_k=max(1, args.top_k),
|
||||
max_new_tokens=max(1, args.max_new_tokens),
|
||||
)
|
||||
_release(base)
|
||||
|
||||
report = {
|
||||
"benchmark": "dirty_strong_rag_compare_4b",
|
||||
"created_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
|
||||
"base_model": str(_path(args.base_model)),
|
||||
"embedding_model": str(encoder_path),
|
||||
"quantization": "4bit_nf4" if not args.no_4bit else "none",
|
||||
"gpu_memory_cap_gb": float(args.gpu_memory_gb),
|
||||
"corpus": corpus_meta,
|
||||
"query_counts": {
|
||||
"all": len(all_rows),
|
||||
"user_dialogue": len(user_rows),
|
||||
"real_repository": len(code_rows),
|
||||
"answerable": sum(int(row["answerable"]) for row in all_rows),
|
||||
"unknown": sum(int(not row["answerable"]) for row in all_rows),
|
||||
"by_category": dict(Counter(str(row["category"]) for row in all_rows)),
|
||||
},
|
||||
"protocol": {
|
||||
"same_tokenizer": True,
|
||||
"same_sampling": "greedy",
|
||||
"top_k": int(args.top_k),
|
||||
"candidate_k": int(args.candidate_k),
|
||||
"embedding_device": "cpu",
|
||||
"embedding_max_length": int(args.embedding_max_length),
|
||||
"embedding_batch_size": int(args.embedding_batch_size),
|
||||
"max_new_tokens": int(args.max_new_tokens),
|
||||
"baseline_note": "local BERT dense retrieval plus fixed transparent reranker; not a trained public cross-encoder",
|
||||
},
|
||||
"index_build": {
|
||||
"records": len(all_records),
|
||||
"indexed_records": len(embedder.records),
|
||||
"seconds": index_seconds,
|
||||
},
|
||||
"system": {
|
||||
"load_vram": load_vram,
|
||||
"all": result,
|
||||
},
|
||||
"limitations": [
|
||||
"The embedding encoder is local bert-base-chinese and CPU-only.",
|
||||
"The reranker is a fixed transparent feature reranker, not a trained public cross-encoder.",
|
||||
"The same strict answer-anchor evaluator is used for Natural Memory comparison.",
|
||||
"Embedding/index build time is reported separately from per-query retrieval and generation.",
|
||||
],
|
||||
}
|
||||
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({"strong_rag": _compact(result), "index_build": report["index_build"]}, ensure_ascii=False, indent=2))
|
||||
print(f"saved={output}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user