"""Low-pressure end-to-end test for the embedded Natural Memory v2 path. This benchmark intentionally uses only the model package's third safetensors memory shard. It does not create SQLite files or exercise disk paging. The long-context cases lower the temporary KV budget so the test measures the model-owned archive/read path without asking a 12 GiB GPU to hold a huge KV. """ from __future__ import annotations import argparse import gc import json import random import time from pathlib import Path from typing import Any import torch from .qwen_integration import load_qwen_dynamic, load_tokenizer def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--model-path", default=r"W:\Flash\model\dynamic_memory_lab\qwen3_5_4b_natural_memory_v2", ) parser.add_argument("--lengths", default="4096,8192,16384,32768") parser.add_argument("--kv-budget", type=int, default=2048) parser.add_argument("--chunk-tokens", type=int, default=1024) parser.add_argument("--max-new-tokens", type=int, default=8) parser.add_argument("--no-4bit", action="store_true") parser.add_argument( "--output", default=r"W:\Flash\model\dynamic_memory_lab\embedded_memory_v2_long_benchmark.json", ) return parser.parse_args() def build_prompt(tokenizer: Any, target_tokens: int, seed: int) -> tuple[str, str, int]: rng = random.Random(seed + target_tokens) answer = f"EMBEDDED-LONG-{target_tokens}-{rng.randrange(100000, 999999)}" needle = f"长期记忆锚点：唯一编号是 {answer}。" filler = ( "这是长文本记忆压力测试中的普通背景段落，包含项目说明、日期、日志和无关备注。" "这些内容不是问题答案，读取时应保留原文但忽略干扰。" ) chunks: list[str] = [] while len(tokenizer(" ".join(chunks + [filler, needle]), add_special_tokens=False)["input_ids"]) < target_tokens: chunks.append(filler) # Keep the needle safely inside the archived prefix even in the smallest # case, so a pass must come from memory rather than the retained window. pivot = max(1, len(chunks) // 3) material = " ".join(chunks[:pivot] + [needle] + chunks[pivot:]) prompt = ( "请阅读下面的长材料，回答末尾问题，只输出编号，不要解释。\n" "---开始材料---\n" f"{material}\n" "---结束材料---\n" "问题：长期记忆锚点的唯一编号是什么？" ) prompt_tokens = len(tokenizer(prompt, add_special_tokens=False)["input_ids"]) return prompt, answer, prompt_tokens def chat_inputs(tokenizer: Any, text: str, device: torch.device) -> dict[str, torch.Tensor]: encoded = tokenizer.apply_chat_template( [{"role": "user", "content": text}], tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True, enable_thinking=False, ) return { key: value.to(device) for key, value in encoded.items() if isinstance(value, torch.Tensor) } def run_case(model: Any, tokenizer: Any, target_tokens: int, args: argparse.Namespace) -> dict[str, Any]: device = model._find_layer_device() model.reset_memory(batch_size=1, device=device) assert model.memory_os_v2 is not None model.memory_os_v2.kv_budget.max_tokens = int(args.kv_budget) prompt, answer, prompt_tokens = build_prompt(tokenizer, target_tokens, 20260904) encoded = chat_inputs(tokenizer, prompt, device) query_text = "长期记忆锚点的唯一编号" query = tokenizer(query_text, add_special_tokens=False, return_tensors="pt") query_ids = query["input_ids"].to(device) query_mask = query.get("attention_mask") if query_mask is None: query_mask = torch.ones_like(query_ids) query_mask = query_mask.to(device) started = time.perf_counter() with torch.inference_mode(): output = model.generate( **encoded, max_new_tokens=args.max_new_tokens, do_sample=False, use_cache=True, update_memory=False, pad_token_id=tokenizer.pad_token_id, memory_query_input_ids=query_ids, memory_query_attention_mask=query_mask, memory_query_text=query_text, ) if device.type == "cuda": torch.cuda.synchronize(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() stats = model.memory_v2_stats() records = model.memory_os_v2.bank.records archived_text_hit = any( answer in tokenizer.decode(record.token_ids.tolist(), skip_special_tokens=True) for record in records.values() if record.memory_type == "context_chunk" and isinstance(record.token_ids, torch.Tensor) ) query_key = model._encode_model_key(query_ids, query_mask)[0] retrieved, decision = model.read_hierarchical_memory( query_key, query_text=query_text, query_token_ids=query_ids[0], top_k_pages=model.memory_config.memory_top_k_pages, top_k_records=model.memory_config.memory_top_k_records, max_hops=model.memory_config.memory_max_hops, ) retrieved_hit = any( answer in tokenizer.decode(record.token_ids.tolist(), skip_special_tokens=True) for record in retrieved if isinstance(record.token_ids, torch.Tensor) ) return { "target_tokens": target_tokens, "prompt_tokens": prompt_tokens, "kv_budget_tokens": args.kv_budget, "chunk_tokens": args.chunk_tokens, "response": response, "expected": answer, "generation_hit": answer in response, "archived_text_hit": archived_text_hit, "retrieved_text_hit": retrieved_hit, "retrieved_records": len(retrieved), "router_stop_reason": decision.stop_reason, "router_hop_count": decision.hop_count, "seconds": elapsed, "records": stats.get("records", 0), "pages": stats.get("pages", 0), "gpu_cache_records": stats.get("gpu_cache_records", 0), "gpu_cache_tokens": stats.get("gpu_cache_tokens", 0), "gpu_cache_device": stats.get("gpu_cache_device", "none"), "passed": bool(archived_text_hit and retrieved_hit), } def main() -> None: args = parse_args() lengths = [int(item.strip()) for item in args.lengths.split(",") if item.strip()] tokenizer = load_tokenizer(args.model_path) model = load_qwen_dynamic(args.model_path, load_in_4bit=not args.no_4bit) model.eval() model.memory_config.context_chunk_tokens = int(args.chunk_tokens) report: dict[str, Any] = { "benchmark": "Natural Memory v2 embedded third-shard long-memory test", "storage_mode": model.memory_config.memory_storage_mode, "tier_store_enabled": bool(model.memory_os_v2 and model.memory_os_v2.bank.tier_store is not None), "model_path": str(Path(args.model_path).resolve()), "lengths": lengths, "quantization": "4bit_nf4" if not args.no_4bit else "none", "kv_budget_tokens": args.kv_budget, "chunk_tokens": args.chunk_tokens, "rows": [], } try: if report["storage_mode"] != "embedded" or report["tier_store_enabled"]: raise RuntimeError("embedded benchmark requires memory_storage_mode=embedded and no tier store") for target_tokens in lengths: row = run_case(model, tokenizer, target_tokens, args) report["rows"].append(row) print(json.dumps(row, ensure_ascii=False)) finally: model.reset_memory(batch_size=1, device=model._find_layer_device()) model.close_memory_storage() del model gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() report["passed_cases"] = sum(bool(row["passed"]) for row in report["rows"]) report["total_cases"] = len(report["rows"]) report["all_passed"] = report["passed_cases"] == report["total_cases"] 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(report, ensure_ascii=False, indent=2)) print(f"saved={output}") if __name__ == "__main__": main()