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,112 @@
+"""Create deterministic train/eval data for the dynamic-memory benchmark."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import random
+from pathlib import Path
+
+
+ATTRIBUTES = ("驻地", "负责人", "维护日", "安全级别", "档案类别")
+# Single-character answers remove shared prefixes and make free-generation
+# exact match a meaningful associative-recall metric.
+VALUES = tuple("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789")
+
+MEMORY_USER_TEMPLATES = (
+    "记住这条资料：{subject}的{attribute}代号是{value}。",
+    "请存储信息——对象{subject}的{attribute}为代号{value}。",
+    "新事实：{subject}的{attribute}代号={value}。请记住。",
+)
+MEMORY_ASSISTANT_TEMPLATES = (
+    "已记录：{subject}的{attribute}代号是{value}。",
+    "好的，{subject}的{attribute}已记为代号{value}。",
+    "收到，已经保存{subject}的{attribute}代号：{value}。",
+)
+QUERY_USER_TEMPLATES = (
+    "查询：{subject}的{attribute}代号是什么？",
+    "请问对象{subject}的{attribute}代号为？",
+    "根据已记信息，{subject}的{attribute}代号是？",
+)
+
+
+def make_records(count: int, *, prefix: str, rng: random.Random) -> list[dict]:
+    records = []
+    for index in range(count):
+        subject = f"{prefix}{index:04d}"
+        attribute = ATTRIBUTES[index % len(ATTRIBUTES)]
+        value = rng.choice(VALUES)
+        fields = {"subject": subject, "attribute": attribute, "value": value}
+        records.append(
+            {
+                "id": f"{prefix.lower()}-{index:04d}",
+                "memory": [
+                    {
+                        "role": "user",
+                        "content": rng.choice(MEMORY_USER_TEMPLATES).format(**fields),
+                    },
+                    {
+                        "role": "assistant",
+                        "content": rng.choice(MEMORY_ASSISTANT_TEMPLATES).format(**fields),
+                    },
+                ],
+                "query": [
+                    {
+                        "role": "user",
+                        "content": rng.choice(QUERY_USER_TEMPLATES).format(**fields),
+                    },
+                {"role": "assistant", "content": value},
+                ],
+                "subject": subject,
+                "attribute": attribute,
+                "answer": value,
+            }
+        )
+    return records
+
+
+def write_jsonl(path: Path, records: list[dict]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(
+        "".join(json.dumps(record, ensure_ascii=False) + "\n" for record in records),
+        encoding="utf-8",
+    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--output-dir", default="dynamic_memory_lab/data")
+    parser.add_argument("--train-count", type=int, default=128)
+    parser.add_argument("--eval-count", type=int, default=64)
+    parser.add_argument("--seed", type=int, default=20260903)
+    args = parser.parse_args()
+    if args.train_count < 1 or args.eval_count < 1:
+        raise ValueError("train-count and eval-count must be positive")
+
+    train = make_records(args.train_count, prefix="训练实体", rng=random.Random(args.seed))
+    evaluation = make_records(args.eval_count, prefix="测试实体", rng=random.Random(args.seed + 1))
+    output_dir = Path(args.output_dir)
+    write_jsonl(output_dir / "benchmark_train.jsonl", train)
+    write_jsonl(output_dir / "benchmark_eval.jsonl", evaluation)
+    (output_dir / "benchmark_manifest.json").write_text(
+        json.dumps(
+            {
+                "seed": args.seed,
+                "train_count": len(train),
+                "eval_count": len(evaluation),
+                "task": "random subject-to-code associative recall",
+                "train_subject_prefix": "训练实体",
+                "eval_subject_prefix": "测试实体",
+                "answer_is_not_derived_from_subject": True,
+            },
+            ensure_ascii=False,
+            indent=2,
+        ),
+        encoding="utf-8",
+    )
+    print(f"train={len(train)} path={output_dir / 'benchmark_train.jsonl'}")
+    print(f"eval={len(evaluation)} path={output_dir / 'benchmark_eval.jsonl'}")
+
+
+if __name__ == "__main__":
+    main()