Add Natural Memory architecture and tooling
This commit is contained in:
WpyQwq
@@ -0,0 +1,162 @@
|
||||
"""Stress the durable Natural Memory v2 page tier without loading Qwen."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import random
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from tempfile import TemporaryDirectory
|
||||
|
||||
import torch
|
||||
|
||||
if __package__ in {None, ""}:
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
from dynamic_memory_lab.memory_os_v2 import MemoryRouterV2, PagedMemoryBankV2
|
||||
from dynamic_memory_lab.tiered_memory_store_v2 import TieredMemoryStoreV2
|
||||
|
||||
|
||||
def run(args: argparse.Namespace) -> dict[str, object]:
|
||||
random.seed(args.seed)
|
||||
torch.manual_seed(args.seed)
|
||||
report_path = Path(args.output)
|
||||
report_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with TemporaryDirectory(prefix="natural-memory-v2-tiered-", dir=str(report_path.parent)) as temp_dir:
|
||||
store_path = Path(temp_dir) / "memory.sqlite"
|
||||
router = MemoryRouterV2(args.hidden_size, router_dim=args.key_dim, num_heads=args.heads)
|
||||
store = TieredMemoryStoreV2(
|
||||
store_path,
|
||||
key_dim=args.key_dim,
|
||||
page_capacity=args.page_capacity,
|
||||
)
|
||||
bank = PagedMemoryBankV2(
|
||||
args.hidden_size,
|
||||
page_capacity=args.page_capacity,
|
||||
max_pages=max(1, (args.records + args.page_capacity - 1) // args.page_capacity + 8),
|
||||
hot_pages=args.hot_pages,
|
||||
top_k_pages=args.top_k_pages,
|
||||
top_k_records=args.top_k_records,
|
||||
router=router,
|
||||
key_dim=args.key_dim,
|
||||
coarse_index_bits=args.coarse_index_bits,
|
||||
tier_store=store,
|
||||
max_resident_pages=args.resident_pages,
|
||||
)
|
||||
target_key = None
|
||||
target_text = ""
|
||||
started = time.perf_counter()
|
||||
for start in range(0, args.records, args.batch_size):
|
||||
batch: list[dict[str, object]] = []
|
||||
for index in range(start, min(args.records, start + args.batch_size)):
|
||||
key = torch.randn(args.hidden_size)
|
||||
if index == args.target_index:
|
||||
target_key = key.clone()
|
||||
target_text = f"tiered-record-{index}"
|
||||
batch.append(
|
||||
{
|
||||
"text": f"tiered-record-{index}",
|
||||
"key": key,
|
||||
"summary": key,
|
||||
"entity": "benchmark",
|
||||
"attribute": f"attribute-{index}",
|
||||
"value": f"value-{index}",
|
||||
"importance": 0.2 if index != args.target_index else 1.0,
|
||||
"confidence": 0.95,
|
||||
"source": "tiered-benchmark",
|
||||
"trusted": True,
|
||||
}
|
||||
)
|
||||
bank.write_batch(batch)
|
||||
write_seconds = time.perf_counter() - started
|
||||
before = bank.stats()
|
||||
store.close()
|
||||
|
||||
reopen_started = time.perf_counter()
|
||||
reopened_store = TieredMemoryStoreV2(
|
||||
store_path,
|
||||
key_dim=args.key_dim,
|
||||
page_capacity=args.page_capacity,
|
||||
)
|
||||
reopened = PagedMemoryBankV2(
|
||||
args.hidden_size,
|
||||
page_capacity=args.page_capacity,
|
||||
max_pages=max(1, (args.records + args.page_capacity - 1) // args.page_capacity + 8),
|
||||
hot_pages=args.hot_pages,
|
||||
top_k_pages=args.top_k_pages,
|
||||
top_k_records=args.top_k_records,
|
||||
router=router,
|
||||
key_dim=args.key_dim,
|
||||
coarse_index_bits=args.coarse_index_bits,
|
||||
tier_store=reopened_store,
|
||||
max_resident_pages=args.resident_pages,
|
||||
)
|
||||
reopen_seconds = time.perf_counter() - reopen_started
|
||||
if target_key is None:
|
||||
raise RuntimeError("target index was not generated")
|
||||
records, decision = reopened.query(
|
||||
query_key=target_key,
|
||||
query_text=target_text,
|
||||
top_k_pages=args.top_k_pages,
|
||||
top_k_records=args.top_k_records,
|
||||
)
|
||||
after = reopened.stats()
|
||||
found = any(record.text == target_text for record in records)
|
||||
target_row = reopened_store.find_by_text(target_text, active_status="active")
|
||||
target_page_id = target_row["page_id"] if target_row is not None else None
|
||||
candidate_pages = reopened._candidate_page_ids(target_key)
|
||||
reopened_store.close()
|
||||
|
||||
report = {
|
||||
"format_version": 2,
|
||||
"records_requested": args.records,
|
||||
"target_index": args.target_index,
|
||||
"page_capacity": args.page_capacity,
|
||||
"coarse_index_bits": args.coarse_index_bits,
|
||||
"resident_pages": args.resident_pages,
|
||||
"write_seconds": write_seconds,
|
||||
"reopen_seconds": reopen_seconds,
|
||||
"before_close": before,
|
||||
"after_reopen": after,
|
||||
"target_recalled_after_restart": found,
|
||||
"target_page_id": target_page_id,
|
||||
"target_page_in_coarse_candidates": target_page_id in candidate_pages if target_page_id else False,
|
||||
"coarse_candidate_count": len(candidate_pages),
|
||||
"decision": {
|
||||
"page_ids": decision.page_ids,
|
||||
"record_ids": decision.record_ids,
|
||||
"hop_count": decision.hop_count,
|
||||
"confidence": decision.confidence,
|
||||
"stop_reason": decision.stop_reason,
|
||||
},
|
||||
}
|
||||
report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
return report
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--records", type=int, default=100_000)
|
||||
parser.add_argument("--target-index", type=int, default=99_999)
|
||||
parser.add_argument("--hidden-size", type=int, default=32)
|
||||
parser.add_argument("--key-dim", type=int, default=16)
|
||||
parser.add_argument("--heads", type=int, default=4)
|
||||
parser.add_argument("--page-capacity", type=int, default=32)
|
||||
parser.add_argument("--resident-pages", type=int, default=64)
|
||||
parser.add_argument("--hot-pages", type=int, default=8)
|
||||
parser.add_argument("--top-k-pages", type=int, default=4)
|
||||
parser.add_argument("--top-k-records", type=int, default=8)
|
||||
parser.add_argument("--coarse-index-bits", type=int, default=12)
|
||||
parser.add_argument("--batch-size", type=int, default=2_000)
|
||||
parser.add_argument("--seed", type=int, default=20260904)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
default="W:/Flash/model/dynamic_memory_lab/tiered_memory_v2_benchmark.json",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(json.dumps(run(parse_args()), ensure_ascii=False, indent=2))
|
||||
Reference in New Issue
Block a user