Add Natural Memory architecture and tooling

2026-09-05 08:53:41 +08:00
parent 0acf8b06ee
commit 516351f0b5
56 changed files with 18319 additions and 0 deletions
@@ -0,0 +1,107 @@
+"""Verify answering after a model restart using only a saved user memory state."""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+from pathlib import Path
+
+import torch
+
+from .benchmark_qwen import _generation_prompt
+from .qwen_integration import load_memory_config, load_qwen_dynamic, load_tokenizer
+from .train_qwen_memory import encode_messages, pad_batch
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model-path", default=".")
+    parser.add_argument(
+        "--adapter",
+        default="dynamic_memory_lab/qwen_memory_adapter_pointer",
+    )
+    parser.add_argument(
+        "--data",
+        default="dynamic_memory_lab/data/benchmark_eval.jsonl",
+    )
+    parser.add_argument("--record-index", type=int, default=0)
+    parser.add_argument(
+        "--memory-state",
+        default="dynamic_memory_lab/data/user_memory_demo.pt",
+    )
+    parser.add_argument("--max-length", type=int, default=128)
+    parser.add_argument("--max-new-tokens", type=int, default=4)
+    parser.add_argument("--no-4bit", action="store_true")
+    args = parser.parse_args()
+
+    records = [
+        json.loads(line)
+        for line in Path(args.data).read_text(encoding="utf-8").splitlines()
+        if line.strip()
+    ]
+    record = records[args.record_index]
+    tokenizer = load_tokenizer(args.model_path)
+    memory_config = load_memory_config(args.adapter)
+
+    model = load_qwen_dynamic(
+        args.model_path,
+        memory_config=memory_config,
+        load_in_4bit=not args.no_4bit,
+    )
+    model.load_memory_adapter(args.adapter)
+    model.eval()
+    device = model._find_layer_device()
+    memory = encode_messages(tokenizer, record["memory"], args.max_length)
+    memory_input, memory_mask, _ = pad_batch([memory], int(tokenizer.pad_token_id))
+    state_path = Path(args.memory_state)
+
+    with torch.inference_mode():
+        model(
+            input_ids=memory_input.to(device),
+            attention_mask=memory_mask.to(device),
+            read_memory=False,
+            update_memory=True,
+            return_memory=True,
+            use_cache=False,
+        )
+        model.save_runtime_memory(state_path)
+    print(f"saved_memory_state={state_path}")
+
+    del model
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    restarted = load_qwen_dynamic(
+        args.model_path,
+        memory_config=memory_config,
+        load_in_4bit=not args.no_4bit,
+    )
+    restarted.load_memory_adapter(args.adapter)
+    restarted.eval()
+    device = restarted._find_layer_device()
+    restarted.load_runtime_memory(state_path, device=device)
+    prompt = _generation_prompt(tokenizer, record["query"][:-1], device)
+    with torch.inference_mode():
+        output = restarted.generate(
+            **prompt,
+            max_new_tokens=args.max_new_tokens,
+            do_sample=False,
+            update_memory=False,
+            use_cache=True,
+            pad_token_id=tokenizer.pad_token_id,
+        )
+    generated = tokenizer.decode(
+        output[0, prompt["input_ids"].shape[1] :],
+        skip_special_tokens=True,
+    ).replace(" ", "").replace("\r", "").replace("\n", "").strip()
+    expected = str(record["answer"])
+    print(f"query_contains_history=false")
+    print(f"expected={expected}")
+    print(f"restarted_generated={generated}")
+    print(f"correct={generated.startswith(expected)}")
+
+
+if __name__ == "__main__":
+    main()