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,66 @@
+"""Check whether teacher-forcing and cached greedy generation agree on token 1."""
+
+from __future__ import annotations
+
+import json
+
+import torch
+
+from .benchmark_qwen import _generation_prompt
+from .qwen_integration import QwenMemoryConfig, load_qwen_dynamic, load_tokenizer
+from .train_qwen_memory import encode_messages, pad_batch
+
+
+def main() -> None:
+    model = load_qwen_dynamic(
+        ".",
+        memory_config=QwenMemoryConfig(mode="blend", blend_init=0.1),
+        load_in_4bit=True,
+    )
+    model.load_memory_adapter("dynamic_memory_lab/qwen_memory_adapter_full")
+    model.eval()
+    tokenizer = load_tokenizer(".")
+    record = json.loads(
+        next(open("dynamic_memory_lab/data/benchmark_eval.jsonl", encoding="utf-8"))
+    )
+    device = model._find_layer_device()
+    memory = encode_messages(tokenizer, record["memory"], 128)
+    memory_input, memory_mask, _ = pad_batch([memory], int(tokenizer.pad_token_id))
+    prompt = _generation_prompt(tokenizer, record["query"][:-1], device)
+
+    with torch.inference_mode():
+        memory_state = 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,
+        ).memory
+        teacher_forcing = model(
+            input_ids=prompt["input_ids"],
+            attention_mask=prompt["attention_mask"],
+            memory_state=memory_state,
+            read_memory=True,
+            update_memory=False,
+            return_memory=True,
+            use_cache=False,
+        )
+        generated = model.generate(
+            **prompt,
+            max_new_tokens=1,
+            do_sample=False,
+            use_cache=True,
+            pad_token_id=tokenizer.pad_token_id,
+        )
+
+    teacher_id = int(teacher_forcing.logits[0, -1].argmax())
+    generated_id = int(generated[0, -1])
+    print(json.dumps({"expected": record["answer"]}, ensure_ascii=True))
+    print(json.dumps({"teacher_id": teacher_id, "teacher_text": tokenizer.decode([teacher_id])}, ensure_ascii=True))
+    print(json.dumps({"generated_id": generated_id, "generated_text": tokenizer.decode([generated_id])}, ensure_ascii=True))
+    print(f"same={teacher_id == generated_id}")
+
+
+if __name__ == "__main__":
+    main()