natural-memory-nm21/benchmark_dirty_real_corpus_4b.py

"""Run a dirty, real-world 4B memory transfer benchmark.

This benchmark is deliberately different from the synthetic fact/value sets:

* the large corpus is the actual Natural Memory repository (source, docs and
  generated engineering artifacts), not fabricated memory IDs;
* a separate slice contains facts explicitly stated by the user in the
  current project conversation;
* query wording is varied by hand into colloquial, paraphrased, temporal and
  conflict-aware forms;
* code questions ask for operational explanations and repository navigation,
  not only arbitrary string recall;
* unknown and stale-value cases are kept in the score.

The script compares the original local Qwen3.5-4B with a no-memory run, a
transparent lexical Chunk RAG run, and Natural Memory v2.  The source corpus
is large, but only bounded target records receive Qwen semantic keys; all
other stored records remain real source chunks and are intentionally kept as
unresolved background noise.  The report records that distinction explicitly
so the result is not mistaken for a fully semantically indexed production
repository.
"""

from __future__ import annotations

import argparse
import gc
import json
import re
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,
    _build_rag_index,
    _rag_prompt,
    _rag_retrieve,
    _record_entry,
    _vram_snapshot,
)
from .benchmark_real_scale_memory_4b import (
    _chat_generate,
    _contains_answer,
    _is_refusal,
    _max_memory,
    _prepare_records,
    _project_targets,
    _set_cuda_process_cap,
    _source_files,
    _sync,
)
from .qwen_integration import load_qwen_base, load_qwen_dynamic, load_tokenizer


PROJECT_ROOT = Path(__file__).resolve().parent
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 _normal(text: Any) -> str:
    return re.sub(r"\s+", "", str(text)).lower()


def _contains_all(text: str, anchors: list[str]) -> bool:
    normalized = _normal(text)
    return bool(anchors) and all(_normal(anchor) in normalized for anchor in anchors)


def _is_refusal_strict(text: str) -> bool:
    normalized = _normal(text)
    if any(marker in normalized for marker in REFUSAL_MARKERS + ("没有列出", "未列出", "没有具体")):
        return True
    if re.search(
        r"(?:没有|不存在|未找到|尚未|并未|不包含|不在|无法).{0,100}"
        r"(?:记录|资料|信息|数据|找到|知道|存在|名为|结果|账单|配置|地址|文件|脚本|客户|事实|实现|证明|记忆|提到|回答|提供|恢复)",
        normalized,
    ):
        return True
    return _is_refusal(text)


def _quality_pass(response: str, row: dict[str, Any]) -> bool:
    if not bool(row.get("answerable", True)):
        return _is_refusal_strict(response)
    anchors = [str(item) for item in row.get("expected_anchors", []) if str(item)]
    if anchors:
        return _contains_all(response, anchors)
    return _contains_answer(response, str(row.get("expected", "")))


def _case(
    case_id: str,
    *,
    domain: str,
    category: str,
    query: str,
    anchors: list[str] | None = None,
    expected: str = "",
    answerable: bool = True,
    target_values: list[str] | None = None,
    note: str = "",
) -> dict[str, Any]:
    return {
        "id": case_id,
        "domain": domain,
        "category": category,
        "query": query,
        "expected": expected or (anchors[0] if anchors else ""),
        "expected_anchors": list(anchors or ([] if not expected else [expected])),
        "answerable": bool(answerable),
        "target_values": list(target_values or []),
        "note": note,
    }


def _user_dialogue_records() -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    """Return a privacy-safe slice grounded in the user's actual messages.

    These are not logs exported from a hidden service.  They are short facts
    explicitly stated by the user in this project conversation.  The old and
    new working directories intentionally share one conflict key so the test
    exercises version replacement rather than a lookup table.
    """

    records = [
        {
            "text": '用户早期明确说过："W:\\Flash\\model" 在这里工作，使用 Conda LLM，开始工程。',
            "entity": "user",
            "attribute": "工作目录",
            "value": r"W:\Flash\model",
            "memory_type": "user_dialogue_fact",
            "importance": 0.85,
            "confidence": 0.96,
            "source": "user_provided_dialogue_earlier",
            "trusted": True,
            "force": True,
        },
        {
            "text": '用户随后明确要求：把项目移到 H:\\Memory；当前工程以 H:\\Memory 为工作目录。',
            "entity": "user",
            "attribute": "工作目录",
            "value": r"H:\Memory",
            "memory_type": "user_dialogue_fact",
            "importance": 0.95,
            "confidence": 0.99,
            "source": "user_provided_dialogue_later",
            "trusted": True,
            "force": True,
        },
        {
            "text": "用户给当前模型定名：Natural Memory v1。",
            "entity": "user",
            "attribute": "模型名称",
            "value": "Natural Memory v1",
            "memory_type": "user_dialogue_fact",
            "importance": 0.9,
            "confidence": 0.98,
            "source": "user_provided_dialogue",
            "trusted": True,
            "force": True,
        },
        {
            "text": "用户要求和原版 qwen3.5 4b 做跑分对比，并纠正为 4B。",
            "entity": "user",
            "attribute": "对照模型",
            "value": "Qwen3.5 4B",
            "memory_type": "user_dialogue_fact",
            "importance": 0.95,
            "confidence": 0.99,
            "source": "user_provided_dialogue",
            "trusted": True,
            "force": True,
        },
        {
            "text": "用户要求：记忆尽量加载到 DRAM，加载不下才进入 RAM。",
            "entity": "user",
            "attribute": "记忆存储优先级",
            "value": "DRAM",
            "memory_type": "user_dialogue_preference",
            "importance": 0.9,
            "confidence": 0.98,
            "source": "user_provided_dialogue",
            "trusted": True,
            "force": True,
        },
        {
            "text": "用户要求：不使用 SQLite 或磁盘分页，完全由第三个权重文件储存所有信息。",
            "entity": "user",
            "attribute": "记忆持久化方案",
            "value": "第三个权重文件，不使用 SQLite 或磁盘分页",
            "memory_type": "user_dialogue_preference",
            "importance": 0.9,
            "confidence": 0.98,
            "source": "user_provided_dialogue",
            "trusted": True,
            "force": True,
        },
        {
            "text": "用户要求模型原生读取记忆：把记忆读取器直接设计成模型架构的一部分，不依赖外部代码指示模型读取。",
            "entity": "user",
            "attribute": "记忆读取方式",
            "value": "模型架构内置记忆读取器",
            "memory_type": "user_dialogue_requirement",
            "importance": 0.95,
            "confidence": 0.99,
            "source": "user_provided_dialogue",
            "trusted": True,
            "force": True,
        },
        {
            "text": "用户要求：不要顶满显存。",
            "entity": "user",
            "attribute": "显存安全要求",
            "value": "不要顶满显存",
            "memory_type": "user_dialogue_requirement",
            "importance": 0.95,
            "confidence": 0.99,
            "source": "user_provided_dialogue",
            "trusted": True,
            "force": True,
        },
        {
            "text": "用户要求：重要信息都要存进去，但不能让模型胡编乱造记忆，还要学会什么时候写入和清理总结。",
            "entity": "user",
            "attribute": "记忆质量要求",
            "value": "重要信息写入，禁止胡编乱造，并学习写入、清理和总结时机",
            "memory_type": "user_dialogue_requirement",
            "importance": 0.95,
            "confidence": 0.98,
            "source": "user_provided_dialogue",
            "trusted": True,
            "force": True,
        },
        {
            "text": "用户明确表示自己更偏向训练，希望从现成预训练模型进行 SFT 或进一步训练。",
            "entity": "user",
            "attribute": "研发取向",
            "value": "训练",
            "memory_type": "user_dialogue_preference",
            "importance": 0.8,
            "confidence": 0.95,
            "source": "user_provided_dialogue",
            "trusted": True,
            "force": True,
        },
        {
            "text": "用户的核心目标是模型原生的记忆能力：不依靠固定外部代码指示模型读取，也不依靠外界文件保存读取逻辑。",
            "entity": "user",
            "attribute": "核心架构目标",
            "value": "模型原生记忆能力",
            "memory_type": "user_dialogue_requirement",
            "importance": 0.95,
            "confidence": 0.99,
            "source": "user_provided_dialogue",
            "trusted": True,
            "force": True,
        },
        {
            "text": "用户要求 Memory Slot 不要过度浓缩，本质是用更高压缩率替代长 KV，但仍要保留碎片对话。",
            "entity": "user",
            "attribute": "记忆压缩策略",
            "value": "高压缩率但保留碎片对话",
            "memory_type": "user_dialogue_requirement",
            "importance": 0.9,
            "confidence": 0.98,
            "source": "user_provided_dialogue",
            "trusted": True,
            "force": True,
        },
        {
            "text": "用户定义的运行分工是：KV 作为当前运行内存，长期历史进入可规模化的 Memory Slot；当前 token 不能对全部 slot 做注意力。",
            "entity": "user",
            "attribute": "KV与Slot分工",
            "value": "KV 负责当前运行内存，Memory Slot 负责长期历史",
            "memory_type": "user_dialogue_requirement",
            "importance": 0.95,
            "confidence": 0.98,
            "source": "user_provided_dialogue",
            "trusted": True,
            "force": True,
        },
        {
            "text": "用户要求尽量由模型自发完成记忆和测试流程，而不是每次都依赖外部固定传参。",
            "entity": "user",
            "attribute": "自动化边界",
            "value": "尽量由模型自发完成",
            "memory_type": "user_dialogue_requirement",
            "importance": 0.9,
            "confidence": 0.97,
            "source": "user_provided_dialogue",
            "trusted": True,
            "force": True,
        },
        {
            "text": "用户明确要求不要只读取一两条记忆，要进行大规模读取和通用样本聊天测试，重点看正确率和显存。",
            "entity": "user",
            "attribute": "评测重点",
            "value": "大规模读取、正确率和显存",
            "memory_type": "user_dialogue_requirement",
            "importance": 0.95,
            "confidence": 0.99,
            "source": "user_provided_dialogue",
            "trusted": True,
            "force": True,
        },
    ]

    rows: list[dict[str, Any]] = []
    variants = [
        ("工作目录", [r"H:\Memory"], "现在这个工程目录到底以哪里为准？", "temporal_conflict"),
        ("工作目录", [r"H:\Memory"], "之前那个 W 盘目录已经不是当前工作区了吧？现在在哪儿干活？", "temporal_conflict"),
        ("工作目录", [r"H:\Memory"], "重启后不要看历史聊天，直接告诉我项目现在放在哪个目录。", "cross_session"),
        ("模型名称", ["Natural Memory v1"], "这个模型之前定的正式名字是什么？", "cross_session"),
        ("模型名称", ["Natural Memory v1"], "我不想翻聊天记录，当前这个记忆模型叫什么？", "paraphrase"),
        ("对照模型", ["Qwen3.5 4B"], "这轮对照实验基线要用哪个 Qwen 规模？", "paraphrase"),
        ("对照模型", ["Qwen3.5 4B"], "我说的原版模型是 4B 的哪个版本？", "paraphrase"),
        ("记忆存储优先级", ["DRAM"], "记忆页应该优先放哪一层内存？", "paraphrase"),
        ("记忆存储优先级", ["DRAM"], "容量不够时才降级，平时优先使用什么？", "paraphrase"),
        ("记忆持久化方案", ["第三个权重文件", "不使用 SQLite", "不使用磁盘分页"], "你记得我对持久化文件和分页的限制吗？", "constraint"),
        ("记忆持久化方案", ["第三个权重文件", "不使用 SQLite"], "这版记忆能不能再偷偷依赖一个数据库？", "constraint"),
        ("记忆读取方式", ["模型架构内置", "记忆读取器"], "我想要的是模型自己读记忆，不是外部脚本拼提示词，对吗？", "constraint"),
        ("显存安全要求", ["不要顶满显存"], "跑大规模测试时最重要的显存原则是什么？", "constraint"),
        ("记忆质量要求", ["禁止胡编乱造", "写入"], "长期记忆为什么不能只追求召回，还要管写错和清理？", "reasoning"),
        ("工作目录", [r"H:\Memory"], "我把工程从 W 盘挪走之后，现在实际在哪个目录继续？", "temporal_conflict"),
        ("工作目录", [r"H:\Memory"], "不要按旧聊天猜，当前项目路径是什么？", "cross_session"),
        ("工作目录", [r"H:\Memory"], "W:\\Flash\\model 和 H:\\Memory 哪个是现在的工作区？", "temporal_conflict"),
        ("模型名称", ["Natural Memory v1"], "我之前给这套记忆模型起的名字，原样是什么？", "cross_session"),
        ("模型名称", ["Natural Memory v1"], "不看历史上下文，你还记得这个项目的模型名吗？", "cross_session"),
        ("对照模型", ["Qwen3.5 4B"], "这次不要弄错规模，原版对照是 4B 还是别的？", "constraint"),
        ("对照模型", ["Qwen3.5 4B"], "跑分时要拿 Natural Memory 和哪个原版 Qwen 比？", "paraphrase"),
        ("记忆存储优先级", ["DRAM"], "长期记忆平时先驻留在哪种内存里？", "paraphrase"),
        ("记忆存储优先级", ["DRAM"], "别一上来就把所有记忆塞显卡，优先级怎么排？", "constraint"),
        ("记忆持久化方案", ["第三个权重文件", "不使用磁盘分页"], "持久化记忆是不是要单独放进第三个权重切片，而不是分页文件？", "constraint"),
        ("记忆持久化方案", ["第三个权重文件", "不使用 SQLite"], "之前说过的存储限制，数据库方案被排除了吗？", "constraint"),
        ("记忆读取方式", ["模型原生记忆能力"], "我的终极要求是让模型自己拥有记忆能力，关键词是什么？", "constraint"),
        ("记忆读取方式", ["模型架构内置", "记忆读取器"], "读取器应该放在模型里面，还是让外部服务每次提醒它？", "constraint"),
        ("显存安全要求", ["不要顶满显存"], "测试再大也不能牺牲哪条显存安全底线？", "constraint"),
        ("记忆质量要求", ["禁止胡编乱造"], "为什么我反复强调写入时不能让模型自己编记忆？", "reasoning"),
        ("研发取向", ["训练"], "相比只做应用拼装，我更想把精力放在哪个方向？", "paraphrase"),
        ("研发取向", ["训练"], "我之前说过自己偏向训练还是只想调用模型？", "cross_session"),
        ("核心架构目标", ["模型原生记忆能力"], "我不想靠固定外部代码读记忆，真正想做的核心是什么？", "constraint"),
        ("记忆压缩策略", ["高压缩率", "碎片对话"], "Memory Slot 是不是应该把所有对话都压成一句话？", "constraint"),
        ("KV与Slot分工", ["KV", "Memory Slot"], "当前上下文和长期历史，在这个架构里分别由什么负责？", "reasoning"),
        ("KV与Slot分工", ["不能对全部 slot 做注意力"], "为什么不能让当前 token 和一百万个 slot 全部算注意力？", "reasoning"),
        ("自动化边界", ["模型自发完成"], "我希望以后少依赖外部固定参数，这个偏好怎么概括？", "paraphrase"),
        ("评测重点", ["大规模读取", "正确率", "显存"], "我要求测试不要只读一两条，真正优先看哪三个东西？", "constraint"),
        ("评测重点", ["大规模读取"], "这次评测的规模要求是什么，不是做几个 demo 就结束？", "constraint"),
    ]
    for index, (attribute, anchors, query, category) in enumerate(variants):
        rows.append(
            _case(
                f"user-{index:03d}",
                domain="user_dialogue",
                category=category,
                query=query,
                anchors=anchors,
                target_values=[anchors[0]],
                note="actual user-stated fact with a hand-written natural paraphrase",
            )
        )

    unknowns = [
        "我之前告诉过你我的血型是什么？",
        "我的家庭住址在记忆里是哪一条？",
        "我最喜欢的编程语言是什么？",
        "我的银行卡尾号是多少？",
        "我在这次对话里有没有说过生日？",
        "请直接告诉我没有记录的护照号码。",
    ]
    for index, query in enumerate(unknowns):
        rows.append(
            _case(
                f"user-unknown-{index:03d}",
                domain="user_dialogue",
                category="unknown_refusal",
                query=query,
                answerable=False,
                note="unknown or deliberately sensitive field absent from the supplied dialogue slice",
            )
        )
    return records, rows


def _source_chunks_with_paths(files: list[Path], *, chars_per_chunk: int = 1800) -> list[dict[str, Any]]:
    chunks: list[dict[str, Any]] = []
    for path in files:
        text = path.read_text(encoding="utf-8", errors="replace")
        relative = path.relative_to(PROJECT_ROOT).as_posix()
        for start in range(0, len(text), chars_per_chunk):
            piece = text[start : start + chars_per_chunk].strip()
            if not piece:
                continue
            chunks.append(
                {
                    "file": relative,
                    "start_char": start,
                    "text": f"真实项目文件：{relative}；字符区间：{start}-{start + len(piece)}\n{piece}",
                }
            )
    if not chunks:
        raise RuntimeError("repository source corpus is empty")
    return chunks


def _dirty_source_files() -> list[Path]:
    """Select real source material while excluding generated benchmark output.

    The repository is intentionally messy at the source/document level, but
    prior score JSONs and synthetic datasets are not allowed to become answer
    leaks.  Model configuration JSON is retained because it is part of the
    shipped project, while tokenizer dumps and checkpoints are excluded.
    """

    selected: list[Path] = []
    for path in _source_files():
        if "__pycache__" in path.parts or "checkpoints" in path.parts or "data" in path.parts:
            continue
        # The failure-driven regression harness is test infrastructure for
        # this benchmark, not part of the repository workload being measured.
        # Keeping it out prevents a newly added test file from changing the
        # evenly sampled symbol target set and contaminating before/after
        # comparisons with duplicate helper names.
        if path.name == "test_failure_driven_v2.py":
            continue
        if path.name.startswith("Natural_Memory_v2_对外技术总结"):
            continue
        if path.suffix.lower() in {".py", ".md"}:
            selected.append(path)
            continue
        if path.suffix.lower() == ".json" and path.name in {
            "config.json",
            "generation_config.json",
            "memory_merge.json",
        }:
            selected.append(path)
    return sorted(selected)


def _sample_evenly(items: list[Any], limit: int) -> list[Any]:
    if len(items) <= limit:
        return list(items)
    return [items[int(index * len(items) / limit)] for index in range(limit)]


def _excerpt(path: Path, needle: str, *, radius: int = 1250) -> str:
    text = path.read_text(encoding="utf-8", errors="replace")
    position = text.lower().find(needle.lower())
    if position < 0:
        raise ValueError(f"grounding phrase not found in {path.name}: {needle}")
    start = max(0, position - radius // 2)
    end = min(len(text), position + radius)
    relative = path.relative_to(PROJECT_ROOT).as_posix()
    return f"真实项目文件：{relative}\n{text[start:end].strip()}"


def _build_code_corpus(
    files: list[Path],
    *,
    record_count: int,
    target_count: int,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]], dict[str, Any]]:
    chunks = _source_chunks_with_paths(files)
    sampled = _sample_evenly(chunks, max(128, record_count))
    records: list[dict[str, Any]] = []
    # Real repository chunks are all retained.  They are background noise for
    # this run; only explicit target facts below receive Qwen semantic keys.
    for index, chunk in enumerate(sampled):
        records.append(
            {
                "text": chunk["text"],
                "entity": chunk["file"],
                "attribute": f"source_chunk:{index}",
                "value": chunk["file"],
                "memory_type": "real_repository_chunk",
                "importance": 0.4,
                "confidence": 0.85,
                "source": "real_repository_snapshot",
                "trusted": True,
                "force": True,
                "key_kind": "random_filler",
            }
        )

    targets = _project_targets(files, max(1, target_count))
    for target in targets:
        records.append(
            {
                "text": target["text"],
                "entity": target["file"],
                "attribute": f"symbol:{target['name']}",
                "value": target["file"],
                "memory_type": "real_repository_symbol",
                "importance": 0.9,
                "confidence": 0.96,
                "source": "real_repository_symbol_index",
                "trusted": True,
                "force": True,
            }
        )

    operational_specs = [
        (
            "README_NATURAL_MEMORY_V2.md",
            "不使用 SQLite",
            "重启后模型不带历史聊天，它到底从哪里恢复记忆？",
            ["memory safetensors", "memory shard"],
            "operational_architecture",
        ),
        (
            "README_NATURAL_MEMORY_V2.md",
            "当前 token 不会对 1M slot 做全量注意力",
            "百万 slot 的时候，当前 token 会不会和所有 slot 做全量注意力？",
            ["Top-K", "全量注意力"],
            "architecture_tradeoff",
        ),
        (
            "README_NATURAL_MEMORY_V2.md",
            "低置信度写入",
            "记忆写入为什么要有 quarantine，而不是所有内容直接变成 active？",
            ["quarantine", "置信度"],
            "write_safety",
        ),
        (
            "memory_os_v2.py",
            "canonical copy remains in process RAM",
            "命中热点记录后，权威记忆副本和 GPU cache 分别放在哪里？",
            ["process RAM", "GPU cache"],
            "memory_tiers",
        ),
        (
            "benchmark_compare_baselines_4b.py",
            "same-protocol 4B comparison",
            "当前 4B 对照脚本的实验对象和协议定位是什么？",
            ["base Qwen", "Chunk RAG", "Natural Memory"],
            "benchmark_protocol",
        ),
        (
            "strong_rag_baseline.py",
            "not a public cross-encoder model",
            "当前 strong RAG 基线是不是一个训练好的公开 cross-encoder？",
            ["不是", "cross-encoder"],
            "baseline_limitation",
        ),
        (
            "qwen_integration.py",
            "DEFAULT_MEMORY_RESET_TOKEN",
            "模型代码里用于清空记忆的默认 reset token 是什么？",
            ["<|fim_prefix|>"],
            "control_path",
        ),
        (
            "qwen_integration.py",
            "_encode_model_key",
            "Natural Memory 的语义地址是怎么从 Qwen 主干得到的？",
            ["_encode_model_key", "hidden state"],
            "router_path",
        ),
    ]
    code_rows: list[dict[str, Any]] = []
    for index, (file_name, needle, query, anchors, category) in enumerate(operational_specs):
        path = PROJECT_ROOT / file_name
        text = _excerpt(path, needle)
        records.append(
            {
                "text": text,
                "entity": file_name,
                "attribute": f"operational:{index}",
                # Keep the source file in ``entity`` for audit and retrieval,
                # while exposing the normalized answer anchors as the value
                # the model should reuse in its response.
                "value": "；".join(anchors),
                "memory_type": "real_repository_operational_fact",
                "importance": 0.95,
                "confidence": 0.98,
                "source": "real_repository_snapshot",
                "trusted": True,
                "force": True,
            }
        )
        code_rows.extend(
            [
                _case(
                    f"code-op-{index:03d}-a",
                    domain="real_repository",
                    category=category,
                    query=query,
                    anchors=anchors,
                    target_values=[file_name],
                    note=f"grounded excerpt from {file_name}",
                ),
                _case(
                    f"code-op-{index:03d}-b",
                    domain="real_repository",
                    category=category,
                    query="我不想只看一个函数名，结合项目资料解释一下：" + query,
                    anchors=anchors,
                    target_values=[file_name],
                    note=f"natural explanatory rewrite grounded in {file_name}",
                ),
            ]
        )

    for index, target in enumerate(targets):
        file_name = str(target["file"])
        symbol = str(target["name"])
        kind = str(target["kind"])
        code_rows.extend(
            [
                _case(
                    f"code-symbol-{index:03d}-a",
                    domain="real_repository",
                    category="symbol_navigation",
                    query=f"我现在不想翻完整个屎山代码库，{symbol} 这个 {kind} 定义落在哪个文件？",
                    anchors=[file_name, symbol],
                    target_values=[file_name],
                    note="real symbol extracted from repository source",
                ),
                _case(
                    f"code-symbol-{index:03d}-b",
                    domain="real_repository",
                    category="symbol_navigation",
                    query=f"帮我定位项目里的 {symbol}，告诉我文件位置，并说明它是不是一个 {kind}。",
                    anchors=[file_name, symbol],
                    target_values=[file_name],
                    note="colloquial paraphrase of a real source lookup",
                ),
            ]
        )

    unknowns = [
        "这个仓库里有没有一个叫 totally_fake_memory_router_999 的函数？",
        "README 里是否记录了不存在的 10 亿页 GPU 常驻配置？",
        "请从项目记忆告诉我作者没有写过的商业客户名称。",
        "这个代码库的 API 密钥是什么？",
        "当前项目是否已经实现了 32B 模型的实测结果？",
        "请找出一个不存在的文件 H:\\Memory\\secret\\private_notes.py。",
        "仓库里有没有名为 MemoryTeleportationController 的类？",
        "这个项目是否已经证明逐 token 无损等价 1M KV？",
        "请告诉我没有出现在仓库里的部署区域。",
        "项目中是否存在一个叫 final_production_elon_demo.py 的脚本？",
        "README 有没有给出真实线上每百万 token 的美元账单？",
        "请从当前代码库恢复用户银行卡尾号。",
    ]
    for index, query in enumerate(unknowns):
        code_rows.append(
            _case(
                f"code-unknown-{index:03d}",
                domain="real_repository",
                category="unknown_refusal",
                query=query,
                answerable=False,
                note="target deliberately absent from the real repository snapshot",
            )
        )
    meta = {
        "source_file_count": len(files),
        "source_characters": sum(path.stat().st_size for path in files),
        "source_chunk_count": len(chunks),
        "stored_real_chunk_records": len(sampled),
        "semantic_target_records": len(targets) + len(operational_specs),
        "record_count": len(records),
        "target_count": len(targets),
        "operational_fact_count": len(operational_specs),
        "chunk_chars": 1800,
        "selection_note": "all stored chunks are real repository text; only symbol and operational target records receive Qwen semantic keys in this bounded run",
    }
    return records, code_rows, meta


def _build_all_records(
    files: list[Path],
    *,
    project_records: int,
    project_targets: int,
) -> tuple[
    list[dict[str, Any]],
    list[dict[str, Any]],
    list[dict[str, Any]],
    list[dict[str, Any]],
    dict[str, Any],
]:
    user_records, user_rows = _user_dialogue_records()
    code_records, code_rows, code_meta = _build_code_corpus(
        files,
        record_count=project_records,
        target_count=project_targets,
    )
    return user_records, user_rows, code_records, code_rows, code_meta


def _retrieval_hit(selected_records: list[dict[str, Any]], row: dict[str, Any]) -> bool:
    targets = [_normal(item) for item in row.get("target_values", []) if str(item)]
    if not targets:
        return False
    for record in selected_records:
        values = [_normal(record.get("value", "")), _normal(record.get("text_preview", ""))]
        if any(target in value for target in targets for value in values):
            return True
    return False


def _failure_stage(row: dict[str, Any]) -> str | None:
    """Attribute an answerable failure to routing, fusion, or control.

    The classification is deliberately observable rather than speculative:
    a routing miss means the target value/file was absent from the selected
    records; a fusion miss means the target was present but required answer
    anchors were not carried into the response; a control miss means the
    model had usable evidence but refused or overrode it.
    """

    if bool(row.get("correct", False)):
        return None
    if not bool(row.get("answerable", True)):
        return "unknown_refusal_control"
    if not bool(row.get("retrieval_target_found", False)):
        return "routing_miss"
    if _is_refusal_strict(str(row.get("response", ""))):
        return "generation_control"
    anchors = [str(item) for item in row.get("expected_anchors", []) if str(item)]
    normalized = _normal(row.get("response", ""))
    matched = sum(int(_normal(anchor) in normalized) for anchor in anchors)
    if anchors and matched < len(anchors):
        return "evidence_fusion"
    return "generation_control"


def _natural_rows(
    model: Any,
    tokenizer: Any,
    records: list[dict[str, Any]],
    rows: list[dict[str, Any]],
    device: torch.device,
    *,
    max_new_tokens: int,
    encode_batch_size: int,
) -> dict[str, Any]:
    model.eval()
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    model.memory_config.memory_top_k_records = 2
    model.memory_os_v2.bank.top_k_records = 2
    model.reset_memory(batch_size=1, device=device)
    prepared = _prepare_records(model, tokenizer, records, device, batch_size=max(1, encode_batch_size))
    model.memory_os_v2.write_batch(prepared)
    output_rows: list[dict[str, Any]] = []
    for index, row in enumerate(rows, 1):
        if index == 1 or index % 32 == 0 or index == len(rows):
            print(f"Natural Memory generation: {index}/{len(rows)}", flush=True)
        generated = _chat_generate(model, tokenizer, str(row["query"]), device, max_new_tokens)
        selected = list(generated.get("selected_records", []))
        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": list(generated.get("selected_values", [])),
                "retrieved_ids": list(generated.get("selected_record_ids", [])),
                "selected_records": selected,
                "prefix_used": bool(generated.get("prefix_used", False)),
                "guard_used": bool(generated.get("guard_used", False)),
                "prefix_tokens": int(generated.get("prefix_tokens", 0)),
                "reader_ms": float(generated.get("memory_read_seconds", 0.0)) * 1000.0,
                "prompt_tokens": int(generated.get("public_prompt_tokens", 0)) + int(generated.get("prefix_tokens", 0)),
                "generated_tokens": int(generated.get("generated_tokens", 0)),
                "total_latency_s": float(generated.get("generation_seconds", 0.0)),
                "decode_tok_s": int(generated.get("generated_tokens", 0)) / max(float(generated.get("generation_seconds", 0.0)), 1e-9),
                "allocated_gb": _vram_snapshot(device)["allocated_gb"],
                "reserved_gb": _vram_snapshot(device)["reserved_gb"],
                "error_stage": _failure_stage({
                    **row,
                    "response": response,
                    "correct": _quality_pass(response, row),
                    "retrieval_target_found": _retrieval_hit(selected, row),
                }),
            }
        )
    return _summarize(output_rows, include_reader=True)


def _base_rows(
    model: Any,
    tokenizer: Any,
    records: list[dict[str, Any]],
    rows: list[dict[str, Any]],
    device: torch.device,
    *,
    mode: str,
    rag_top_k: int,
    max_new_tokens: int,
) -> dict[str, Any]:
    index = _build_rag_index(records) if mode == "chunk_rag" else []
    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"Qwen {mode} generation: {row_index}/{len(rows)}", flush=True)
        query = str(row["query"])
        retrieve_ms = 0.0
        retrieved: list[dict[str, Any]] = []
        if mode == "chunk_rag":
            retrieved, retrieve_ms = _rag_retrieve(index, query, top_k=rag_top_k)
            prompt = _rag_prompt(query, retrieved)
        else:
            prompt = query
        generated = _base_generate(model, tokenizer, prompt, device, max_new_tokens=max_new_tokens)
        response = str(generated.get("response", ""))
        output_rows.append(
            {
                **row,
                "response": response,
                "correct": _quality_pass(response, row),
                "retrieval_target_found": _retrieval_hit(
                    [
                        {
                            "value": item.get("value", ""),
                            "text_preview": item.get("text", ""),
                        }
                        for item in retrieved
                    ],
                    row,
                ),
                "retrieved_values": [item.get("value", "") for item in retrieved],
                "retrieved_ids": [],
                "selected_records": [],
                "prefix_used": bool(mode == "chunk_rag" and 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 _summarize(rows: list[dict[str, Any]], *, include_reader: bool) -> dict[str, Any]:
    answerable = [row for row in rows if bool(row.get("answerable", True))]
    unknown = [row for row in rows if not bool(row.get("answerable", True))]
    by_category: dict[str, dict[str, int]] = {}
    for row in rows:
        category = str(row.get("category", "unknown"))
        summary = by_category.setdefault(category, {"cases": 0, "correct": 0})
        summary["cases"] += 1
        summary["correct"] += int(bool(row.get("correct", False)))
    result: dict[str, Any] = {
        "cases": len(rows),
        "answerable_cases": len(answerable),
        "answerable_correct": sum(int(row.get("correct", False)) for row in answerable),
        "answerable_accuracy": sum(int(row.get("correct", False)) for row in answerable) / max(1, len(answerable)),
        "unknown_cases": len(unknown),
        "unknown_correct": sum(int(row.get("correct", False)) for row in unknown),
        "unknown_refusal_accuracy": sum(int(row.get("correct", False)) for row in unknown) / max(1, len(unknown)),
        "overall_correct": sum(int(row.get("correct", False)) for row in rows),
        "overall_accuracy": sum(int(row.get("correct", False)) for row in rows) / max(1, len(rows)),
        "retrieval_target_found": sum(int(row.get("retrieval_target_found", False)) for row in answerable),
        "answerable_retrieval_recall": sum(int(row.get("retrieval_target_found", False)) for row in answerable) / max(1, len(answerable)),
        "answerable_failure_stage_counts": dict(
            sorted(
                Counter(
                    stage
                    for row in answerable
                    if (stage := row.get("error_stage")) is not None
                ).items()
            )
        ),
        "unknown_failure_stage_counts": dict(
            sorted(
                Counter(
                    stage
                    for row in unknown
                    if (stage := row.get("error_stage")) is not None
                ).items()
            )
        ),
        "mean_prompt_tokens": mean(float(row.get("prompt_tokens", 0)) for row in rows) if rows else 0.0,
        "mean_total_latency_ms": mean(float(row.get("total_latency_s", 0.0)) for row in rows) * 1000.0 if rows else 0.0,
        "mean_decode_tok_s": mean(float(row.get("decode_tok_s", 0.0)) for row in rows) if rows else 0.0,
        "peak_allocated_gb": max((float(row.get("allocated_gb") or 0.0) for row in rows), default=0.0),
        "peak_reserved_gb": max((float(row.get("reserved_gb") or 0.0) for row in rows), default=0.0),
        "by_category": by_category,
        "rows": rows,
    }
    if include_reader:
        result["mean_reader_ms"] = mean(float(row.get("reader_ms", 0.0)) for row in rows) if rows else 0.0
        result["prefix_used_rate"] = sum(int(row.get("prefix_used", False)) for row in rows) / max(1, len(rows))
        result["mean_prefix_tokens"] = mean(float(row.get("prefix_tokens", 0)) for row in rows) if rows else 0.0
    else:
        result["mean_retriever_ms"] = mean(float(row.get("retriever_ms", 0.0)) for row in rows) if rows else 0.0
    return result


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("--memory-model", default=r"H:\Memory\V2_dpskw\qwen3_5_4b_natural_memory_v2")
    parser.add_argument("--output", default=str(PROJECT_ROOT / "dirty_real_corpus_compare_4b.json"))
    parser.add_argument("--project-records", type=int, default=8192)
    parser.add_argument("--project-targets", type=int, default=128)
    parser.add_argument("--rag-top-k", type=int, default=2)
    parser.add_argument("--max-new-tokens", type=int, default=32)
    parser.add_argument(
        "--encode-batch-size",
        type=int,
        default=1,
        help="Qwen semantic-key batch; 1 is intentional to keep the 12 GiB GPU safe",
    )
    parser.add_argument("--gpu-memory-gb", type=float, default=10.0)
    parser.add_argument("--no-4bit", action="store_true")
    parser.add_argument(
        "--skip-chunk-rag",
        action="store_true",
        help="run only the requested base Qwen no-memory baseline and Natural Memory",
    )
    parser.add_argument(
        "--skip-no-memory",
        action="store_true",
        help="skip the base-model generation pass for fast Natural Memory iteration",
    )
    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)}"
    )

    max_memory = _max_memory(args.gpu_memory_gb)
    use_4bit = not args.no_4bit
    qwen_no_memory = None
    qwen_chunk_rag = None
    base_load_vram = None
    if not args.skip_no_memory:
        print("loading original Qwen3.5 4B")
        base = load_qwen_base(_path(args.base_model), load_in_4bit=use_4bit, max_memory=max_memory)
        base.eval()
        base_device = base.get_input_embeddings().weight.device
        base_load_vram = _vram_snapshot(base_device)
        qwen_no_memory = _base_rows(
            base,
            tokenizer,
            all_records,
            all_rows,
            base_device,
            mode="no_memory",
            rag_top_k=args.rag_top_k,
            max_new_tokens=max(1, args.max_new_tokens),
        )
        if not args.skip_chunk_rag:
            qwen_chunk_rag = _base_rows(
                base,
                tokenizer,
                all_records,
                all_rows,
                base_device,
                mode="chunk_rag",
                rag_top_k=max(1, args.rag_top_k),
                max_new_tokens=max(1, args.max_new_tokens),
            )
        _release(base)
        base = None

    print("loading Natural Memory v2 4B")
    memory = load_qwen_dynamic(_path(args.memory_model), load_in_4bit=use_4bit, max_memory=max_memory)
    memory.configure_memory_grounding_guard(tokenizer)
    memory_device = memory._find_layer_device()
    natural_user = _natural_rows(
        memory,
        tokenizer,
        user_records,
        user_rows,
        memory_device,
        max_new_tokens=max(1, args.max_new_tokens),
        encode_batch_size=max(1, args.encode_batch_size),
    )
    memory.memory_os_v2 = None
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    memory.memory_os_v2 = memory._new_memory_os_v2(memory.memory.hidden_size)
    memory.reset_memory(batch_size=1, device=memory_device)
    natural_code = _natural_rows(
        memory,
        tokenizer,
        code_records,
        code_rows,
        memory_device,
        max_new_tokens=max(1, args.max_new_tokens),
        encode_batch_size=max(1, args.encode_batch_size),
    )
    natural_all = _summarize(natural_user["rows"] + natural_code["rows"], include_reader=True)
    natural_load_vram = _vram_snapshot(memory_device)
    memory_stats = memory.memory_os_v2.stats()
    _release(memory)

    systems: dict[str, Any] = {
        "natural_memory_v2": {
            "load_vram": natural_load_vram,
            "user_dialogue": natural_user,
            "real_repository": natural_code,
            "all": natural_all,
            "memory_stats": memory_stats,
        },
    }
    if qwen_no_memory is not None:
        systems["qwen35_4b_no_memory"] = {"load_vram": base_load_vram, "all": qwen_no_memory}
    if qwen_chunk_rag is not None:
        systems["qwen35_4b_chunk_rag"] = {"load_vram": base_load_vram, "all": qwen_chunk_rag}

    report = {
        "benchmark": "dirty_real_corpus_compare_4b",
        "created_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
        "base_model": str(_path(args.base_model)),
        "memory_model": str(_path(args.memory_model)),
        "quantization": "4bit_nf4" if use_4bit else "none",
        "gpu_memory_cap_gb": float(args.gpu_memory_gb),
        "data_provenance": {
            "real_repository": True,
            "repository_root": str(PROJECT_ROOT),
            "user_dialogue": "short facts explicitly stated by the user in this project conversation",
            "synthetic_hardset_used": False,
            "hand_written_paraphrases": True,
            "privacy_note": "No external private conversation export was found or used; the user slice is limited to project requirements already present in this conversation.",
        },
        "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",
            "max_new_tokens": int(args.max_new_tokens),
            "rag_top_k": int(args.rag_top_k),
            "memory_top_k_records": 2,
            "memory_storage": "embedded/process RAM; bounded Natural Memory GPU cache; no SQLite or disk paging",
            "vram_note": "per-query current allocated/reserved snapshots; not torch.cuda.max_memory_allocated high-water marks",
        },
        "systems": systems,
        "interpretation": [
            "This is a transfer/stress benchmark over real repository text plus a user-provided dialogue slice; it is not a public benchmark score.",
            "The repository is intentionally dirty: source, docs and prior engineering artifacts are mixed together.",
            "Natural Memory stores all sampled real chunks, but only explicit symbol and operational targets receive Qwen semantic keys in this bounded run; this is a documented limitation, not hidden coverage.",
            "Natural-language paraphrases are hand-written stress variants, not claimed human-authored logs.",
            "Exact-answer anchors measure whether the model used the evidence; they do not replace human assessment of explanation quality.",
        ],
    }
    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({name: _compact(system["all"]) for name, system in report["systems"].items()}, ensure_ascii=False, indent=2))
    print(f"saved={output}")


if __name__ == "__main__":
    main()