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,64 @@
+"""Evaluate a Qwen dynamic-memory adapter on streaming JSONL records."""
+
+from __future__ import annotations
+
+import argparse
+
+import torch
+
+from .qwen_integration import load_qwen_dynamic, load_tokenizer
+from .train_qwen_memory import encode_messages, load_records, pad_batch
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", default=".")
+    parser.add_argument("--data", default="dynamic_memory_lab/data/demo_stream.jsonl")
+    parser.add_argument("--adapter", default="dynamic_memory_lab/qwen_memory_adapter")
+    parser.add_argument("--max-length", type=int, default=512)
+    parser.add_argument("--no-4bit", action="store_true")
+    args = parser.parse_args()
+
+    tokenizer = load_tokenizer(args.model_path)
+    model = load_qwen_dynamic(args.model_path, load_in_4bit=not args.no_4bit)
+    model.load_memory_adapter(args.adapter)
+    model.eval()
+    device = model._find_layer_device()
+    pad_id = int(tokenizer.pad_token_id)
+    records = load_records(args.data)
+    correct = 0
+
+    for record in records:
+        memory_input, memory_mask, _ = pad_batch(
+            [encode_messages(tokenizer, record["memory"], args.max_length)], pad_id
+        )
+        query_input, query_mask, query_labels = pad_batch(
+            [encode_messages(tokenizer, record["query"], args.max_length)], pad_id
+        )
+        with torch.no_grad():
+            memory_output = model(
+                input_ids=memory_input.to(device),
+                attention_mask=memory_mask.to(device),
+                read_memory=False,
+                update_memory=True,
+            )
+            output = model(
+                input_ids=query_input.to(device),
+                attention_mask=query_mask.to(device),
+                memory_state=memory_output.memory,
+                read_memory=True,
+                update_memory=False,
+            )
+        labels = query_labels.to(device)
+        shifted_labels = labels[..., 1:]
+        predictions = output.logits[..., :-1, :].argmax(dim=-1)
+        target_positions = shifted_labels != -100
+        sequence_ok = bool((predictions[target_positions] == shifted_labels[target_positions]).all())
+        correct += int(sequence_ok)
+        print(f"sequence_ok={sequence_ok}")
+
+    print(f"exact_sequence_accuracy={correct / len(records):.3f}")
+
+
+if __name__ == "__main__":
+    main()