natural-memory-nm21/benchmark_compare_baselines_4b.py

"""Run a same-protocol 4B comparison: base Qwen, Chunk RAG, and Natural Memory.

The benchmark uses the same general-memory and real-repository query sets for
all systems.  It reports answer correctness, refusal correctness, prompt/read
overhead, generation speed, and peak VRAM.  Memory is kept in process RAM and
the GPU placement cap is shared by every model load.
"""

from __future__ import annotations

import argparse
import gc
import json
import math
import re
import time
from pathlib import Path
from statistics import mean
from typing import Any

import torch

from .benchmark_real_scale_memory_4b import (
    _build_general_records,
    _build_project_records,
    _chat_generate,
    _contains_answer,
    _is_refusal,
    _max_memory,
    _native_cases,
    _path,
    _prepare_records,
    _set_cuda_process_cap,
    _source_files,
    _sync,
)
from .qwen_integration import load_qwen_base, load_qwen_dynamic, load_tokenizer
from .stream_chat_qwen_memory import _chat_tensor


PROJECT_ROOT = Path(__file__).resolve().parent
TERM_PATTERN = re.compile(r"[\u4e00-\u9fff]|[A-Za-z0-9_\-]+")


def _terms(text: str) -> set[str]:
    return set(TERM_PATTERN.findall(str(text).lower()))


def _vram_snapshot(device: torch.device) -> dict[str, float | None]:
    if device.type != "cuda":
        return {"allocated_gb": None, "reserved_gb": None}
    return {
        "allocated_gb": torch.cuda.memory_allocated(device) / 1024**3,
        "reserved_gb": torch.cuda.memory_reserved(device) / 1024**3,
    }


def _record_entry(record: dict[str, Any]) -> dict[str, Any]:
    text = str(record.get("text", ""))
    return {
        "text": text,
        "terms": _terms(text),
        "entity": str(record.get("entity", "")),
        "attribute": str(record.get("attribute", "")),
        "value": str(record.get("value", "")),
    }


def _build_rag_index(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
    return [_record_entry(record) for record in records if str(record.get("text", "")).strip()]


def _rag_retrieve(
    index: list[dict[str, Any]],
    query: str,
    *,
    top_k: int,
) -> tuple[list[dict[str, Any]], float]:
    query_terms = _terms(query)
    query_lower = query.strip().lower()
    started = time.perf_counter()
    scored: list[tuple[float, int, dict[str, Any]]] = []
    for position, item in enumerate(index):
        shared = len(query_terms.intersection(item["terms"]))
        score = shared / math.sqrt(max(1, len(query_terms) * len(item["terms"])))
        entity = item["entity"].strip().lower()
        attribute = item["attribute"].strip().lower()
        if len(entity) >= 4 and entity in query_lower:
            score += 5.0
        if attribute and attribute in query_lower:
            score += 1.0
        scored.append((score, -position, item))
    scored.sort(key=lambda value: (value[0], value[1]), reverse=True)
    elapsed_ms = (time.perf_counter() - started) * 1000.0
    return [item for _, _, item in scored[: max(1, int(top_k))]], elapsed_ms


def _rag_prompt(query: str, records: list[dict[str, Any]]) -> str:
    evidence = "\n".join(
        f"[证据 {index}] {record['text']}"
        for index, record in enumerate(records, 1)
    )
    return (
        "以下是检索器返回的记忆证据。只能使用证据中明确出现的事实，"
        "先核对实体、属性和已确认值；多个候选并存时不要把候选拼成一个事实；"
        "新旧冲突时优先最新且来源更可靠的证据；找不到目标时请明确说不知道，"
        "不要用相似用户的信息代替。\n"
        "---记忆证据开始---\n"
        + evidence
        + "\n---记忆证据结束---\n问题："
        + query
    )


@torch.inference_mode()
def _base_generate(
    model: Any,
    tokenizer: Any,
    prompt: str,
    device: torch.device,
    *,
    max_new_tokens: int,
) -> dict[str, Any]:
    encoded = _chat_tensor(tokenizer, prompt)
    encoded = {key: value.to(device) for key, value in encoded.items()}
    if device.type == "cuda":
        torch.cuda.reset_peak_memory_stats(device)
    _sync(device)
    started = time.perf_counter()
    output = model.generate(
        **encoded,
        max_new_tokens=max_new_tokens,
        do_sample=False,
        use_cache=True,
        pad_token_id=tokenizer.pad_token_id,
    )
    _sync(device)
    elapsed = time.perf_counter() - started
    response_ids = output[0, encoded["input_ids"].shape[1] :]
    response = tokenizer.decode(
        response_ids.detach().cpu().tolist(),
        skip_special_tokens=True,
    ).strip()
    generated_tokens = int(response_ids.numel())
    row = {
        "status": "ok",
        "response": response,
        "prompt_tokens": int(encoded["input_ids"].shape[1]),
        "generated_tokens": generated_tokens,
        "total_latency_s": elapsed,
        "decode_tok_s": generated_tokens / max(elapsed, 1e-9),
    }
    row.update(_vram_snapshot(device))
    del output
    return row


def _quality_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 {
        "cases": len(rows),
        "answerable_cases": len(answerable),
        "answerable_correct": sum(int(row["correct"]) for row in answerable),
        "answerable_accuracy": sum(int(row["correct"]) for row in answerable) / max(1, len(answerable)),
        "unknown_cases": len(unknown),
        "unknown_correct": sum(int(row["correct"]) for row in unknown),
        "unknown_refusal_accuracy": sum(int(row["correct"]) for row in unknown) / max(1, len(unknown)),
        "mean_prompt_tokens": mean(row["prompt_tokens"] for row in rows) if rows else 0.0,
        "mean_total_latency_ms": mean(row["total_latency_s"] for row in rows) * 1000.0 if rows else 0.0,
        "mean_decode_tok_s": mean(row["decode_tok_s"] for row in rows) if rows else 0.0,
        "peak_allocated_gb": max((row.get("allocated_gb") or 0.0 for row in rows), default=0.0),
        "peak_reserved_gb": max((row.get("reserved_gb") or 0.0 for row in rows), default=0.0),
    }


def _run_base_system(
    model: Any,
    tokenizer: Any,
    rows: list[dict[str, Any]],
    index: list[dict[str, Any]],
    device: torch.device,
    *,
    mode: str,
    rag_top_k: int,
    max_new_tokens: int,
) -> dict[str, Any]:
    output_rows: list[dict[str, Any]] = []
    for row in rows:
        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,
        )
        expected = str(row.get("expected", ""))
        answerable = bool(row.get("answerable", True))
        correct = (
            _contains_answer(generated["response"], expected)
            if answerable
            else _is_refusal(generated["response"])
        )
        retrieved_target = answerable and any(
            _contains_answer(item["value"], expected)
            or _contains_answer(item["text"], expected)
            for item in retrieved
        )
        output_rows.append(
            {
                "id": row.get("id", ""),
                "query": query,
                "expected": expected,
                "answerable": answerable,
                "correct": bool(correct),
                "retrieved_target": bool(retrieved_target),
                "retrieved_count": len(retrieved),
                "retriever_ms": retrieve_ms,
                "retrieved_values": [item["value"] for item in retrieved],
                **generated,
            }
        )
    summary = _quality_summary(output_rows)
    summary["mean_retriever_ms"] = mean(row["retriever_ms"] for row in output_rows) if output_rows else 0.0
    reader_values = [float(row["reader_ms"]) for row in output_rows if "reader_ms" in row]
    if reader_values:
        summary["mean_reader_ms"] = mean(reader_values)
    summary["answerable_retrieval_recall"] = (
        sum(int(row["retrieved_target"]) for row in output_rows if row["answerable"])
        / max(1, sum(int(row["answerable"]) for row in output_rows))
    )
    summary["rows"] = output_rows
    return summary


def _run_natural_memory(
    model: Any,
    tokenizer: Any,
    general_records: list[dict[str, Any]],
    general_rows: list[dict[str, Any]],
    project_records: list[dict[str, Any]],
    project_rows: list[dict[str, Any]],
    device: torch.device,
    *,
    generation_general: int,
    generation_project: int,
    max_new_tokens: int,
    encode_batch_size: int,
) -> dict[str, Any]:
    model.eval()
    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)

    output: dict[str, Any] = {}
    prepared_general = _prepare_records(
        model,
        tokenizer,
        general_records,
        device,
        batch_size=max(1, encode_batch_size),
    )
    model.memory_os_v2.write_batch(prepared_general)
    general_rows = general_rows[: max(1, generation_general)]
    general_output: list[dict[str, Any]] = []
    for row in general_rows:
        generated = _chat_generate(
            model,
            tokenizer,
            str(row["query"]),
            device,
            max_new_tokens,
        )
        expected = str(row.get("expected", ""))
        answerable = bool(row.get("answerable", True))
        correct = (
            _contains_answer(generated["response"], expected)
            if answerable
            else _is_refusal(generated["response"])
        )
        general_output.append(
            {
                "id": row.get("id", ""),
                "query": row["query"],
                "expected": expected,
                "answerable": answerable,
                "correct": bool(correct),
                "prompt_tokens": int(generated["public_prompt_tokens"] + generated["prefix_tokens"]),
                "generated_tokens": int(generated["generated_tokens"]),
                "total_latency_s": float(generated["generation_seconds"]),
                "decode_tok_s": int(generated["generated_tokens"]) / max(float(generated["generation_seconds"]), 1e-9),
                "reader_ms": float(generated["memory_read_seconds"]) * 1000.0,
                "prefix_used": bool(generated["prefix_used"]),
                "prefix_tokens": int(generated["prefix_tokens"]),
                "retrieved_values": list(generated["selected_values"]),
                "retrieved_ids": list(generated["selected_record_ids"]),
                "allocated_gb": _vram_snapshot(device)["allocated_gb"],
                "reserved_gb": _vram_snapshot(device)["reserved_gb"],
                "response": generated["response"],
            }
        )
    output["general"] = _quality_summary(general_output)
    output["general"]["rows"] = general_output

    model.memory_os_v2 = model._new_memory_os_v2(model.memory.hidden_size)
    model.reset_memory(batch_size=1, device=device)
    prepared_project = _prepare_records(
        model,
        tokenizer,
        project_records,
        device,
        batch_size=max(1, encode_batch_size),
    )
    model.memory_os_v2.write_batch(prepared_project)
    project_rows = project_rows[: max(1, generation_project)]
    project_output: list[dict[str, Any]] = []
    for row in project_rows:
        generated = _chat_generate(
            model,
            tokenizer,
            str(row["query"]),
            device,
            max_new_tokens,
        )
        expected = str(row.get("expected", ""))
        correct = _contains_answer(generated["response"], expected)
        snapshot = _vram_snapshot(device)
        project_output.append(
            {
                "id": row.get("id", ""),
                "query": row["query"],
                "expected": expected,
                "answerable": True,
                "correct": bool(correct),
                "prompt_tokens": int(generated["public_prompt_tokens"] + generated["prefix_tokens"]),
                "generated_tokens": int(generated["generated_tokens"]),
                "total_latency_s": float(generated["generation_seconds"]),
                "decode_tok_s": int(generated["generated_tokens"]) / max(float(generated["generation_seconds"]), 1e-9),
                "reader_ms": float(generated["memory_read_seconds"]) * 1000.0,
                "prefix_used": bool(generated["prefix_used"]),
                "prefix_tokens": int(generated["prefix_tokens"]),
                "retrieved_values": list(generated["selected_values"]),
                "retrieved_ids": list(generated["selected_record_ids"]),
                "allocated_gb": snapshot["allocated_gb"],
                "reserved_gb": snapshot["reserved_gb"],
                "response": generated["response"],
            }
        )
    output["project"] = _quality_summary(project_output)
    output["project"]["rows"] = project_output
    return output


def _release(model: Any) -> None:
    del model
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


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("--data-root", default=str(PROJECT_ROOT / "data"))
    parser.add_argument("--output", default=str(PROJECT_ROOT / "baseline_compare_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("--generation-general", type=int, default=128)
    parser.add_argument("--generation-project", type=int, default=64)
    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=16)
    parser.add_argument("--gpu-memory-gb", type=float, default=10.0)
    parser.add_argument("--no-4bit", action="store_true")
    parser.add_argument(
        "--reuse-original-report",
        default=None,
        help="reuse original_qwen results from an earlier report and rerun only Natural Memory",
    )
    args = parser.parse_args()
    _set_cuda_process_cap(args.gpu_memory_gb)

    tokenizer = load_tokenizer(_path(args.base_model))
    general_cases = _native_cases(_path(args.data_root), max(1, int(args.general_cases)))
    general_raw_records, general_queries = _build_general_records(general_cases)
    general_generation_rows = [
        row for row in general_queries if row["split"] == "eval"
    ][: max(1, int(args.generation_general))]
    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=512,
    )
    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))])
    ]
    general_index = _build_rag_index(general_raw_records)
    project_index = _build_rag_index(project_raw_records)
    use_4bit = not args.no_4bit
    max_memory = _max_memory(args.gpu_memory_gb)

    reused_original_report = None
    if args.reuse_original_report:
        reused_original_report = json.loads(
            _path(args.reuse_original_report).read_text(encoding="utf-8")
        )
        original_system = reused_original_report.get("systems", {}).get("original_qwen")
        if not isinstance(original_system, dict):
            raise ValueError("reuse report does not contain systems.original_qwen")
        base_load_vram = original_system.get("load_vram")
        base_general = original_system["general"]
        base_project = original_system["project"]
        print("reusing original Qwen results")
    else:
        print("loading original Qwen 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)
        base_general = {
            "no_memory": _run_base_system(
                base,
                tokenizer,
                general_generation_rows,
                general_index,
                base_device,
                mode="no_memory",
                rag_top_k=args.rag_top_k,
                max_new_tokens=max(1, args.max_new_tokens),
            ),
            "chunk_rag": _run_base_system(
                base,
                tokenizer,
                general_generation_rows,
                general_index,
                base_device,
                mode="chunk_rag",
                rag_top_k=max(1, args.rag_top_k),
                max_new_tokens=max(1, args.max_new_tokens),
            ),
        }
        base_project = {
            "no_memory": _run_base_system(
                base,
                tokenizer,
                project_generation_rows,
                project_index,
                base_device,
                mode="no_memory",
                rag_top_k=args.rag_top_k,
                max_new_tokens=max(1, args.max_new_tokens),
            ),
            "chunk_rag": _run_base_system(
                base,
                tokenizer,
                project_generation_rows,
                project_index,
                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)
        # Drop the caller's reference as well.  Otherwise the original Qwen
        # remains resident while Natural Memory is loaded below, making the
        # latter's VRAM measurement include two complete 4B models.
        base = None

    print("loading Natural Memory 4B")
    memory = load_qwen_dynamic(
        _path(args.memory_model),
        load_in_4bit=use_4bit,
        max_memory=max_memory,
    )
    memory_device = memory._find_layer_device()
    natural = _run_natural_memory(
        memory,
        tokenizer,
        general_raw_records,
        general_generation_rows,
        project_raw_records,
        project_generation_rows,
        memory_device,
        generation_general=args.generation_general,
        generation_project=args.generation_project,
        max_new_tokens=max(1, args.max_new_tokens),
        encode_batch_size=max(1, args.encode_batch_size),
    )
    natural_load_vram = _vram_snapshot(memory_device)
    _release(memory)

    report = {
        "benchmark": "baseline_compare_4b",
        "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),
        "general_cases": len(general_cases),
        "general_records": len(general_raw_records),
        "project": project_meta,
        "protocol": {
            "generation_general": len(general_generation_rows),
            "generation_project": len(project_generation_rows),
            "rag_top_k": int(args.rag_top_k),
            "max_new_tokens": int(args.max_new_tokens),
            "same_tokenizer": True,
            "same_sampling": "greedy",
            "memory_storage": "process_ram; selected Natural Memory records promoted to bounded GPU cache",
            "reused_original_report": str(_path(args.reuse_original_report)) if args.reuse_original_report else None,
        },
        "systems": {
            "original_qwen": {
                "load_vram": base_load_vram,
                "general": base_general,
                "project": base_project,
            },
            "natural_memory": {
                "load_vram": natural_load_vram,
                "general": natural["general"],
                "project": natural["project"],
            },
        },
        "limitations": [
            "Chunk RAG uses a CPU lexical candidate index with exact entity/attribute bonuses; it is a transparent baseline, not a hosted embedding service.",
            "The general corpus is the local native-memory benchmark; the project corpus is the current repository source and documentation.",
            "Latency for Natural Memory is split into reader_ms and generation is measured by the existing model wrapper; this first comparison prioritizes correctness and VRAM.",
        ],
    }
    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")
    compact = {
        name: {
            "general": {key: value for key, value in system["general"].items() if key != "rows"},
            "project": {key: value for key, value in system["project"].items() if key != "rows"},
        }
        for name, system in report["systems"].items()
    }
    print(json.dumps(compact, ensure_ascii=False, indent=2))
    print(f"saved={output}")


if __name__ == "__main__":
    main()