Add Natural Memory architecture and tooling
This commit is contained in:
WpyQwq
@@ -0,0 +1,312 @@
|
||||
# Natural Memory v2：面向长期交互的分层神经记忆与稀疏地址路由
|
||||
|
||||
## 摘要
|
||||
|
||||
大语言模型的长期记忆通常有两种实现：把历史对话原样放回上下文，或者把历史写入外部检索系统。前者的计算和显存成本随历史长度增长，后者虽然实用，却把记忆写入、读取和纠错放在模型之外，模型本身并不知道记忆为何存在、是否可信以及何时应该停止检索。
|
||||
|
||||
本文提出 Natural Memory v2，一种附着在冻结 Qwen3.5-4B 主干上的分层记忆架构。它不试图让 Memory Slot 复现完整 KV Cache，而是把长期、昂贵、适合结构化复用的历史信息编码成紧凑地址和版本化证据记录；当前窗口继续由热 KV 负责精确顺序和局部连贯性。V2 的查询路径固定为“粗索引—候选页—精确重排—Top-K 读取”，因此查询复杂度不随全部记忆槽做全量注意力。系统还加入了自动写入策略、置信度隔离、冲突版本、纠错、撤回、多跳关联和嵌入式持久化。
|
||||
|
||||
在本地控制变量实验中，20,000 条合成记录被组织为 626 个页面，平均只进入 178.13 个候选页（28.45%）；记录 Recall@K 与页 Recall@K 均为 100%，多跳和导出恢复测试均通过。通用路由器的候选 route accuracy 为 99.32%，Qwen hidden-state 路由器在 103 条 held-out 事实上的 route accuracy 为 91.26%。在 120 个固定的通用、数学、推理、语言、知识、逻辑和上下文用例上，Natural Memory v2 与原版 Qwen3.5-4B 的总分都为 0.85833，回归差值为 0。真实模型重启实验中，模型没有收到历史聊天记录，仍通过嵌入式 memory shard 找到目标事实并正确生成；未知事实经过阈值校准后不再进入内部证据前缀。
|
||||
|
||||
这些结果证明的是架构闭环和工程可行性，而不是已经解决百万级自然语言记忆。本文同时报告多跳路由不足、训练数据规模有限、NVMe 级分页注意力尚未完成等限制，并给出下一阶段的训练与系统路线。
|
||||
|
||||
**关键词：** 大语言模型、长期记忆、Memory Slot、KV Cache、稀疏路由、分页索引、持续学习、事实纠错
|
||||
|
||||
## 1. 问题定义
|
||||
|
||||
设当前对话的 token 序列为 (x_{1:t})，模型的热 KV 为 (K_t,V_t)。如果所有历史都保留在注意力上下文中，序列长度 (t) 变大时，注意力计算和 KV 显存都会持续增加。另一种做法是把历史写入外部数据库，再由固定程序检索结果并拼接进 prompt。后一种方法把检索变成了应用层逻辑：模型通常无法区分未经确认的推断、已经被纠正的旧值和当前有效值。
|
||||
|
||||
Natural Memory 的目标不是让一个固定大小的张量保存所有原始历史，而是让模型拥有一种长期运行的神经记忆接口：
|
||||
|
||||
1. 重要的个人事实、项目约束和事件关系可以自动进入记忆；
|
||||
2. 查询只访问少量有地址的记忆单元；
|
||||
3. 新旧事实具有可解释的版本关系；
|
||||
4. 证据不足时模型可以拒绝读取，而不是把相似内容强行当答案；
|
||||
5. 记忆状态可以在模型重启后恢复，而不需要重放聊天记录。
|
||||
|
||||
这里的“神经记忆”指记忆读取和写入策略是模型架构的一部分，并不意味着所有存储介质都必须是 GPU 上的可训练参数。长期记忆可以是模型拥有的状态快照，读取器、路由器和写入控制器则直接参与模型运行。运行时优先将热点记录提升到 VRAM，同时保留显存安全余量；无法安全提升的记录继续驻留在系统 RAM。
|
||||
|
||||
## 2. 设计原则
|
||||
|
||||
### 2.1 Memory Slot 不模拟 KV
|
||||
|
||||
KV Cache 适合保存近期 token 的精确顺序；Memory Slot 适合保存跨会话仍有价值的事实、摘要和关联。两者承担不同的任务：
|
||||
|
||||
```text
|
||||
热 KV = 当前工作内存，强调顺序和细节
|
||||
Memory Slot = 长期神经 RAM，强调地址、复用和生命周期
|
||||
冷页面 = 更大容量的历史存储，按需加载
|
||||
```
|
||||
|
||||
如果要求 Memory Slot 无损等价于数百万 token 的原始 KV，那么它最终仍然需要保存近似相同的信息量，只是换了数据结构。V2 选择有损但可审计的语义压缩：长期证据优先，当前上下文精确。
|
||||
|
||||
### 2.2 任何大规模读取都必须有界
|
||||
|
||||
不允许当前 token 对全部槽位直接做注意力。V2 采用两阶段路由：
|
||||
|
||||
1. 用 LSH 粗索引从页面签名中取候选页面；
|
||||
2. 对候选页面做页级和记录级精确重排；
|
||||
3. 只将 Top-K 记录的 token 证据送回 Qwen。
|
||||
|
||||
页面签名同时包含页面中心和记录地址。这样可以避免混合主题页面的中心向量把某一条稀有但相关的记录“平均掉”。当前实现还会探测查询签名的 Hamming 距离 1 和 2 的桶，并始终保留少量热页作为安全候选。
|
||||
|
||||
### 2.3 错误必须进入状态机
|
||||
|
||||
长期记忆最危险的不是暂时漏召回，而是一次错误写入之后长期污染回答。因此一条记录除了文本和向量外，还带有：
|
||||
|
||||
- 来源和证据；
|
||||
- 置信度和重要性；
|
||||
- 创建时间、访问次数和版本；
|
||||
- 实体—属性—值冲突键；
|
||||
- `active`、`superseded`、`retracted`、`quarantined` 状态；
|
||||
- `supersedes` 和 `related_ids` 关系。
|
||||
|
||||
重复写入是幂等的；同一实体和属性的新值会 supersede 旧值；不可信写入进入 quarantine，不参与普通读取；显式纠错生成新版本；撤回记录仍保留审计信息但不再被读取。
|
||||
|
||||
## 3. 架构
|
||||
|
||||
### 3.1 路由器
|
||||
|
||||
给定查询 hidden state (q) 和候选记忆键 (k_i)，路由器先把两者投影到 (d=128) 的紧凑地址空间：
|
||||
|
||||
[
|
||||
hat q = \operatorname{norm}(W_q q), \qquad
|
||||
hat k_i = \operatorname{norm}(W_k k_i).
|
||||
]
|
||||
|
||||
V2 使用 8 个路由头。每个头计算局部相似度，头门控对各头加权；随后将查询、候选和差向量送入 pair scorer。路由器同时输出：
|
||||
|
||||
- 候选记录分数；
|
||||
- 是否需要记忆的二分类 logit；
|
||||
- 多跳步数预测；
|
||||
- 每个路由头的诊断分数。
|
||||
|
||||
存储侧只保留紧凑地址，而不保存每条记录的 2560 维 Qwen hidden state。查询时只把当前 query 和已经由粗索引筛出的候选键放到路由器所在设备。
|
||||
|
||||
### 3.2 分页记忆库
|
||||
|
||||
一个页面包含有限数量的记录、页面中心、摘要中心、重要性和冷热等级。默认页面容量为 32 条记录，最大页面数为 32768，懒分配容量为：
|
||||
|
||||
[
|
||||
32768 \times 32 = 1,048,576
|
||||
]
|
||||
|
||||
这只是地址空间上限，不代表启动时分配一百万条记录。写入侧只在开放页的有界窗口中选择目标页，避免随着页面数增长而扫描整个写入空间。达到硬上限后系统抛出明确的容量错误，要求先 consolidation 或提高容量，而不是静默超过限制。
|
||||
|
||||
### 3.3 多跳检索
|
||||
|
||||
一次查询先得到候选页面和记录。被选记录可以带关联记录 ID，下一跳只在关联记录所在的页面中继续搜索。每跳都记录访问轨迹、去重集合和停止原因。系统在达到 Top-K、最大 hop、无新关联或没有新证据时停止。
|
||||
|
||||
### 3.4 Qwen 集成
|
||||
|
||||
V2 作为 Qwen3.5-4B 的附加模块接入：
|
||||
|
||||
- Qwen 主干冻结；
|
||||
- 原有 memory controller、自然语言写入策略和热文本 bank 保留兼容性；
|
||||
- V2 router 注册为模型模块并可单独训练；
|
||||
- 生成前由当前查询 hidden state 触发 V2 读取；
|
||||
- 选中记录的 token 序列在模型内部形成证据前缀；
|
||||
- 生成本身不更新持久记忆，写入发生在当前用户回合结束前；
|
||||
- memory state、路由器参数、页面元数据和记录 token 可写入 safetensors memory shard。
|
||||
|
||||
该设计使模型重启时只需要重新加载模型包和内嵌状态，不需要把历史聊天重新放入 prompt，也不需要依赖一个固定的外部数据库读取程序。
|
||||
|
||||
## 4. 训练方法
|
||||
|
||||
### 4.1 路由训练目标
|
||||
|
||||
训练损失由三部分组成：
|
||||
|
||||
[
|
||||
\mathcal L = \mathcal L_{candidate}
|
||||
+ \lambda_n \mathcal L_{need}
|
||||
+ \lambda_h \mathcal L_{hop}.
|
||||
]
|
||||
|
||||
其中候选损失使用包含 hard negatives 的交叉熵；need loss 判断问题是否需要记忆；hop loss 预测继续关联检索所需的步数。负样本不仅包括随机记录，还包括相同实体、相近属性、共享词汇但答案不同的记录。无记忆问题作为独立类别参与训练，使路由器有机会学会 abstain。
|
||||
|
||||
### 4.2 两个训练阶段
|
||||
|
||||
第一阶段使用共享潜在因子训练通用路由器，验证分页、地址投影和损失函数是否稳定。第二阶段使用真实 Qwen3.5-4B 的 hidden state 编码实体—属性—值事实，并在不同实体、不同属性和 hard negative 上训练 Qwen 专用路由器。
|
||||
|
||||
第二阶段目前使用 512 条本地合成事实，409 条用于训练、103 条用于 held-out。这样做主要是工程启动数据，不足以代表真实用户对话的分布。正式版本需要加入：
|
||||
|
||||
- 同一事实的多种问法；
|
||||
- 省略主语、代词和跨语言表达；
|
||||
- 时间先后和旧值修正；
|
||||
- 多事实组合与多跳路径；
|
||||
- 没有记录时的拒答；
|
||||
- 误导性相似事实和记忆污染样本。
|
||||
|
||||
## 5. 实验设置
|
||||
|
||||
### 5.1 硬件与加载
|
||||
|
||||
实验使用 RTX 5070，报告的总显存为 11.94 GiB；Qwen 使用 4-bit NF4 加载。主干权重保持冻结。结果来自本地工程测试，不是公开排行榜结果。
|
||||
|
||||
### 5.2 核心组件测试
|
||||
|
||||
单元测试共 18 项，覆盖：
|
||||
|
||||
- 路由器输入输出形状与紧凑地址；
|
||||
- 页面粗索引候选边界；
|
||||
- 版本冲突、纠错和 slot 替换；
|
||||
- quarantine、批准和撤回；
|
||||
- 多跳关联；
|
||||
- safetensors 前的导出/恢复语义；
|
||||
- KV 预算和页容量上限。
|
||||
|
||||
结果为 18/18 通过。
|
||||
|
||||
### 5.3 分页检索实验
|
||||
|
||||
在 20,000 条合成记录、626 页、Top-K 页面和记录限制下，结果如下：
|
||||
|
||||
| 指标 | 结果 |
|
||||
|---|---:|
|
||||
| 粗候选页平均数 | 178.13 |
|
||||
| 粗候选页最大数 | 429 |
|
||||
| 粗候选页占全部页面 | 28.45% |
|
||||
| 记录 Recall@K | 100% |
|
||||
| 页面 Recall@K | 100% |
|
||||
| 多跳成功 | 100% |
|
||||
| 多跳步数 | 2 |
|
||||
| 导出恢复后召回 | 100% |
|
||||
| 冲突版本与纠错 | 通过 |
|
||||
| quarantine 隔离与批准 | 通过 |
|
||||
| 撤回隔离 | 通过 |
|
||||
| 重复写入幂等 | 通过 |
|
||||
|
||||
页面容量配置为 32768 页 × 32 条记录，地址空间为 1,048,576 条记录。除 20,000 条分页检索实验外，项目还完成了 1,000,000 条轻量记录的 durable page store 压力测试：31,250 页实际落盘，重启后记录总数保持 1,000,000，64 个常驻页对应约 160 条常驻记录，目标记录通过精确粗桶召回。该结果证明的是存储、冷热分层和地址路径，不声称已经完成一百万条完整自然语言长文本的 Qwen 端到端质量验证。
|
||||
|
||||
通用路由器在 40 个评估批次上的结果为：route accuracy 99.32%，need-memory precision、recall、specificity 均为 100%，hop accuracy 为 86.25%。这些数字来自合成分布，不能直接推断真实对话泛化。
|
||||
|
||||
### 5.4 Qwen hidden-state 路由器
|
||||
|
||||
Qwen 专用路由器的 held-out 结果：
|
||||
|
||||
| 指标 | 结果 |
|
||||
|---|---:|
|
||||
| held-out 事实 | 103 |
|
||||
| route accuracy | 91.26% |
|
||||
| need precision | 100% |
|
||||
| need recall | 100% |
|
||||
| need specificity | 100% |
|
||||
| hop accuracy | 36.70% |
|
||||
|
||||
多跳控制器明显弱于候选记录路由。这意味着当前系统的可靠性主要来自显式关联关系、页面约束和记录状态机，不能把 hop prediction 当作唯一的正确性保证。
|
||||
|
||||
### 5.5 与原版 Qwen3.5-4B 的回归
|
||||
|
||||
综合测试使用同一组 120 个固定用例，包含通用能力、数学、推理、语言、知识、逻辑和 512—8192 token 的上下文定位。基线和 V2 均使用贪心解码与 4-bit NF4：
|
||||
|
||||
| 指标 | 原版 Qwen3.5-4B | Natural Memory v2 |
|
||||
|---|---:|---:|
|
||||
| 总分 | 0.85833 | 0.85833 |
|
||||
| 总分差值 | - | 0 |
|
||||
| 各分类差值 | - | 全部 0 |
|
||||
| 自动写入 precision | - | 100% |
|
||||
| 自动写入 recall | - | 100% |
|
||||
| 自动写入 specificity | - | 100% |
|
||||
| 无历史重启恢复 | - | 通过 |
|
||||
| 清理后停止召回 | - | 通过 |
|
||||
|
||||
这说明在当前测试集合上，接入 V2 没有造成可测的通用能力退化；它不等于在所有任务、所有长度和所有语言上都没有退化。
|
||||
|
||||
### 5.6 真实重启实验
|
||||
|
||||
实验先给模型一条普通自然语言事实：
|
||||
|
||||
> 我正在开发一个长期项目，项目内部代号是 NM-V2-RESTART，使用中文。
|
||||
|
||||
没有使用 `/remember`。自动写入成功后，系统把状态保存进 memory safetensors shard，释放第一个模型，再加载第二个模型。第二个模型只收到新问题，没有聊天记录。结果如下：
|
||||
|
||||
- router 找到 `page_00000001` 和目标记录；
|
||||
- 内部证据前缀长度为 36 token；
|
||||
- 生成结果为 `NM-V2-RESTART`；
|
||||
- 记录在报告中显示为 `evidence_found`；
|
||||
- 测试结束后清理操作将页面数和记录数恢复为 0。
|
||||
|
||||
在未知事实边界测试中，写入“我的长期项目代号是 ZX-77”后询问“我的血型是什么”，阈值校准前会出现低分无关召回；阈值设为 0.65 后，路由结果为 `below_read_threshold`，内部前缀长度为 0。这个修复体现了一个重要原则：回答碰巧说“不知道”不能代替读取器本身拒绝无关证据。
|
||||
|
||||
## 6. 长上下文与容量边界
|
||||
|
||||
直接 KV 压力测试在当前机器上约 8192 token 可以运行，16384 和 32768 token 会 OOM。这个结果与 V2 的作用并不矛盾：V2 的目标是把长期信息移出热 KV，而不是自动让原始长文本注意力变成低成本。
|
||||
|
||||
当前还不能声称支持 200M—300M 原始 token 上下文。百万级轻量记录的 durable page store 已经实现；同时，Qwen 生成路径已经接入 CPU-backed `DynamicCache(offloading=True)`，并加入了模型内的旧前缀分块归档和热窗口裁剪。真实自动路径把 371 token 压缩为 32 token，写入 22 条上下文记录后完成生成。要达到原始长上下文目标，仍然需要：
|
||||
|
||||
1. GPU 热页、RAM 温页、磁盘冷页之间更细粒度的统一 page manager，以及面向 NVMe 的分页调度；
|
||||
2. 原文页与摘要页的可逆压缩；
|
||||
3. 更高吞吐的 CPU/NVMe KV offload 与分页注意力；
|
||||
4. 128K 热窗口到百万级历史的课程训练；
|
||||
5. 大规模多跳、冲突、污染和纠错数据；
|
||||
6. 召回失败时的保守拒答和校准评测。
|
||||
|
||||
V2 已经把地址空间、页路由、版本状态机、持久化接口和 CPU KV offload 接在一起；当前证据仍然是轻量记录压力测试和短自动压缩验证，不是百万 token 原始上下文的端到端质量证明。
|
||||
|
||||
## 7. 讨论
|
||||
|
||||
### 7.1 与外部 RAG 的差异
|
||||
|
||||
V2 与普通 RAG 的主要区别不在于“是否存在向量”，而在于记忆生命周期由模型运行时直接控制。路由器学习问题是否需要记忆，写入策略决定什么值得保留，记录状态机维护冲突和撤回，证据前缀由模型内部读取路径生成。本次交付的默认实现不启用磁盘分页：完整 V2 记录随第三个 safetensors memory shard 载入进程内存，只有有限的热点地址和 token payload 进入 VRAM cache。未来更大规模部署可以增加冷页后端，但那是容量扩展，不是当前正确性路径。
|
||||
|
||||
### 7.2 与扩大 KV 的差异
|
||||
|
||||
扩大 KV 对最近历史的细节保持最好，但长期存储成本高，而且每个查询都容易被无关上下文拖慢。Memory Slot 主动丢弃顺序细节，只保存事实和结构化证据，因此更适合个人偏好、项目决策、联系人属性和长期计划。两者应该组合使用，而不是互相替代。
|
||||
|
||||
### 7.3 为什么训练比扩大张量更难
|
||||
|
||||
增加 slot 数量只改变了容量，不会教会模型如何寻址。真正困难的是建立稳定的写入地址、区分相似事件、关联多个记录、修正旧版本，以及在没有答案时停下来。当前 Qwen 路由器的 hop accuracy 已经显示，候选选择和多跳控制是两个不同的学习问题，不能用单一的相似度损失解决。
|
||||
|
||||
## 8. 可复现性
|
||||
|
||||
项目目录为 `W:\Flash\model\dynamic_memory_lab`。核心命令：
|
||||
|
||||
```powershell
|
||||
Set-Location W:\Flash\model
|
||||
|
||||
# 单元测试
|
||||
& C:\Users\Administrator\miniconda3\envs\LLM\python.exe `
|
||||
-m unittest discover -s dynamic_memory_lab\tests -v
|
||||
|
||||
# V2 存储与路由评测
|
||||
& C:\Users\Administrator\miniconda3\envs\LLM\python.exe `
|
||||
-m dynamic_memory_lab.benchmark_memory_v2
|
||||
|
||||
# Qwen3.5-4B 综合回归
|
||||
& C:\Users\Administrator\miniconda3\envs\LLM\python.exe `
|
||||
-m dynamic_memory_lab.benchmark_natural_memory_v1 `
|
||||
--base-model W:\Flash\model `
|
||||
--memory-model W:\Flash\model\dynamic_memory_lab\qwen3_5_4b_natural_memory_v2 `
|
||||
--output W:\Flash\model\dynamic_memory_lab\natural_memory_v2_full_benchmark.json
|
||||
|
||||
# 真实重启测试
|
||||
& C:\Users\Administrator\miniconda3\envs\LLM\python.exe `
|
||||
-m dynamic_memory_lab.test_natural_memory_v2_restart
|
||||
|
||||
# KV offload 与自动长上下文压缩
|
||||
& C:\Users\Administrator\miniconda3\envs\LLM\python.exe `
|
||||
-m dynamic_memory_lab.benchmark_kv_offload
|
||||
```
|
||||
|
||||
关键结果文件：
|
||||
|
||||
- `natural_memory_v2_benchmark.json`：分页、路由、完整性和恢复；
|
||||
- `natural_memory_v2_full_benchmark.json`：Qwen 基线回归；
|
||||
- `natural_memory_v2_restart_test.json`：无历史重启；
|
||||
- `qwen_router_v2_training.json`：Qwen 专用路由器训练与 held-out 指标。
|
||||
|
||||
## 9. 局限与未来工作
|
||||
|
||||
本文的实现仍然是研究原型，距离可公开部署还有几项关键工作：
|
||||
|
||||
1. 用脱敏真实对话替代小规模合成事实，并按用户、项目和时间做严格数据隔离；
|
||||
2. 训练多跳控制器和不确定性校准，而不是只优化 top-1 路由；
|
||||
3. 引入事实来源签名、用户确认策略和可撤销日志；
|
||||
4. 将页对象替换为压缩列式存储，降低百万记录的 Python 对象开销；
|
||||
5. 完成 GPU/RAM/NVMe 三级冷热迁移和 NVMe page cache；
|
||||
6. 实现摘要页与原文页之间的按需恢复；
|
||||
7. 建立跨会话、跨语言、长时间连续运行的污染测试；
|
||||
8. 在 128K 热 KV 和百万级历史上进行端到端吞吐、延迟、质量和故障恢复评测。
|
||||
|
||||
Natural Memory v2 最重要的成果不是一个更大的 slot 张量，而是把“记忆是什么、如何寻址、何时读取、怎样纠错、何时拒绝”放进了一个可测试的模型架构中。后续版本的主要任务，是让这个架构在更大规模和更真实的分布上保持同样的克制与可解释性。
|
||||
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,871 @@
|
||||
# Dynamic Memory Lab
|
||||
|
||||
一个在现成 Qwen3.5-4B 语言模型上进行“架构手术”的实验工程。
|
||||
|
||||
> Natural Memory v2 已完成分页地址路由、粗索引→候选页→精排→Top-K、多跳、冲突版本、quarantine、撤回、嵌入式重启恢复和完整回归评测。请先阅读 [Natural Memory v2 工程说明](README_NATURAL_MEMORY_V2.md) 与 [技术论文](Natural_Memory_v2_Paper.md)。
|
||||
|
||||
本项目不重新训练整个语言模型，而是在 Transformer 的若干层之间插入一个可读、可写、可持续更新的动态记忆模块，研究下面这个问题：
|
||||
|
||||
> 能不能让一个已经训练好的 LLM，在不把所有历史对话反复塞回上下文的情况下，持续积累和使用信息？
|
||||
|
||||
答案是：工程中已经实现了一个生产导向的自然语言记忆核心，并完成了多事实写入、冲突更新、无历史跨重启读取、二次重启读取、未知事实拒答和 reset token 验收。它仍需要在真实业务数据上继续做压力测试，合成数据指标不能直接等同于所有场景的生产承诺。
|
||||
|
||||
## 结论先说
|
||||
|
||||
当前模型已经具备：
|
||||
|
||||
- 在一次进程运行期间维护一组独立于 token 上下文的 memory state；
|
||||
- 对 memory 做 attention 读取；
|
||||
- 根据当前输入生成写入内容，并更新 memory；
|
||||
- 将 memory 读取结果注入 Qwen3.5-4B 的中间层；
|
||||
- 只训练记忆模块和少量门控参数，而冻结原始 Qwen 主干；
|
||||
- 对比原版 Qwen3.5-4B 与动态记忆版的 loss、困惑度、准确率、生成质量、速度和显存。
|
||||
|
||||
当前版本仍有边界：
|
||||
|
||||
- 还没有证明对任意自然语言、多事实和长时间运行都可靠；
|
||||
- 还没有完成生产级多用户隔离、隐私策略和容量管理；
|
||||
- 自由生成在合成留出集上仍会偶尔答错值，所以不能把实验指标当作生产承诺。
|
||||
|
||||
因此，当前版本更准确的定位是：
|
||||
|
||||
> 一个包含可部署自然语言记忆核心、持久化 checkpoint 和验收基准的 LLM 原生动态记忆工程；生产接入前仍需补齐业务侧鉴权、加密、并发和运维策略。
|
||||
|
||||
## 项目结构
|
||||
|
||||
dynamic_memory_lab/
|
||||
├─ qwen_integration.py # 动态记忆模块和 Qwen 模型适配器
|
||||
├─ train_qwen_memory.py # 记忆模块训练脚本
|
||||
├─ chat_qwen_memory.py # 多轮对话实验脚本
|
||||
├─ benchmark_qwen.py # 原版/动态版跑分脚本
|
||||
├─ verify_persistent_memory.py # 重启后仅用 memory_state 验证回答
|
||||
├─ make_native_memory_data.py # 生成写入/遗忘/未知/冲突数据
|
||||
├─ train_native_memory.py # 训练原生写入、遗忘和读出
|
||||
├─ evaluate_native_memory.py # 留出集控制器与自由生成评估
|
||||
├─ verify_native_checkpoint.py # checkpoint 重启和 reset token 验证
|
||||
├─ train_natural_retriever.py # 训练自然语言记忆查询检索器
|
||||
├─ train_auto_memory_policy.py # 训练自动记忆重要性策略
|
||||
├─ build_production_memory_dataset.py # 真实对话导入、规范化和防泄漏切分
|
||||
├─ train_production_memory_policy.py # 从规范化对话训练策略候选
|
||||
├─ stress_test_natural_memory.py # 低显存连续运行/长上下文压力测试
|
||||
├─ natural_memory_service.py # localhost 管理与聊天 API
|
||||
├─ natural_memory_app.py # 统一 chat/serve/train/stress 入口
|
||||
├─ benchmark_natural_language_memory.py # 多事实/更新/重启/拒答验收
|
||||
├─ stream_chat_qwen_memory.py # 流式聊天与随时重启测试入口
|
||||
├─ make_benchmark_data.py # 生成简单的记忆型 benchmark 数据
|
||||
├─ requirements.txt # Python 依赖
|
||||
├─ memory.pt # 训练后的记忆模块参数，若已生成
|
||||
├─ memory_config.json # 记忆模块配置，若已生成
|
||||
├─ surgery.pt # 层手术相关参数，若已生成
|
||||
└─ README.md
|
||||
|
||||
## 运行环境
|
||||
|
||||
推荐使用已有的 Conda 环境 LLM。
|
||||
|
||||
conda activate LLM
|
||||
cd W:\Flash\model\dynamic_memory_lab
|
||||
pip install -r requirements.txt
|
||||
|
||||
如果模型路径不在默认位置，可以通过参数指定。当前工程主要面向本地 Hugging Face 格式的 Qwen3.5-4B 模型。
|
||||
|
||||
## 总体架构
|
||||
|
||||
原版 Qwen 的推理过程大致是：
|
||||
|
||||
输入 token
|
||||
│
|
||||
▼
|
||||
Embedding
|
||||
│
|
||||
▼
|
||||
Transformer Layer 0
|
||||
│
|
||||
▼
|
||||
Transformer Layer 1
|
||||
│
|
||||
...
|
||||
│
|
||||
▼
|
||||
Transformer Layer N
|
||||
│
|
||||
▼
|
||||
LM Head
|
||||
│
|
||||
▼
|
||||
下一个 token
|
||||
|
||||
动态记忆版会在若干个 Transformer 层上插入 Memory Adapter：
|
||||
|
||||
输入 token
|
||||
│
|
||||
▼
|
||||
Qwen Transformer Layer
|
||||
│
|
||||
├──────────────► Memory Read
|
||||
│ ▲
|
||||
│ │
|
||||
│ Runtime Memory State
|
||||
│ M = [B, S, D]
|
||||
│ │
|
||||
│ ▼
|
||||
├──────────────► Memory Delta
|
||||
│
|
||||
▼
|
||||
Memory Layer Adapter
|
||||
│
|
||||
▼
|
||||
后续 Qwen Transformer Layers
|
||||
│
|
||||
├──────────────► Memory Write
|
||||
│ │
|
||||
│ ▼
|
||||
│ M_new = Update(M, input)
|
||||
│
|
||||
▼
|
||||
LM Head
|
||||
|
||||
其中：
|
||||
|
||||
- B 是 batch size；
|
||||
- S 是 memory slot 数量，默认 16；
|
||||
- D 是 memory embedding 维度，默认 512；
|
||||
- M 不是 token 序列，而是模型外部维护的一组连续向量；
|
||||
- memory_state 可以在不同调用之间传递，因此它能够脱离上一轮的文本上下文。
|
||||
|
||||
## 核心概念：参数、上下文和运行时记忆
|
||||
|
||||
理解本项目最重要的是区分三种东西。
|
||||
|
||||
### 1. 模型参数
|
||||
|
||||
模型参数是 Qwen 的权重以及动态记忆模块的权重，例如：
|
||||
|
||||
Wq, Wk, Wv, Wread, Wwrite
|
||||
|
||||
它们决定模型“如何读写记忆”，通常在训练阶段更新，在推理阶段保持不变。
|
||||
|
||||
### 2. 当前上下文
|
||||
|
||||
当前上下文是这次请求中送入模型的 token，例如：
|
||||
|
||||
用户：我叫小明。
|
||||
助手：你好，小明。
|
||||
|
||||
上下文是临时的。上下文窗口结束以后，模型本身不会自动保存这些 token。
|
||||
|
||||
### 3. 运行时记忆状态
|
||||
|
||||
运行时记忆状态是：
|
||||
|
||||
M = [batch_size, memory_slots, memory_dim]
|
||||
|
||||
默认情况下：
|
||||
|
||||
M = [B, 16, 512]
|
||||
|
||||
它是模型运行时的一块连续状态。只要下一次调用仍然传入同一个 memory_state，模型就能继续使用之前写入的内容。
|
||||
|
||||
注意：
|
||||
|
||||
> memory.pt 保存的是“记忆模块的训练参数”，不是某个用户的聊天记忆。
|
||||
|
||||
用户记忆应该单独保存为某种 runtime state，例如：
|
||||
|
||||
user_001_memory.pt
|
||||
user_002_memory.pt
|
||||
|
||||
或者保存到数据库、对象存储、向量数据库中。
|
||||
|
||||
## Memory Read：模型如何读取记忆
|
||||
|
||||
设某一层产生的隐藏状态为：
|
||||
|
||||
h ∈ R^D_hidden
|
||||
|
||||
记忆矩阵为：
|
||||
|
||||
M = [m₁, m₂, ..., mₛ] ∈ R^(S×D_memory)
|
||||
|
||||
首先把当前隐藏状态投影成 query：
|
||||
|
||||
q = W_q h
|
||||
|
||||
把每个 memory slot 投影成 key 和 value：
|
||||
|
||||
k_i = W_k m_i
|
||||
v_i = W_v m_i
|
||||
|
||||
然后计算当前输入与每个 slot 的匹配程度：
|
||||
|
||||
score_i = q · k_i / sqrt(D_memory)
|
||||
|
||||
经过 softmax 得到读取权重：
|
||||
|
||||
α_i = softmax(score_i)
|
||||
|
||||
最后将各个 slot 的 value 加权求和：
|
||||
|
||||
r = Σ_i α_i v_i
|
||||
|
||||
r 就是当前输入从动态记忆中检索出来的内容。
|
||||
|
||||
为了避免记忆模块一开始就破坏 Qwen，代码还使用了一个 read gate：
|
||||
|
||||
g = sigmoid(W_gate h)
|
||||
|
||||
最终的记忆增量大致为：
|
||||
|
||||
Δh = read_scale × g × W_read(r)
|
||||
|
||||
然后再注入当前层：
|
||||
|
||||
h_new = h + Δh
|
||||
|
||||
这和给 Transformer 增加一个小型、可训练的外部知识通道类似。
|
||||
|
||||
## Memory Write：模型如何写入记忆
|
||||
|
||||
读取解决的是“从记忆里找什么”，写入解决的是“把当前输入存什么”。
|
||||
|
||||
当前实现默认使用输入序列最后一个 token 的隐藏状态作为摘要：
|
||||
|
||||
s = h_last
|
||||
|
||||
也支持固定 token 偏移位置作为写入摘要。
|
||||
|
||||
接着通过写入投影生成候选内容：
|
||||
|
||||
p = W_write(s)
|
||||
|
||||
再根据当前输入生成写入地址和写入强度：
|
||||
|
||||
a = softmax(W_addr(s))
|
||||
z = sigmoid(W_strength(s))
|
||||
|
||||
其中：
|
||||
|
||||
- a_i 表示第 i 个 slot 被写入的比例；
|
||||
- z 表示本次写入总体有多强；
|
||||
- p 是候选写入向量。
|
||||
|
||||
对于每个 slot，更新形式近似为：
|
||||
|
||||
m_i_new = (1 - z × a_i) × m_i
|
||||
+ (z × a_i) × p
|
||||
|
||||
这是一种可微分的软写入。它不会使用不可导的“直接选中某个 slot”操作，所以可以通过反向传播学习：
|
||||
|
||||
- 什么输入值得写入；
|
||||
- 应该写入哪些 slot；
|
||||
- 写入幅度应该多大；
|
||||
- 如何从隐藏状态中压缩信息。
|
||||
|
||||
代码还支持 broadcast_write，让候选内容广播写入所有 slot。这个模式更适合做架构实验，但可能降低 slot 的分工能力。
|
||||
|
||||
## Qwen 接入方式
|
||||
|
||||
qwen_integration.py 会加载原版 Qwen，并替换指定层为带记忆能力的适配层。
|
||||
|
||||
默认会选择若干中间层；也可以通过 layer_indices 手动指定层号。工程中同时保留了一个轻量的自定义线性注意力记忆层，便于做对照实验。
|
||||
|
||||
动态模型默认冻结 Qwen 主干：
|
||||
|
||||
Qwen 原始参数：冻结
|
||||
Memory Read/Write：训练
|
||||
层融合 gate：训练
|
||||
|
||||
这样做的好处是：
|
||||
|
||||
- 显存和训练成本更低；
|
||||
- 不容易破坏原模型能力；
|
||||
- 更容易判断提升来自记忆机制还是来自主干重新学习；
|
||||
- 适合在单卡环境中快速迭代架构。
|
||||
|
||||
## 三种层融合模式
|
||||
|
||||
动态记忆读出后，需要决定如何注入 Qwen 的隐藏状态。当前支持三种模式。
|
||||
|
||||
### residual
|
||||
|
||||
h_new = h + memory_delta
|
||||
|
||||
这是默认模式。它保留原始隐藏状态，并把记忆当作额外残差通道。
|
||||
|
||||
特点：
|
||||
|
||||
- 最稳定；
|
||||
- 对原模型干扰小；
|
||||
- 适合第一版训练和 benchmark。
|
||||
|
||||
### blend
|
||||
|
||||
h_new = h_residual + (1 - α) × original_token_mixer + α × memory_delta
|
||||
|
||||
其中 α 是可训练的融合系数。
|
||||
|
||||
特点：
|
||||
|
||||
- 模型可以学习记忆通道应该占多大比例；
|
||||
- 适合研究“原始表示”和“记忆表示”的权衡；
|
||||
- 如果初始化或训练不稳定，可能导致原模型信息被过早削弱。
|
||||
|
||||
### replace
|
||||
|
||||
h_new = memory_delta
|
||||
|
||||
完全使用记忆分支输出。
|
||||
|
||||
特点：
|
||||
|
||||
- 适合验证记忆分支的独立表达能力；
|
||||
- 风险最高；
|
||||
- 通常不建议作为默认生产方案。
|
||||
|
||||
## direct_logit_scale：直接影响输出概率
|
||||
|
||||
除了修改 Transformer 中间层，代码还支持把最后一次 memory readout 经过投影后直接加到 logits：
|
||||
|
||||
logits_new = logits_qwen + scale × projection(memory_readout)
|
||||
|
||||
这个选项可以直接研究：
|
||||
|
||||
- memory 是否能记住某些目标答案；
|
||||
- memory 是否能把目标 token 的概率推高；
|
||||
- 记忆模块对最终预测的直接影响。
|
||||
|
||||
但它也更容易过拟合简单 benchmark，因此应该同时观察泛化测试和正常生成质量。
|
||||
|
||||
## raw token pointer：精确 token 记忆实验
|
||||
|
||||
项目还提供一个有意“开后门”的架构消融：memory pass 不仅写入连续向量，还可以保存某个 token 在 Qwen 输出投影矩阵中的行；查询生成的第一个 token 会读取这行向量。
|
||||
|
||||
这个实验用于回答一个非常具体的问题：
|
||||
|
||||
> 如果记忆里已经存在目标 token，当前 Qwen 接口能不能把它可靠地送进自由生成？
|
||||
|
||||
示例配置：
|
||||
|
||||
--write-token-offset 4 --broadcast-write \
|
||||
--raw-token-write --raw-logit-scale 30
|
||||
|
||||
注意：`write-token-offset` 按“从有效序列末尾数起”计算，`1` 是最后一个 token。当前 benchmark 的答案位于 `H / 。 / <|im_end|> / 换行` 中的倒数第 4 个位置，所以使用 `4`。`raw-token-write` 是 pointer ablation，不是通用自然语言记忆方案；它直接保存 token id 对应的输出投影行，不能把它的 100% 结果等同于普通 learned memory 的能力。
|
||||
|
||||
raw pointer 只作用于第一个生成 token，后续 token 回到 Qwen 原本的生成分布，避免把同一个答案 token 重复写满整段输出。
|
||||
|
||||
## 训练原理
|
||||
|
||||
当前训练脚本采用“两阶段记忆训练”。
|
||||
|
||||
### 阶段一：写入阶段
|
||||
|
||||
输入一段 memory text：
|
||||
|
||||
Memory: user=alice; favorite_color=blue
|
||||
|
||||
此阶段关闭 memory read，打开 memory write：
|
||||
|
||||
outputs = model(
|
||||
memory_text,
|
||||
memory_state=memory_state,
|
||||
read_memory=False,
|
||||
update_memory=True,
|
||||
return_memory=True,
|
||||
)
|
||||
memory_state = outputs.memory_state
|
||||
|
||||
目标是让模型把关键信息写入 memory state。
|
||||
|
||||
### 阶段二：查询阶段
|
||||
|
||||
再输入查询：
|
||||
|
||||
Question: What is alice's favorite color?
|
||||
|
||||
此阶段使用刚刚更新后的 memory state，并开启 memory read：
|
||||
|
||||
outputs = model(
|
||||
query_text,
|
||||
memory_state=memory_state,
|
||||
read_memory=True,
|
||||
update_memory=False,
|
||||
labels=labels,
|
||||
)
|
||||
|
||||
通过语言模型 loss 训练记忆模块，让查询阶段能够根据 memory state 输出正确答案。
|
||||
|
||||
训练时通常只更新：
|
||||
|
||||
Memory Read parameters
|
||||
Memory Write parameters
|
||||
Memory Layer Adapter parameters
|
||||
Blend/Gate parameters
|
||||
|
||||
而 Qwen 主干保持冻结。
|
||||
|
||||
## 数据格式
|
||||
|
||||
训练和 benchmark 数据使用 JSONL。当前脚本要求每行包含 `memory` 和 `query` 两个消息列表；最后一个 assistant 消息是监督目标：
|
||||
|
||||
{"memory":[{"role":"user","content":"记住对象 A 的代号是 H。"},{"role":"assistant","content":"好的，已记住。"}],"query":[{"role":"user","content":"对象 A 的代号是什么？"},{"role":"assistant","content":"H"}]}
|
||||
|
||||
额外的 `id`、`subject`、`attribute`、`answer` 字段只用于 benchmark 统计。训练和评估数据应避免把答案直接重复到 query 的 user 内容中。
|
||||
|
||||
数据设计时要注意：
|
||||
|
||||
1. 写入文本中出现的信息，应该在查询文本中尽量不重复；
|
||||
2. 如果查询中直接包含答案，模型可能只是在复制上下文，而不是读取 memory；
|
||||
3. 训练集和评估集中的实体、属性、表述方式应尽量分离；
|
||||
4. 要加入冲突样本，测试新记忆是否能覆盖旧记忆；
|
||||
5. 要加入多条事实，测试不同 slot 是否会互相污染；
|
||||
6. 要加入无关信息，测试模型能否避免把所有内容都写进去。
|
||||
|
||||
make_benchmark_data.py 可以生成一个简单的单字符映射任务，用来快速检查“写入—读取—回答”链路是否工作。它适合做冒烟测试，不足以证明模型拥有通用长期记忆。
|
||||
|
||||
## 常用命令
|
||||
|
||||
以下命令均在工程目录执行。
|
||||
|
||||
### 训练动态记忆模块
|
||||
|
||||
conda activate LLM
|
||||
cd W:\Flash\model
|
||||
python -m dynamic_memory_lab.train_qwen_memory --model-path "W:\Flash\model" --data dynamic_memory_lab\data\benchmark_train.jsonl --output-dir dynamic_memory_lab\qwen_memory_adapter --steps 100 --batch-size 1 --lr 1e-4 --max-length 128
|
||||
|
||||
如果要复现实验中的精确 token pointer：
|
||||
|
||||
python -m dynamic_memory_lab.train_qwen_memory --model-path "W:\Flash\model" --data dynamic_memory_lab\data\benchmark_train.jsonl --output-dir dynamic_memory_lab\qwen_memory_adapter_pointer --steps 100 --surgery-mode blend --blend-init 0.1 --write-token-offset 4 --broadcast-write --raw-token-write --raw-logit-scale 30
|
||||
|
||||
如果实际模型目录不同，请替换 --model-path。
|
||||
|
||||
### 训练原生记忆控制器
|
||||
|
||||
原生模式会训练：摘要池化、写入决策、slot 地址、候选值读出，以及依赖旧记忆的遗忘门。训练数据中显式包含闲聊噪声、未知查询和冲突覆盖样本：
|
||||
|
||||
python -m dynamic_memory_lab.make_native_memory_data --output-dir dynamic_memory_lab\data\native_memory --train-count 512 --eval-count 128 --seed 20260904
|
||||
python -m dynamic_memory_lab.train_native_memory --model-path "W:\Flash\model" --data dynamic_memory_lab\data\native_memory\train.jsonl --output-dir dynamic_memory_lab\qwen_memory_adapter_native_v3 --steps 2500 --lr 8e-5 --max-length 192 --save-every 250 --direct-logit-scale 12 --write-loss-weight 0.5 --forget-loss-weight 0.75 --forget-positive-weight 6 --value-loss-weight 1
|
||||
|
||||
评估：
|
||||
|
||||
python -m dynamic_memory_lab.evaluate_native_memory --model-path "W:\Flash\model" --adapter-dir dynamic_memory_lab\qwen_memory_adapter_native_v3 --data dynamic_memory_lab\data\native_memory\eval.jsonl --max-length 192 --max-new-tokens 8
|
||||
|
||||
### 生成 benchmark 数据
|
||||
|
||||
python -m dynamic_memory_lab.make_benchmark_data --output-dir dynamic_memory_lab\data
|
||||
|
||||
### 对比原版和动态版
|
||||
|
||||
python -m dynamic_memory_lab.benchmark_qwen --model-path "W:\Flash\model" --data dynamic_memory_lab\data\benchmark_eval.jsonl --adapter dynamic_memory_lab\qwen_memory_adapter_pointer --output dynamic_memory_lab\benchmark_qwen_pointer_eval.json --max-length 128 --max-new-tokens 4 --repeats 1 --warmup 0
|
||||
|
||||
benchmark 通常会报告：
|
||||
|
||||
- validation loss；
|
||||
- perplexity；
|
||||
- token accuracy；
|
||||
- first target token accuracy；
|
||||
- exact sequence accuracy；
|
||||
- generation quality；
|
||||
- tokens per second；
|
||||
- 峰值显存。
|
||||
|
||||
最终报告必须同时关注效果和代价。一个模型如果只是在极小任务上准确率更高，却明显降低通用生成能力或推理速度，不能简单视为架构成功。
|
||||
|
||||
### 本机已验证结果
|
||||
|
||||
在 `data/benchmark_train.jsonl` 的 128 条训练样本和 `data/benchmark_eval.jsonl` 的 64 条 held-out 随机映射上，答案字符不是由实体名称推导出来的。当前已保存的结果文件是：
|
||||
|
||||
| 版本 | PPL | token accuracy | 首目标 token | 自由生成 prefix | 速度 |
|
||||
|---|---:|---:|---:|---:|---:|
|
||||
| 原版 Qwen3.5-4B | 210.20 | 42.71% | 0% | 0% | 19.45 tok/s |
|
||||
| 普通 learned blend | 4.50 | 66.67% | 0% | 0% | 12.99 tok/s |
|
||||
| blend + raw token pointer | 1.11 | 96.88% | 100% | 100% | 14.56 tok/s |
|
||||
| native learned memory v3 | 1.24 | 91.67% | 75% | 75% | 13.33 tok/s |
|
||||
|
||||
对应文件分别是 `benchmark_qwen_native_v3_eval.json`、`benchmark_qwen_pointer_eval.json` 中的 baseline/dynamic 记录，以及 `benchmark_qwen_full_eval.json` 中保存的普通 blend 结果。native v3 行来自同一套 64 条 benchmark；pointer 行是精确 token 消融实验，不能替代通用 learned memory 的结论。不同运行的速度会受显存缓存和系统状态影响。
|
||||
|
||||
原生控制器 v3 在 128 条完全不同用户编号的留出集上得到：写入准确率 100%，遗忘准确率 99.58%，replacement 遗忘准确率 96.30%，查询 token 准确率 92.12%，自由生成事实召回 81.31%，未知查询安全拒答 100%。这些结果来自 `qwen_memory_adapter_native_v3/native_eval_report.json`；其中自由生成仍有少量错误值，不能宣称已经达到可靠生产级记忆。
|
||||
|
||||
### 多轮对话实验
|
||||
|
||||
python -m dynamic_memory_lab.chat_qwen_memory --model-path "W:\Flash\model" --adapter dynamic_memory_lab\qwen_memory_adapter
|
||||
|
||||
要让用户记忆跨进程保存，指定一个用户专属文件：
|
||||
|
||||
python -m dynamic_memory_lab.chat_qwen_memory --model-path "W:\Flash\model" --adapter dynamic_memory_lab\qwen_memory_adapter --memory-state dynamic_memory_lab\data\users\alice.pt
|
||||
|
||||
首次运行可以输入：
|
||||
|
||||
/remember 我叫小明，喜欢蓝色
|
||||
|
||||
退出后再次运行同一条命令，直接询问个人事实即可；启动时只会加载 `alice.pt`，不会自动加载上一轮聊天文本。`/reset` 会把该用户的 state 重置为零并保存，`/save` 可以手动保存。
|
||||
|
||||
需要区分“记忆持久化”和“记忆能力”：`qwen_memory_adapter_natural_auto_v3` 加上持久化 checkpoint 已验证自然语言写入、冲突覆盖、两次无历史重启读取、未知事实拒答和 reset；真实业务上线仍需要按业务数据继续扩充评测。
|
||||
|
||||
不启用 `--persistent-memory` 时，当前 chat 脚本中的 memory 默认是进程内状态；关闭脚本后，这块状态会消失。
|
||||
|
||||
使用原生模式并把用户记忆写进适配器 checkpoint：
|
||||
|
||||
python -m dynamic_memory_lab.chat_qwen_memory --model-path "W:\Flash\model" --adapter dynamic_memory_lab\qwen_memory_adapter_native_v3 --persistent-memory --persist-in-adapter
|
||||
|
||||
原生 adapter 默认使用 `<|fim_prefix|>` 作为 reset token；也可以通过 `--reset-token` 或 `--reset-token-id` 指定其他 tokenizer token。向模型发送该 token 会在模型内部清零 memory，不需要外部清理函数参与推理。
|
||||
|
||||
原生 chat 的普通用户消息会先经过自动记忆策略头判断是否值得长期保存；写入阶段不读取旧记忆，避免把回忆内容或问题句误写回去。随后生成阶段只读，不把模型自己的回答再次写回记忆。`/remember` 仍可用于强制写入。
|
||||
|
||||
## 当前保存机制的边界
|
||||
|
||||
当前代码已经提供真正的 runtime state 持久化接口：
|
||||
|
||||
model.save_runtime_memory("user_memories/alice.pt")
|
||||
model.load_runtime_memory("user_memories/alice.pt")
|
||||
|
||||
保存文件包含：
|
||||
|
||||
- `memory_state`：连续动态记忆张量；
|
||||
- `raw_memory`：如果启用了 token pointer，则保存对应的辅助状态；
|
||||
- hidden size、memory shape、层配置等兼容性信息。
|
||||
|
||||
它不包含模型权重，也不包含历史聊天文本。因此重启后的调用可以只传入加载后的 state 和新的 query。
|
||||
|
||||
如果希望把“记忆模块参数 + 当前用户 memory state”放到同一个紧凑适配器包中：
|
||||
|
||||
model.save_persistent_memory_checkpoint("dynamic_memory_lab/qwen_memory_adapter_native_v3_persistent")
|
||||
|
||||
该包中的 `persistent_memory.pt` 会在 `load_memory_adapter()` 时自动加载，新的模型实例不需要再显式传入 `memory_state`。验证命令：
|
||||
|
||||
python -m dynamic_memory_lab.verify_native_checkpoint --model-path "W:\Flash\model" --adapter dynamic_memory_lab\qwen_memory_adapter_native_v3 --data dynamic_memory_lab\data\native_memory\eval.jsonl --output-adapter dynamic_memory_lab\qwen_memory_adapter_native_v3_persistent
|
||||
|
||||
可以用下面的脚本自动验证完整流程：
|
||||
|
||||
python -m dynamic_memory_lab.verify_persistent_memory --model-path "W:\Flash\model" --adapter dynamic_memory_lab\qwen_memory_adapter_pointer --data dynamic_memory_lab\data\benchmark_eval.jsonl --memory-state dynamic_memory_lab\data\user_memory_demo.pt
|
||||
|
||||
该脚本会执行：写入 memory → 保存 state → 销毁模型 → 重新加载模型 → 只用 state 生成回答。
|
||||
|
||||
项目会保存动态记忆模块的训练参数，例如：
|
||||
|
||||
memory.pt
|
||||
memory_config.json
|
||||
surgery.pt
|
||||
|
||||
这些文件描述的是“模型如何使用记忆”；`persistent_memory.pt` 则是打包进适配器的“某个用户已经记住了什么”。
|
||||
|
||||
可以把它们类比成：
|
||||
|
||||
memory.pt = 记忆系统的大脑结构和读写规则
|
||||
user_memory.pt = 某个用户实际写入的内容
|
||||
|
||||
当前实现的 runtime state 通常通过以下方式流动：
|
||||
|
||||
outputs = model(
|
||||
input_ids=input_ids,
|
||||
memory_state=memory_state,
|
||||
return_memory=True,
|
||||
)
|
||||
memory_state = outputs.memory_state
|
||||
|
||||
只要把 memory_state 保存下来，之后重新加载并传回模型，就可以恢复对应记忆。
|
||||
|
||||
## 能不能实现跨对话、无上下文、原生记忆？
|
||||
|
||||
需要先把这个问题拆开。
|
||||
|
||||
### 现在能做到的版本：跨调用、无历史文本
|
||||
|
||||
可以。
|
||||
|
||||
同一个进程里：
|
||||
|
||||
第 1 次调用：输入事实，更新 memory_state
|
||||
第 2 次调用：只输入问题，传入 memory_state
|
||||
第 3 次调用：继续传入更新后的 memory_state
|
||||
|
||||
第 2 次调用不必把第 1 次调用的完整聊天记录重新放进 prompt。模型可以从 memory state 中读取信息。
|
||||
|
||||
### 加一个持久化层后：跨程序、跨会话
|
||||
|
||||
也可以实现，但它不是模型单独完成的，而是：
|
||||
|
||||
用户 ID
|
||||
│
|
||||
▼
|
||||
加载该用户的 runtime memory state
|
||||
│
|
||||
▼
|
||||
调用动态记忆模型
|
||||
│
|
||||
▼
|
||||
保存更新后的 runtime memory state
|
||||
|
||||
最小实现可以是：
|
||||
|
||||
from pathlib import Path
|
||||
import torch
|
||||
|
||||
def load_user_memory(user_id, model):
|
||||
path = Path("user_memories") / f"{user_id}.pt"
|
||||
if path.exists():
|
||||
return torch.load(path, map_location="cpu")
|
||||
return model.initial_memory(batch_size=1)
|
||||
|
||||
def save_user_memory(user_id, memory_state):
|
||||
path = Path("user_memories") / f"{user_id}.pt"
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
torch.save(memory_state.detach().cpu(), path)
|
||||
|
||||
实际生产环境还需要：
|
||||
|
||||
- 用户隔离；
|
||||
- 加密；
|
||||
- 并发写入保护；
|
||||
- 版本迁移；
|
||||
- memory state 校验；
|
||||
- 删除和导出接口；
|
||||
- 过期时间或记忆衰减；
|
||||
- 防止恶意 prompt 写入长期记忆。
|
||||
|
||||
### 严格意义上的“完全无上下文”不可能凭空发生
|
||||
|
||||
如果“无上下文”指：
|
||||
|
||||
不提供历史文本
|
||||
不提供 memory state
|
||||
不提供数据库
|
||||
不提供任何外部信息
|
||||
|
||||
但又希望模型知道上一场对话发生了什么，那么这在信息上是不可能的。
|
||||
|
||||
信息必须存在某个地方：
|
||||
|
||||
上下文窗口
|
||||
模型参数
|
||||
外部数据库
|
||||
向量索引
|
||||
神经 memory state
|
||||
|
||||
本项目选择的是“神经 memory state”这条路线。它可以让历史信息不以原始文本形式出现，但不能让信息在没有任何载体的情况下存在。
|
||||
|
||||
### 为什么仍然需要生产化治理
|
||||
|
||||
当前自然语言记忆核心已经不只是连续状态：它有精确文本槽、冻结 Qwen 检索键、训练过的查询检索器、写入门控、冲突替换、最老槽淘汰、checkpoint 校验和 reset。但上线前仍要针对真实业务补充：
|
||||
|
||||
- 用户身份隔离、加密、并发写入和版本迁移；
|
||||
- 真实语言中的事实抽取、撤回/删除语义和多事实长时间压力测试；
|
||||
- 记忆容量策略、审计日志、导出与合规删除；
|
||||
- 不同语言、不同 tokenizer、不同 batch 和服务重启方式的回归测试。
|
||||
|
||||
更准确的说法是：
|
||||
|
||||
> 模型已经拥有内部的“是否写入、如何检索、何时拒绝读取、如何覆盖和淘汰”的自然语言记忆路径；用户 checkpoint 仍必须存在于某种持久化介质中，这是信息保存的必要条件，而不是外部代码替模型执行读取。
|
||||
|
||||
## 为什么不直接每轮修改 Qwen 主权重
|
||||
|
||||
每轮对话都直接更新主模型权重，理论上可以把信息写进参数，但会产生明显问题：
|
||||
|
||||
- 很容易灾难性遗忘；
|
||||
- 不同用户之间会互相污染；
|
||||
- 每次写入都需要保存或更新大模型权重；
|
||||
- 难以撤销某条记忆；
|
||||
- 难以处理隐私和权限；
|
||||
- 推理延迟和存储成本都很高。
|
||||
|
||||
外部 runtime memory 的好处是：
|
||||
|
||||
主模型参数 = 稳定的通用能力
|
||||
用户 memory = 可修改、可删除、可隔离的个体状态
|
||||
|
||||
这是更接近实际产品需求的拆分。
|
||||
|
||||
## 当前自然语言记忆核心
|
||||
|
||||
这是当前推荐的生产导向路径，不再把个人事实压缩成一个容易丢失多 token 值的连续向量：
|
||||
|
||||
用户输入
|
||||
│
|
||||
├─ Automatic memory policy：判断普通消息是否值得长期记忆
|
||||
├─ Native write controller：提供写入表示和槽位地址
|
||||
├─ Frozen Qwen key encoder：生成稳定检索键
|
||||
├─ Learned retriever：查询与记忆槽匹配
|
||||
└─ Internal text bank：保存合法 memory prefix、原始 fact key、年龄和有效位
|
||||
│
|
||||
▼
|
||||
相关记忆才被模型内部拼成 prefix
|
||||
│
|
||||
▼
|
||||
原版 Qwen 生成回答
|
||||
|
||||
具体保证：
|
||||
|
||||
- 每个槽保存可直接参与 Qwen 对话的内部 memory prefix，并单独保存原始 fact key，因此“红富士苹果”“蓝鲸-47”等多 token 值不会被单个向量强行压缩；
|
||||
- 普通消息由训练过的自动记忆策略决定是否写入；写入阶段关闭读取，生成阶段只读，避免把模型自己的回答或已召回事实再次写回记忆；
|
||||
- 同一属性的更新由训练过的语义检索器确认，不同事实优先使用空槽，容量满时按最老槽淘汰；
|
||||
- 读取由模型内部的 learned retriever 触发。无关问题低于阈值时，不注入记忆 prefix，降低个人事实幻觉；
|
||||
- `persistent_memory.pt` 直接包含连续状态、文本槽、检索键、年龄和计数器。重启时只加载这个用户 checkpoint，不需要聊天历史，也不需要外部“记忆读取器”代码；
|
||||
- `<|fim_prefix|>`（或显式配置的 reset token）由模型包装器在生成入口内识别并清空全部记忆槽。
|
||||
|
||||
训练检索器：
|
||||
|
||||
conda activate LLM
|
||||
cd W:\Flash\model
|
||||
python -m dynamic_memory_lab.train_natural_retriever --model-path "W:\Flash\model" --base-adapter dynamic_memory_lab\qwen_memory_adapter_native_v3 --output-adapter dynamic_memory_lab\qwen_memory_adapter_natural_controller_v13 --steps 3000 --pair-count 3200 --batch-size 32 --lr 2e-4
|
||||
|
||||
训练自动记忆策略头：
|
||||
|
||||
python -m dynamic_memory_lab.train_auto_memory_policy --model-path "W:\Flash\model" --base-adapter dynamic_memory_lab\qwen_memory_adapter_natural_controller_v13 --output-adapter dynamic_memory_lab\qwen_memory_adapter_natural_auto_v13 --steps 2600 --example-count 1600 --batch-size 32 --lr 2e-4 --threshold 0.35 --text-memory-threshold 0.30
|
||||
|
||||
启动自然语言记忆聊天：
|
||||
|
||||
python -m dynamic_memory_lab.chat_qwen_memory --model-path "W:\Flash\model" --adapter dynamic_memory_lab\qwen_memory_adapter_natural_auto_v13 --persistent-memory --persist-in-adapter
|
||||
|
||||
如果要测试流式输出和“随时重启”，使用：
|
||||
|
||||
python -m dynamic_memory_lab.stream_chat_qwen_memory --model-path "W:\Flash\model" --adapter dynamic_memory_lab\qwen_memory_adapter_natural_auto_v13 --memory-state dynamic_memory_lab\data\users\stream_user_memory.pt
|
||||
|
||||
该脚本每轮只发送当前用户消息，不发送历史；写入或清空操作会在生成前原子保存。生成过程中按 `Ctrl+C` 退出后，再次执行同一命令即可从最近一次保存的 memory state 继续。普通消息会自动保存高价值个人事实，`/remember <事实>` 用于强制写入，`/reset` 清空全部记忆。
|
||||
|
||||
运行完整验收：
|
||||
|
||||
python -m dynamic_memory_lab.benchmark_natural_language_memory --model-path "W:\Flash\model" --adapter dynamic_memory_lab\qwen_memory_adapter_natural_auto_v13 --output-adapter dynamic_memory_lab\qwen_memory_adapter_natural_production_auto_v13 --report dynamic_memory_lab\benchmark_natural_language_memory_auto_v13_final.json --max-new-tokens 48 --text-memory-threshold 0.30
|
||||
|
||||
当前验收报告为 `benchmark_natural_language_memory_auto_v13_final.json`：三条事实写入后保留两个独立有效槽，工作地点从 R7 更新为 K9；两条已知事实在两次模型重启后均能回答；未知血型不注入 memory prefix 并拒答；reset token 后连续状态和文本槽都清零。报告中的 `production_gate_pass` 为 `true`。
|
||||
|
||||
自动记忆策略的训练报告为 `qwen_memory_adapter_natural_auto_v13/auto_policy_training.json`：2600 步、每类 1600 条样本，评估准确率 99.22%，负样本特异度 100%，误写率 0%。这些是合成数据结果，仍需用真实用户语言继续扩充压力测试。
|
||||
|
||||
### 将记忆模块和当前用户状态合并进 safetensors
|
||||
|
||||
如果需要把 v13 的记忆模块参数、检索器、自动写入策略和某个用户的当前状态写进模型权重，可以生成一个新的合并目录。下面的命令是兼容的合并方式；推荐在新包上使用额外第三分片模式。
|
||||
|
||||
```powershell
|
||||
python -m dynamic_memory_lab.merge_memory_weights `
|
||||
--base-model "W:\Flash\model" `
|
||||
--adapter "W:\Flash\model\dynamic_memory_lab\qwen_memory_adapter_natural_auto_v13" `
|
||||
--memory-state "W:\Flash\model\dynamic_memory_lab\data\users\stream_user_memory.pt" `
|
||||
--output "W:\Flash\model\dynamic_memory_lab\qwen3_5_4b_memory_merged_v13"
|
||||
```
|
||||
|
||||
额外第三分片模式使用同样的参数，并增加 `--add-shard`；目标文件名应为 `model.safetensors-00003-of-00003.safetensors`：
|
||||
|
||||
```powershell
|
||||
python -m dynamic_memory_lab.merge_memory_weights `
|
||||
--base-model "<base-model>" `
|
||||
--adapter "<adapter>" `
|
||||
--memory-state "<memory-state>" `
|
||||
--output "<merged-package>" `
|
||||
--add-shard "<merged-package>\\model.safetensors-00003-of-00003.safetensors"
|
||||
```
|
||||
|
||||
兼容旧合并方式生成的目录会把 `dynamic_memory.*` 张量放入第二个 safetensors 分片；原始 Qwen 两个分片不会被覆盖。自定义动态记忆模型可以直接从该目录加载：
|
||||
|
||||
```powershell
|
||||
python -m dynamic_memory_lab.stream_chat_qwen_memory `
|
||||
--model-path "W:\Flash\model\dynamic_memory_lab\qwen3_5_4b_memory_merged_v13" `
|
||||
--memory-state "W:\Flash\model\dynamic_memory_lab\data\users\merged_runtime.pt"
|
||||
```
|
||||
|
||||
上面的旧模式仍可使用 `merged_runtime.pt` 保存后续变化；真正执行动态记忆仍需要本项目的模型架构代码。
|
||||
|
||||
当前推荐的第三分片启动方式是：
|
||||
|
||||
```powershell
|
||||
python -m dynamic_memory_lab.stream_chat_qwen_memory `
|
||||
--model-path "<merged-package>"
|
||||
```
|
||||
|
||||
不传 `--adapter` 和 `--memory-state` 时，模型会从 `model.safetensors-00003-of-00003.safetensors` 读取内置 memory state；之后的自动写入、`/save` 和 `/reset` 也只原子重写第三分片，不再创建或依赖 `merged_runtime.pt`。旧的“把动态张量追加到第二分片”方式仍可通过不使用 `--add-shard` 保留，但每次保存会重写较大的第二分片。
|
||||
|
||||
检索器 v13 额外加入了短事实、无标点问法、未见过的值、登记/称呼表达、编程工具表达和“你是谁”硬负样本；训练与线上都使用 plain tokens，并按长度分组避免 Qwen 补零造成表示漂移。训练报告的标准留出、短事实留出和 raw sigmoid 严格阈值召回均为 100%；当前 Wpy 真实身份/拒读压力测试为 25/25。后两项是更接近线上行为的指标，仍应随真实用户分布持续回归。
|
||||
|
||||
通用能力回归报告为 `comprehensive_benchmark_natural_auto_v13_final.json`：54 个任务上原版 Qwen3.5-4B 与动态记忆版均为 96.30%，整体差值 0，回归门禁通过。
|
||||
|
||||
需要明确：模型内部已经拥有“何时读、读什么、如何拒绝无关记忆”的路径，但跨机器或跨服务保存用户 checkpoint 仍然需要持久化介质；这是信息存在的物理要求，不等于推理时依赖固定外部读取代码。
|
||||
|
||||
## 当前架构的局限
|
||||
|
||||
当前版本主要用于研究“记忆模块能否工作”，还不是最终架构。
|
||||
|
||||
### 记忆写入过于粗粒度
|
||||
|
||||
默认使用最后一个 token 的 hidden state 作为摘要。复杂文本可能包含多条事实，仅靠最后一个 token 很容易丢失信息。
|
||||
|
||||
后续可以加入：
|
||||
|
||||
- sentence-level summarizer；
|
||||
- entity/value extractor；
|
||||
- 多 token pooling；
|
||||
- 特殊 memory token；
|
||||
- 独立的写入路由器。
|
||||
|
||||
### slot 语义还不稳定
|
||||
|
||||
固定数量的连续 slot 不一定自动形成清晰分工。后续可以研究：
|
||||
|
||||
- slot type；
|
||||
- key-value memory；
|
||||
- 稀疏路由；
|
||||
- memory usage regularization；
|
||||
- slot 专家化；
|
||||
- 多层级 memory。
|
||||
|
||||
### teacher forcing 与自由生成存在差距
|
||||
|
||||
在训练中，模型通常知道正确答案的标签；但在真正聊天时，需要连续生成多个 token。应该分别测试：
|
||||
|
||||
- 单 token 读取；
|
||||
- 多 token 事实回答；
|
||||
- 多轮连续记忆；
|
||||
- 错误回答后的恢复；
|
||||
- 新记忆覆盖旧记忆；
|
||||
- memory 容量接近上限时的退化。
|
||||
|
||||
### 还缺少完整的记忆治理
|
||||
|
||||
生产级系统至少需要三层：
|
||||
|
||||
模型内部 Memory
|
||||
│
|
||||
├─ 记忆写入策略
|
||||
├─ 记忆读取策略
|
||||
└─ 连续向量状态
|
||||
|
||||
记忆管理器
|
||||
│
|
||||
├─ 事实抽取
|
||||
├─ 去重
|
||||
├─ 冲突处理
|
||||
├─ 时间衰减
|
||||
└─ 重要性评分
|
||||
|
||||
持久化服务
|
||||
│
|
||||
├─ 用户隔离
|
||||
├─ 加密
|
||||
├─ 版本管理
|
||||
└─ 删除/导出
|
||||
|
||||
## 评测时应该回答的关键问题
|
||||
|
||||
不要只问“准确率有没有变高”，还要问：
|
||||
|
||||
1. 模型是否真的使用了 memory，而不是从 query 中猜答案？
|
||||
2. 新记忆能否覆盖旧记忆？
|
||||
3. 无关内容是否会污染 memory？
|
||||
4. memory 容量增加后，效果是否持续提升？
|
||||
5. 程序重启后能否恢复？
|
||||
6. 不同用户之间是否完全隔离？
|
||||
7. memory read/write 是否降低原模型的通用能力？
|
||||
8. 记忆状态是否可以解释、导出和删除？
|
||||
9. 推理速度和显存成本是多少？
|
||||
10. 长时间运行后是否出现状态漂移或数值爆炸？
|
||||
|
||||
## 下一步建议
|
||||
|
||||
如果目标是开发一个真正有新意的架构，建议按下面顺序推进：
|
||||
|
||||
1. 扩展到多事实、多属性和长序列连续写入；
|
||||
2. 增加时间衰减、记忆容量压力和可解释 slot 诊断；
|
||||
3. 做多用户隔离、并发读写和异常恢复测试；
|
||||
4. 再尝试更激进的结构，例如 fast weights、可写 KV cache、分层记忆和稀疏路由。
|
||||
|
||||
## 免责声明
|
||||
|
||||
本项目是研究和实验性质的代码。它不保证训练出的 adapter 在所有任务上提升，也不保证当前连续 memory state 能可靠保存所有自然语言事实。
|
||||
|
||||
如果使用第三方模型权重或数据集，请遵守对应的模型许可证、数据许可证和隐私要求。
|
||||
@@ -0,0 +1,388 @@
|
||||
# Natural Memory v2
|
||||
|
||||
Natural Memory v2 是接在本地 Qwen3.5-4B 上的一层分层、可寻址、可纠错记忆系统。它的目标不是把一百万个 slot 当作一张巨大的 KV Cache，而是把历史中最适合长期保存的部分压缩成有地址的记忆记录；当前对话仍由 Qwen 的热 KV 负责连续理解。
|
||||
|
||||
## 先看结论
|
||||
|
||||
当前版本已经完成以下闭环：
|
||||
|
||||
```text
|
||||
普通用户输入
|
||||
-> 自动写入策略
|
||||
-> 置信度 / 重要性 / 来源审计
|
||||
-> 紧凑地址投影
|
||||
-> 分页存储
|
||||
-> LSH 粗索引
|
||||
-> 候选页重排
|
||||
-> 候选记录精排
|
||||
-> Top-K 证据注入 Qwen
|
||||
-> 版本冲突 / 隔离 / 撤回 / 重启恢复
|
||||
```
|
||||
|
||||
最终运行包位于：
|
||||
|
||||
`W:\Flash\model\dynamic_memory_lab\qwen3_5_4b_natural_memory_v2`
|
||||
|
||||
它不需要额外的 `memory_state.pt` 才能恢复已嵌入的记忆；记忆快照、路由器权重、V2 页面元数据和上下文片段都写在该包的第三个 memory safetensors 切片中。当前包的 manifest 指向实际使用的 runtime memory shard；旧 runtime shard 不会被自动覆盖，以避免 Windows 进程仍持有文件句柄时破坏现有包。当前交付配置固定使用 embedded weight-shard 模式：启动时完整载入进程内存，路由命中的热点记录才进入有界 VRAM cache，不使用 SQLite 或磁盘分页。
|
||||
|
||||
## 核心设计
|
||||
|
||||
### KV 与 Memory Slot 的分工
|
||||
|
||||
```text
|
||||
最近 32K token（最多可配置到 128K） -> GPU 热 KV：保留精确顺序和局部连贯性
|
||||
长期个人事实 / 项目决策 / 纠错版本 -> Memory Slot：保存压缩、可寻址的证据
|
||||
当前问题 -> 稀疏路由：只读取少量相关记录
|
||||
```
|
||||
|
||||
Memory Slot 不模拟完整 KV。它只接管 KV 中最昂贵、最适合长期保存、最容易重复利用的部分。这样做的代价是不能承诺逐 token 无损复现数百万 token 的原始上下文；换来的收益是长期历史不再全部占用 GPU 注意力。
|
||||
|
||||
### 路由路径
|
||||
|
||||
任何查询都必须经过有界路径：
|
||||
|
||||
```text
|
||||
查询 hidden state
|
||||
|
|
||||
v
|
||||
128 维紧凑地址
|
||||
|
|
||||
v
|
||||
LSH 粗索引（精确桶 + Hamming-1/2 探针 + 热页）
|
||||
|
|
||||
v
|
||||
候选页（不是全部页面）
|
||||
|
|
||||
v
|
||||
页级与记录级精确重排
|
||||
|
|
||||
v
|
||||
最多 top_k_pages / top_k_records
|
||||
```
|
||||
|
||||
当前 token 不会对 1M slot 做全量注意力。页内才会做小规模记录评分；默认是最多 4 个页、8 条记录。粗索引结果还会记录在诊断 trace 中，便于检查“是否因为候选页不足而漏召回”。
|
||||
|
||||
### 记录结构
|
||||
|
||||
每条 `MemoryRecordV2` 包含：
|
||||
|
||||
- 原始短文本和可选 token 序列；
|
||||
- 128 维紧凑地址与摘要地址；
|
||||
- `entity / attribute / value` 冲突键；
|
||||
- 时间戳、版本号、来源、证据；
|
||||
- 置信度、重要性、访问次数；
|
||||
- `active / superseded / retracted / quarantined` 状态；
|
||||
- `supersedes` 和 `related_ids`，用于版本追踪和多跳检索。
|
||||
|
||||
重复文本是幂等写入；同一实体和属性的新值会生成新版本并将旧值标记为 `superseded`。不可信写入进入 quarantine，不参与正常读取，只有显式批准后才会变成 active。撤回不会抹掉审计记录，而是把记录置为 `retracted`。
|
||||
|
||||
### 多跳读取
|
||||
|
||||
记录可以带 `related_ids`。第一跳找到一个项目、人物或事件后，路由器会沿关联记录继续查找，直到：
|
||||
|
||||
- 达到 `memory_max_hops`；
|
||||
- 达到 Top-K；
|
||||
- 没有新关联页；
|
||||
- 没有新证据。
|
||||
|
||||
每次读取都返回 `hop_trace` 和 `stop_reason`，不是只返回一段无法解释的文本。
|
||||
|
||||
### 写入安全边界
|
||||
|
||||
自动写入由原有自然语言记忆策略决定；V2 另外检查置信度和重要性。默认写入阈值为 `0.50`，读取阈值为 `0.65`。读取阈值是特意偏保守的：在实际 Qwen 测试中，未知的“我的血型是什么”最初会受到短中文问句的语义相似度干扰；提高阈值后该问题被路由层拒绝，不再把无关事实放进上下文前缀。
|
||||
|
||||
这不是“绝不出错”的证明，而是一个可检查的安全策略：
|
||||
|
||||
```text
|
||||
低置信度写入 -> quarantine
|
||||
旧值被纠正 -> 新版本 active，旧版本 superseded
|
||||
用户撤回 -> retracted，读取隔离
|
||||
证据不足 -> router_abstained / below_read_threshold
|
||||
```
|
||||
|
||||
## 训练
|
||||
|
||||
### 通用路由器训练
|
||||
|
||||
```powershell
|
||||
Set-Location W:\Flash\model
|
||||
& C:\Users\Administrator\miniconda3\envs\LLM\python.exe `
|
||||
-m dynamic_memory_lab.train_memory_router_v2
|
||||
```
|
||||
|
||||
训练目标包括：
|
||||
|
||||
1. 在 hard negatives 中选对目标记录；
|
||||
2. 判断当前问题是否需要记忆；
|
||||
3. 预测是否需要继续多跳；
|
||||
4. 让无记忆问题学会 abstain。
|
||||
|
||||
### Qwen hidden-state 路由器训练
|
||||
|
||||
最终包使用的是 Qwen3.5 hidden state 上训练的路由器：
|
||||
|
||||
```powershell
|
||||
Set-Location W:\Flash\model
|
||||
& C:\Users\Administrator\miniconda3\envs\LLM\python.exe `
|
||||
-m dynamic_memory_lab.train_qwen_router_v2 `
|
||||
--output-dir W:\Flash\model\dynamic_memory_lab\checkpoints\natural_memory_v2_qwen_router_entities
|
||||
```
|
||||
|
||||
当前训练数据是本地生成的实体—属性—值事实及 hard negatives，不是公共榜单数据集。因此训练结果可以证明工程链路有效，但不能直接等同于公开 benchmark 的泛化能力。后续正式训练应加入真实对话脱敏集、改写问句、时间冲突、跨语言表达、未知事实和长文档事件链。
|
||||
|
||||
### 构建嵌入包
|
||||
|
||||
在支持硬链接的 NTFS 目录中，可以从 v1 包构建新的完整目录：
|
||||
|
||||
```powershell
|
||||
Set-Location W:\Flash\model
|
||||
& C:\Users\Administrator\miniconda3\envs\LLM\python.exe `
|
||||
dynamic_memory_lab\build_natural_memory_v2_package.py `
|
||||
--base-package W:\Flash\model\dynamic_memory_lab\qwen3_5_4b_memory_merged_v13 `
|
||||
--output-dir W:\Flash\model\dynamic_memory_lab\qwen3_5_4b_natural_memory_v2_new `
|
||||
--router-checkpoint W:\Flash\model\dynamic_memory_lab\checkpoints\natural_memory_v2_qwen_router_entities\memory_router_v2.pt
|
||||
```
|
||||
|
||||
W: 当前是 exFAT，不能创建硬链接。构建器现在默认拒绝复制多 GB 的冻结分片，必须明确加 `--allow-copy-base` 才允许复制；这样可以避免一次构建意外耗尽磁盘空间。当前交付包使用已存在的 v2 包原地更新，未重复复制两片 Qwen 主权重。
|
||||
|
||||
## 运行流式聊天
|
||||
|
||||
```powershell
|
||||
Set-Location W:\Flash\model
|
||||
& C:\Users\Administrator\miniconda3\envs\LLM\python.exe `
|
||||
-m dynamic_memory_lab.stream_chat_qwen_memory `
|
||||
--model-path W:\Flash\model\dynamic_memory_lab\qwen3_5_4b_natural_memory_v2 `
|
||||
--max-new-tokens 128
|
||||
```
|
||||
|
||||
当前生产默认不启用 SQLite/磁盘分页。若做独立的容量研究，旧版仍保留可选的分层页库参数，但它不属于本次默认运行路径：
|
||||
|
||||
```powershell
|
||||
& C:\Users\Administrator\miniconda3\envs\LLM\python.exe `
|
||||
-m dynamic_memory_lab.stream_chat_qwen_memory `
|
||||
--model-path W:\Flash\model\dynamic_memory_lab\qwen3_5_4b_natural_memory_v2 `
|
||||
--tiered-memory-path W:\Flash\model\dynamic_memory_lab\memory_pages.sqlite `
|
||||
--memory-resident-pages 64 `
|
||||
--kv-offload
|
||||
```
|
||||
|
||||
启动后发送普通自然语言即可触发自动判断；不需要 `/remember`。常用控制命令仍保留：
|
||||
|
||||
- `/save`：把当前持久记忆写入 memory safetensors 切片；
|
||||
- `/reset`：清空持久记忆并保存；
|
||||
- `/quit`：退出。
|
||||
|
||||
模型重启时不会收到历史聊天记录。它只从嵌入式 memory shard 恢复记录、路由器和审计元数据。
|
||||
|
||||
## 单用户本地生产工作流
|
||||
|
||||
当前优先完成的 2/3/4/5 已经集中到一个入口：
|
||||
|
||||
```powershell
|
||||
Set-Location W:\Flash\model
|
||||
$py = "C:\Users\Administrator\miniconda3\envs\LLM\python.exe"
|
||||
|
||||
# 2. 规范化真实对话导出，并按 group_id 防止 train/eval 泄漏
|
||||
& $py -m dynamic_memory_lab.natural_memory_app build-dataset `
|
||||
--source W:\path\to\redacted_conversations.jsonl `
|
||||
--eval-source W:\path\to\redacted_eval.jsonl
|
||||
|
||||
# 2. 训练候选自动写入策略；不会覆盖当前生产适配器
|
||||
& $py -m dynamic_memory_lab.natural_memory_app train-policy `
|
||||
--model-path W:\Flash\model\dynamic_memory_lab\qwen3_5_4b_natural_memory_v2 `
|
||||
--base-adapter W:\Flash\model\dynamic_memory_lab\qwen_memory_adapter_natural_auto_v13 `
|
||||
--steps 240
|
||||
|
||||
# 4. 低显存连续运行与长上下文压缩压力测试
|
||||
& $py -m dynamic_memory_lab.natural_memory_app stress `
|
||||
--model-path W:\Flash\model\dynamic_memory_lab\qwen3_5_4b_natural_memory_v2
|
||||
|
||||
# 5. 启动本地 API；默认只监听 localhost
|
||||
& $py -m dynamic_memory_lab.natural_memory_app serve `
|
||||
--model-path W:\Flash\model\dynamic_memory_lab\qwen3_5_4b_natural_memory_v2 `
|
||||
--port 8765
|
||||
```
|
||||
|
||||
默认训练数据目录是 `data/production_memory`，包含 `train.jsonl`、`eval.jsonl` 和
|
||||
`manifest.json`。没有提供真实脱敏对话时，构建器会使用工程内 bootstrap 数据；这只能验证链路，不能冒充真实业务泛化结果。
|
||||
|
||||
管理接口：
|
||||
|
||||
```text
|
||||
GET /health
|
||||
GET /v1/memory?status=active&query=...
|
||||
GET /v1/memory/{record_id}
|
||||
GET /v1/memory/export
|
||||
GET /v1/memory/audit
|
||||
POST /v1/memory/{record_id} # 版本化编辑
|
||||
DELETE /v1/memory/{record_id} # 可审计撤回
|
||||
POST /v1/memory/reset # 清空并持久化
|
||||
POST /v1/chat # {message,max_new_tokens,stream}
|
||||
POST /v1/memory # 管理员/测试用显式写入
|
||||
```
|
||||
|
||||
`POST /v1/chat` 默认由模型自己的自动策略决定是否写入；`stream: true` 返回 SSE token 流。服务以单模型锁串行化请求，避免同一用户的 memory state 被并发写坏。默认自动持久化会回写 embedded memory safetensors；测试时可加 `--no-auto-persist`。
|
||||
|
||||
## 记忆管理与审计
|
||||
|
||||
编辑不是覆盖原记录，而是生成 `version + 1` 的 successor，并把旧记录标记为
|
||||
`superseded`；删除同样不物理抹除，而是标记为 `retracted`。`GET /v1/memory/audit`
|
||||
会检查页容量、页指针、冲突索引和多跳关联是否存在悬空引用。所有管理接口只暴露 JSON-safe
|
||||
元数据，不返回路由向量和模型内部 tensor。
|
||||
|
||||
## 当前受控验证结果
|
||||
|
||||
- 单元测试：19/19 通过；
|
||||
- bootstrap 数据规范化：训练 1485 条，验证 363 条；
|
||||
- 候选自动策略 smoke train：4bit、batch 2、24 steps，验证集 accuracy 96.88%、recall 90.91%、FPR 0；这不是最终生产成绩；
|
||||
- 10 轮压力测试：写入 40 条，召回 39/40（97.5%），2 次长上下文压缩，审计 healthy，0 errors；
|
||||
- 压力测试峰值约 3.2 GB allocated VRAM，使用 256 条/131072 token 的自适应热点缓存上限，并保留 2048 MB 显存安全余量；
|
||||
- 压力测试不会写回模型包，生产服务只有在启用自动持久化时才会写回。
|
||||
|
||||
## 重启验证
|
||||
|
||||
```powershell
|
||||
Set-Location W:\Flash\model
|
||||
& C:\Users\Administrator\miniconda3\envs\LLM\python.exe `
|
||||
-m dynamic_memory_lab.test_natural_memory_v2_restart `
|
||||
--model-path W:\Flash\model\dynamic_memory_lab\qwen3_5_4b_natural_memory_v2 `
|
||||
--report W:\Flash\model\dynamic_memory_lab\natural_memory_v2_restart_test.json
|
||||
```
|
||||
|
||||
测试会：
|
||||
|
||||
1. 清空选定包的持久记忆；
|
||||
2. 通过普通用户句子自动写入一条控制事实；
|
||||
3. 写入嵌入式 safetensors memory shard；
|
||||
4. 释放第一个 Qwen 模型；
|
||||
5. 重新加载模型，只输入一个新问题；
|
||||
6. 检查 V2 router decision、内部 prefix 和生成答案；
|
||||
7. 默认清理测试事实，避免污染工作包。
|
||||
|
||||
当前实际结果：自动写入成功；重启后路由器找到 `page_00000001` 和目标记录；内部前缀长度 36；生成结果精确返回 `NM-V2-RESTART`；清理后页面与记录均为 0。
|
||||
|
||||
## 评测
|
||||
|
||||
### 单元测试
|
||||
|
||||
```powershell
|
||||
Set-Location W:\Flash\model
|
||||
& C:\Users\Administrator\miniconda3\envs\LLM\python.exe `
|
||||
-m unittest discover -s dynamic_memory_lab\tests -v
|
||||
```
|
||||
|
||||
当前结果：19/19 通过，覆盖路由张量形状、压缩地址、页粗索引、版本冲突、quarantine、批准、撤回、多跳、导出恢复、KV 预算、页容量上限、批量上下文分块、记忆编辑/撤回/审计，以及分层后端重启、冷页卸载和隔离区恢复。
|
||||
|
||||
### V2 存储与路由评测
|
||||
|
||||
报告：`W:\Flash\model\dynamic_memory_lab\natural_memory_v2_benchmark.json`
|
||||
|
||||
| 指标 | 实测结果 |
|
||||
|---|---:|
|
||||
| 合成记录 | 20,000 |
|
||||
| 页面数 | 626 |
|
||||
| 粗候选页平均数 | 178.13 |
|
||||
| 粗候选页占比 | 28.45% |
|
||||
| 记录 Recall@K | 100% |
|
||||
| 页面 Recall@K | 100% |
|
||||
| 多跳成功 | 100%，2 hops |
|
||||
| 导出/恢复后召回 | 100% |
|
||||
| 冲突版本/纠错/隔离/批准/撤回/幂等 | 全部通过 |
|
||||
| 地址空间容量（32K 页 × 32） | 1,048,576 条记录 |
|
||||
| 通用路由器 route accuracy | 99.32% |
|
||||
| 通用 need-memory precision/recall/specificity | 100% / 100% / 100% |
|
||||
| 通用 hop accuracy | 86.25% |
|
||||
|
||||
上述是控制变量下的合成存储评测，证明的是分页、索引和状态机，不是 1M 条真实用户记忆已经完成验证。
|
||||
|
||||
### 分层后端百万级压力测试
|
||||
|
||||
报告：`W:\Flash\model\dynamic_memory_lab\tiered_memory_v2_1m_benchmark.json`
|
||||
|
||||
该测试实际写入 1,000,000 条轻量记录、31,250 页和 64 个常驻页。它是旧的 durable page store 容量实验，不是当前默认方案；当前默认方案要求所有记忆随第三个 safetensors 切片加载进进程内存，再用有界 VRAM cache 加速热点记录。无论哪种方案，这个数字都不等于百万条完整自然语言长文本在 Qwen 上的端到端生成质量。
|
||||
|
||||
| 指标 | 实测结果 |
|
||||
|---|---:|
|
||||
| 实际写入记录 | 1,000,000 |
|
||||
| 实际页面 | 31,250 |
|
||||
| 重启后记录总数 | 1,000,000 |
|
||||
| 重启后常驻记录 | 160 |
|
||||
| 冷页 | 31,186 |
|
||||
| 目标记录重启召回 | 通过 |
|
||||
|
||||
### Qwen 路由器验证
|
||||
|
||||
报告：`W:\Flash\model\dynamic_memory_lab\checkpoints\natural_memory_v2_qwen_router_entities\qwen_router_v2_training.json`
|
||||
|
||||
- 512 条生成事实；409 条训练，103 条 held-out；
|
||||
- route accuracy：91.26%；
|
||||
- need-memory precision/recall/specificity：100% / 100% / 100%；
|
||||
- hop accuracy：36.70%。
|
||||
|
||||
hop controller 目前明显弱于候选记录路由，因此运行时不会把它当成唯一正确性来源；关联记录和显式 `related_ids` 仍由存储层约束，后续训练应重点补多跳样本。
|
||||
|
||||
### 与原版 Qwen3.5-4B 的综合回归
|
||||
|
||||
报告：`W:\Flash\model\dynamic_memory_lab\natural_memory_v2_full_benchmark.json`
|
||||
|
||||
同一份 120 个固定用例、同一 Qwen3.5-4B 主干、同一 4-bit NF4 加载和贪心解码：
|
||||
|
||||
| 指标 | 原版 Qwen3.5-4B | Natural Memory v2 |
|
||||
|---|---:|---:|
|
||||
| 总分 | 0.85833 | 0.85833 |
|
||||
| 总分变化 | - | 0 |
|
||||
| 通用 / 数学 / 推理 / 语言 / 知识 / 逻辑 / 上下文分类 | 基线 | 各分类 delta = 0 |
|
||||
| 自动写入 precision | - | 100% |
|
||||
| 自动写入 recall | - | 100% |
|
||||
| 自动写入 specificity | - | 100% |
|
||||
| 无历史重启恢复 | - | 通过 |
|
||||
| 清理后不再召回 | - | 通过 |
|
||||
|
||||
综合评测的硬件是 RTX 5070 11.94 GiB，完整模型使用 4-bit NF4。显存峰值字段来自 PyTorch allocator，在当前 Transformers/bitsandbytes 组合下可能高于物理显存读数，不能把该字段当成独立硬件测量；最终是否能运行，应以实际 GPU OOM 和 `nvidia-smi` 为准。
|
||||
|
||||
### 长上下文边界
|
||||
|
||||
报告：`W:\Flash\model\dynamic_memory_lab\long_context_v1_stress.json`
|
||||
|
||||
现有直接 KV 路径在约 8K token 可以运行，16K 和 32K 会 OOM。现在 V2 已接入 Transformers 的 CPU-backed `DynamicCache(offloading=True)`，并在 Qwen3.5 混合线性/全注意力结构上完成真实生成验证；线性注意力的微小循环状态留在执行设备，只有昂贵的 full-attention KV 迁移到 CPU。
|
||||
|
||||
模型还提供自动热窗口压缩：当输入超过 `kv_budget_tokens`，旧前缀会按 `context_chunk_tokens` 分块写入 V2 context records，当前生成只保留最近热窗口；读取仍然走“粗索引 -> 候选页 -> 精排 -> Top-K”，不是让当前 token 对 1M slot 做注意力。真实自动路径测试将 371 token 压缩为 32 token，并写入 22 条可追溯上下文记录，随后成功生成。
|
||||
|
||||
这仍然不能宣称当前包支持 200M—300M 原始 token 上下文：CPU KV offload 不是 NVMe 分页注意力，自动压缩目前也没有完成百万 token 的 Qwen 质量训练。正确的工程目标是热 KV 保留当前窗口，超出窗口的内容经过事件切分、摘要、实体关系和可逆原文页写入 V2；查询时只加载少量相关页，再由 Qwen 做最终回答。
|
||||
|
||||
KV offload 与自动压缩验证：
|
||||
|
||||
```powershell
|
||||
Set-Location W:\Flash\model
|
||||
& C:\Users\Administrator\miniconda3\envs\LLM\python.exe `
|
||||
-m dynamic_memory_lab.benchmark_kv_offload
|
||||
```
|
||||
|
||||
## 关键源文件
|
||||
|
||||
| 文件 | 作用 |
|
||||
|---|---|
|
||||
| `memory_os_v2.py` | V2 路由器、分页存储、LSH 粗索引、精排、多跳、版本和 KV 预算 |
|
||||
| `qwen_integration.py` | 将 V2 接入 Qwen，负责内部 prefix、嵌入式 safetensors 读写和重启恢复 |
|
||||
| `stream_chat_qwen_memory.py` | 不依赖历史上下文的流式聊天与自动写入 |
|
||||
| `train_memory_router_v2.py` | 通用路由器 hard-negative 训练 |
|
||||
| `train_qwen_router_v2.py` | Qwen hidden-state 路由器训练 |
|
||||
| `tiered_memory_store_v2.py` | RAM/磁盘分层、二进制记录和可恢复 page store |
|
||||
| `benchmark_tiered_memory_v2.py` | 1M 级 durable page store 压力测试 |
|
||||
| `benchmark_memory_v2.py` | 独立存储/路由/完整性评测 |
|
||||
| `benchmark_natural_memory_v1.py` | 兼容 v1/v2 的 Qwen 综合回归评测 |
|
||||
| `test_natural_memory_v2_restart.py` | 真实模型释放、重载、无历史召回测试 |
|
||||
| `tests/test_memory_os_v2.py` | V2 核心单元测试 |
|
||||
|
||||
## 当前限制与下一阶段
|
||||
|
||||
已经实现的是可运行的 V2 内核和百万级 durable 存储后端；模型质量和长上下文仍有明确边界：
|
||||
|
||||
1. 真实 Qwen 路由器只用 512 条本地合成事实训练，需继续做跨实体、改写、冲突、时间和未知事实泛化；
|
||||
2. 多跳控制器的 held-out accuracy 只有 36.70%，需要专门的多跳 curriculum；
|
||||
3. 当前交付包使用 embedded weight-shard 模式，所有 V2 记录随第三个 safetensors 切片加载进 RAM；热点记录优先进入 VRAM，最多 256 条/131,072 token，并动态保留 2 GiB 显存安全余量；显存不足时自动留在 RAM，避免把模型推理和 KV 顶满；
|
||||
4. `memory_max_pages=32768` 代表 32K 页 × 32 条记录的地址上限；1M 轻量记录已经完成存储压力测试，但不是 1M 条完整长文本的 Qwen 端到端质量验证；
|
||||
5. 还需要加入更强的压缩摘要、原文页、事件时间线、事实置信度校准和用户级撤销日志；
|
||||
6. 还需要长序列课程训练、NVMe 级分页和更大规模端到端评测，才能验证 128K 热 KV 与百万级历史的实际吞吐和质量。
|
||||
|
||||
这些限制是设计边界，不是用一个“无限上下文”数字掩盖的未验证假设。
|
||||
+15
@@ -0,0 +1,15 @@
|
||||
"""Dynamic Memory Lab: small, reproducible architecture experiments."""
|
||||
|
||||
from .model import DynamicMemoryConfig, DynamicMemoryLM
|
||||
from .qwen_integration import QwenDynamicMemoryModel, QwenMemoryConfig, load_qwen_base, load_qwen_dynamic
|
||||
from .tiered_memory_store_v2 import TieredMemoryStoreV2
|
||||
|
||||
__all__ = [
|
||||
"DynamicMemoryConfig",
|
||||
"DynamicMemoryLM",
|
||||
"QwenDynamicMemoryModel",
|
||||
"QwenMemoryConfig",
|
||||
"load_qwen_base",
|
||||
"load_qwen_dynamic",
|
||||
"TieredMemoryStoreV2",
|
||||
]
|
||||
@@ -0,0 +1,203 @@
|
||||
"""Low-pressure end-to-end test for the embedded Natural Memory v2 path.
|
||||
|
||||
This benchmark intentionally uses only the model package's third safetensors
|
||||
memory shard. It does not create SQLite files or exercise disk paging. The
|
||||
long-context cases lower the temporary KV budget so the test measures the
|
||||
model-owned archive/read path without asking a 12 GiB GPU to hold a huge KV.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import random
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from .qwen_integration import load_qwen_dynamic, load_tokenizer
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--model-path",
|
||||
default=r"W:\Flash\model\dynamic_memory_lab\qwen3_5_4b_natural_memory_v2",
|
||||
)
|
||||
parser.add_argument("--lengths", default="4096,8192,16384,32768")
|
||||
parser.add_argument("--kv-budget", type=int, default=2048)
|
||||
parser.add_argument("--chunk-tokens", type=int, default=1024)
|
||||
parser.add_argument("--max-new-tokens", type=int, default=8)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
default=r"W:\Flash\model\dynamic_memory_lab\embedded_memory_v2_long_benchmark.json",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def build_prompt(tokenizer: Any, target_tokens: int, seed: int) -> tuple[str, str, int]:
|
||||
rng = random.Random(seed + target_tokens)
|
||||
answer = f"EMBEDDED-LONG-{target_tokens}-{rng.randrange(100000, 999999)}"
|
||||
needle = f"长期记忆锚点：唯一编号是 {answer}。"
|
||||
filler = (
|
||||
"这是长文本记忆压力测试中的普通背景段落，包含项目说明、日期、日志和无关备注。"
|
||||
"这些内容不是问题答案，读取时应保留原文但忽略干扰。"
|
||||
)
|
||||
chunks: list[str] = []
|
||||
while len(tokenizer(" ".join(chunks + [filler, needle]), add_special_tokens=False)["input_ids"]) < target_tokens:
|
||||
chunks.append(filler)
|
||||
# Keep the needle safely inside the archived prefix even in the smallest
|
||||
# case, so a pass must come from memory rather than the retained window.
|
||||
pivot = max(1, len(chunks) // 3)
|
||||
material = " ".join(chunks[:pivot] + [needle] + chunks[pivot:])
|
||||
prompt = (
|
||||
"请阅读下面的长材料，回答末尾问题，只输出编号，不要解释。\n"
|
||||
"---开始材料---\n"
|
||||
f"{material}\n"
|
||||
"---结束材料---\n"
|
||||
"问题：长期记忆锚点的唯一编号是什么？"
|
||||
)
|
||||
prompt_tokens = len(tokenizer(prompt, add_special_tokens=False)["input_ids"])
|
||||
return prompt, answer, prompt_tokens
|
||||
|
||||
|
||||
def chat_inputs(tokenizer: Any, text: str, device: torch.device) -> dict[str, torch.Tensor]:
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
[{"role": "user", "content": text}],
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
return_tensors="pt",
|
||||
return_dict=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
return {
|
||||
key: value.to(device)
|
||||
for key, value in encoded.items()
|
||||
if isinstance(value, torch.Tensor)
|
||||
}
|
||||
|
||||
|
||||
def run_case(model: Any, tokenizer: Any, target_tokens: int, args: argparse.Namespace) -> dict[str, Any]:
|
||||
device = model._find_layer_device()
|
||||
model.reset_memory(batch_size=1, device=device)
|
||||
assert model.memory_os_v2 is not None
|
||||
model.memory_os_v2.kv_budget.max_tokens = int(args.kv_budget)
|
||||
prompt, answer, prompt_tokens = build_prompt(tokenizer, target_tokens, 20260904)
|
||||
encoded = chat_inputs(tokenizer, prompt, device)
|
||||
query_text = "长期记忆锚点的唯一编号"
|
||||
query = tokenizer(query_text, add_special_tokens=False, return_tensors="pt")
|
||||
query_ids = query["input_ids"].to(device)
|
||||
query_mask = query.get("attention_mask")
|
||||
if query_mask is None:
|
||||
query_mask = torch.ones_like(query_ids)
|
||||
query_mask = query_mask.to(device)
|
||||
started = time.perf_counter()
|
||||
with torch.inference_mode():
|
||||
output = model.generate(
|
||||
**encoded,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
do_sample=False,
|
||||
use_cache=True,
|
||||
update_memory=False,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
memory_query_input_ids=query_ids,
|
||||
memory_query_attention_mask=query_mask,
|
||||
memory_query_text=query_text,
|
||||
)
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize(device)
|
||||
elapsed = time.perf_counter() - started
|
||||
response_ids = output[0, encoded["input_ids"].shape[1] :]
|
||||
response = tokenizer.decode(response_ids.detach().cpu().tolist(), skip_special_tokens=True).strip()
|
||||
stats = model.memory_v2_stats()
|
||||
records = model.memory_os_v2.bank.records
|
||||
archived_text_hit = any(
|
||||
answer in tokenizer.decode(record.token_ids.tolist(), skip_special_tokens=True)
|
||||
for record in records.values()
|
||||
if record.memory_type == "context_chunk" and isinstance(record.token_ids, torch.Tensor)
|
||||
)
|
||||
query_key = model._encode_model_key(query_ids, query_mask)[0]
|
||||
retrieved, decision = model.read_hierarchical_memory(
|
||||
query_key,
|
||||
query_text=query_text,
|
||||
query_token_ids=query_ids[0],
|
||||
top_k_pages=model.memory_config.memory_top_k_pages,
|
||||
top_k_records=model.memory_config.memory_top_k_records,
|
||||
max_hops=model.memory_config.memory_max_hops,
|
||||
)
|
||||
retrieved_hit = any(
|
||||
answer in tokenizer.decode(record.token_ids.tolist(), skip_special_tokens=True)
|
||||
for record in retrieved
|
||||
if isinstance(record.token_ids, torch.Tensor)
|
||||
)
|
||||
return {
|
||||
"target_tokens": target_tokens,
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"kv_budget_tokens": args.kv_budget,
|
||||
"chunk_tokens": args.chunk_tokens,
|
||||
"response": response,
|
||||
"expected": answer,
|
||||
"generation_hit": answer in response,
|
||||
"archived_text_hit": archived_text_hit,
|
||||
"retrieved_text_hit": retrieved_hit,
|
||||
"retrieved_records": len(retrieved),
|
||||
"router_stop_reason": decision.stop_reason,
|
||||
"router_hop_count": decision.hop_count,
|
||||
"seconds": elapsed,
|
||||
"records": stats.get("records", 0),
|
||||
"pages": stats.get("pages", 0),
|
||||
"gpu_cache_records": stats.get("gpu_cache_records", 0),
|
||||
"gpu_cache_tokens": stats.get("gpu_cache_tokens", 0),
|
||||
"gpu_cache_device": stats.get("gpu_cache_device", "none"),
|
||||
"passed": bool(archived_text_hit and retrieved_hit),
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
lengths = [int(item.strip()) for item in args.lengths.split(",") if item.strip()]
|
||||
tokenizer = load_tokenizer(args.model_path)
|
||||
model = load_qwen_dynamic(args.model_path, load_in_4bit=not args.no_4bit)
|
||||
model.eval()
|
||||
model.memory_config.context_chunk_tokens = int(args.chunk_tokens)
|
||||
report: dict[str, Any] = {
|
||||
"benchmark": "Natural Memory v2 embedded third-shard long-memory test",
|
||||
"storage_mode": model.memory_config.memory_storage_mode,
|
||||
"tier_store_enabled": bool(model.memory_os_v2 and model.memory_os_v2.bank.tier_store is not None),
|
||||
"model_path": str(Path(args.model_path).resolve()),
|
||||
"lengths": lengths,
|
||||
"quantization": "4bit_nf4" if not args.no_4bit else "none",
|
||||
"kv_budget_tokens": args.kv_budget,
|
||||
"chunk_tokens": args.chunk_tokens,
|
||||
"rows": [],
|
||||
}
|
||||
try:
|
||||
if report["storage_mode"] != "embedded" or report["tier_store_enabled"]:
|
||||
raise RuntimeError("embedded benchmark requires memory_storage_mode=embedded and no tier store")
|
||||
for target_tokens in lengths:
|
||||
row = run_case(model, tokenizer, target_tokens, args)
|
||||
report["rows"].append(row)
|
||||
print(json.dumps(row, ensure_ascii=False))
|
||||
finally:
|
||||
model.reset_memory(batch_size=1, device=model._find_layer_device())
|
||||
model.close_memory_storage()
|
||||
del model
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
report["passed_cases"] = sum(bool(row["passed"]) for row in report["rows"])
|
||||
report["total_cases"] = len(report["rows"])
|
||||
report["all_passed"] = report["passed_cases"] == report["total_cases"]
|
||||
output = Path(args.output)
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(json.dumps(report, ensure_ascii=False, indent=2))
|
||||
print(f"saved={output}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,128 @@
|
||||
"""Smoke-test the model-owned KV offload and long-context compaction paths.
|
||||
|
||||
The test intentionally uses the real Qwen package. It verifies that
|
||||
Transformers' CPU-backed DynamicCache can generate through Qwen3.5's hybrid
|
||||
linear/full-attention stack and that an over-budget prompt is archived into
|
||||
Natural Memory before only the recent hot window is passed to generation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
if __package__ in {None, ""}:
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
from dynamic_memory_lab.qwen_integration import load_qwen_dynamic, load_tokenizer
|
||||
from dynamic_memory_lab.stream_chat_qwen_memory import _chat_tensor
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--model-path",
|
||||
default=r"W:\Flash\model\dynamic_memory_lab\qwen3_5_4b_natural_memory_v2",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--report",
|
||||
default=r"W:\Flash\model\dynamic_memory_lab\kv_offload_compaction_test.json",
|
||||
)
|
||||
parser.add_argument("--compact-hot-tokens", type=int, default=32)
|
||||
parser.add_argument("--max-new-tokens", type=int, default=4)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
model_path = Path(args.model_path)
|
||||
tokenizer = load_tokenizer(model_path)
|
||||
model = load_qwen_dynamic(model_path, load_in_4bit=not args.no_4bit)
|
||||
model.eval()
|
||||
model.memory_config.kv_offload = True
|
||||
device = model._find_layer_device()
|
||||
|
||||
answer_inputs = {
|
||||
key: value.to(device)
|
||||
for key, value in _chat_tensor(tokenizer, "请用一句话说明你支持什么。").items()
|
||||
}
|
||||
generated = model.generate(
|
||||
**answer_inputs,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
answer = tokenizer.decode(generated[0].detach().cpu().tolist(), skip_special_tokens=True).strip()
|
||||
|
||||
if model.memory_os_v2 is None:
|
||||
raise RuntimeError("the selected package does not contain Natural Memory v2")
|
||||
model.memory_os_v2.kv_budget.max_tokens = int(args.compact_hot_tokens)
|
||||
model.memory_os_v2.kv_budget.keep_recent_tokens = min(
|
||||
model.memory_os_v2.kv_budget.keep_recent_tokens,
|
||||
model.memory_os_v2.kv_budget.max_tokens,
|
||||
)
|
||||
model.memory_config.context_chunk_tokens = max(4, min(16, args.compact_hot_tokens // 2))
|
||||
long_text = " ".join(
|
||||
[
|
||||
"历史上下文片段用于验证 Natural Memory 的自动分页压缩。",
|
||||
"这段内容应该被写入长期上下文页面，而不是继续占用当前热 KV。",
|
||||
]
|
||||
* 12
|
||||
)
|
||||
encoded = tokenizer(long_text, add_special_tokens=False, return_tensors="pt")
|
||||
long_ids = encoded["input_ids"].to(device)
|
||||
long_mask = encoded.get("attention_mask")
|
||||
if long_mask is None:
|
||||
long_mask = torch.ones_like(long_ids)
|
||||
long_mask = long_mask.to(device)
|
||||
before = model.memory_v2_stats()
|
||||
generated_long = model.generate(
|
||||
input_ids=long_ids,
|
||||
attention_mask=long_mask,
|
||||
max_new_tokens=1,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
compaction = dict(model.runtime.context_compaction or {})
|
||||
after = model.memory_v2_stats()
|
||||
compaction_pass = bool(
|
||||
compaction["compacted"]
|
||||
and compaction["archived_records"] > 0
|
||||
and max(compaction["retained_tokens"]) <= args.compact_hot_tokens
|
||||
and int(compaction["retained_tokens"][0]) <= args.compact_hot_tokens
|
||||
and after["active_records"] >= before["active_records"] + compaction["archived_records"]
|
||||
)
|
||||
|
||||
report = {
|
||||
"model_path": str(model_path),
|
||||
"transformers_cache": "DynamicCache(offloading=True)",
|
||||
"kv_offload_pass": True,
|
||||
"answer": answer,
|
||||
"compaction": compaction,
|
||||
"compaction_pass": compaction_pass,
|
||||
"long_generation_tokens": int(generated_long.shape[1]),
|
||||
"memory_before": before,
|
||||
"memory_after": after,
|
||||
"device": str(device),
|
||||
"cuda_memory_allocated_bytes": (
|
||||
int(torch.cuda.memory_allocated(device)) if torch.cuda.is_available() else 0
|
||||
),
|
||||
}
|
||||
report_path = Path(args.report)
|
||||
report_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(json.dumps(report, ensure_ascii=False, indent=2))
|
||||
if not compaction_pass:
|
||||
raise SystemExit("context compaction verification failed")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,190 @@
|
||||
"""Practical long-context stress probe for the local Qwen/Natural Memory setup."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import random
|
||||
import statistics
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from .qwen_integration import load_qwen_base, load_qwen_dynamic, load_tokenizer
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--base-model", default=r"W:\Flash\model")
|
||||
parser.add_argument(
|
||||
"--memory-model",
|
||||
default=r"W:\Flash\model\dynamic_memory_lab\qwen3_5_4b_memory_merged_v13",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
default=r"W:\Flash\model\dynamic_memory_lab\long_context_v1_stress.json",
|
||||
)
|
||||
parser.add_argument("--lengths", default="8192,16384,32768")
|
||||
parser.add_argument("--max-new-tokens", type=int, default=4)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def build_prompt(tokenizer: Any, target_tokens: int, seed: int) -> tuple[str, str, int]:
|
||||
rng = random.Random(seed + target_tokens)
|
||||
answer = f"LONGCTX-{target_tokens}-{rng.randrange(100000, 999999)}"
|
||||
needle = f"唯一目标记录：检索编码为 {answer}。"
|
||||
filler = (
|
||||
"这段材料是用于压力测试的背景文本。它包含版本、日志、普通备注和无关的项目描述，"
|
||||
"不包含目标编码。阅读时应忽略这些干扰内容，只寻找唯一目标记录。"
|
||||
)
|
||||
chunks: list[str] = []
|
||||
while len(tokenizer(" ".join(chunks + [filler, needle]), add_special_tokens=False)["input_ids"]) < target_tokens:
|
||||
chunks.append(filler)
|
||||
half = len(chunks) // 2
|
||||
material = " ".join(chunks[:half] + [needle] + chunks[half:])
|
||||
prompt = (
|
||||
"请阅读下面的长材料，只输出唯一目标记录中的检索编码，不要解释。\n"
|
||||
"---开始---\n"
|
||||
f"{material}\n"
|
||||
"---结束---\n"
|
||||
"问题：唯一目标记录中的检索编码是什么？"
|
||||
)
|
||||
prompt_tokens = len(tokenizer(prompt, add_special_tokens=False)["input_ids"])
|
||||
return prompt, answer, prompt_tokens
|
||||
|
||||
|
||||
def make_inputs(tokenizer: Any, prompt: str, device: torch.device) -> dict[str, torch.Tensor]:
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
[{"role": "user", "content": prompt}],
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
return_tensors="pt",
|
||||
return_dict=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
return {
|
||||
key: value.to(device)
|
||||
for key, value in encoded.items()
|
||||
if isinstance(value, torch.Tensor)
|
||||
}
|
||||
|
||||
|
||||
def probe_model(model: Any, tokenizer: Any, lengths: list[int], *, dynamic: bool, max_new_tokens: int) -> dict[str, Any]:
|
||||
device = model._find_layer_device() if dynamic else model.get_input_embeddings().weight.device
|
||||
rows = []
|
||||
for length in lengths:
|
||||
prompt, answer, prompt_tokens = build_prompt(tokenizer, length, 20260904)
|
||||
if dynamic:
|
||||
model.reset_memory(batch_size=1, device=device)
|
||||
if device.type == "cuda":
|
||||
torch.cuda.reset_peak_memory_stats(device)
|
||||
row: dict[str, Any] = {
|
||||
"target_tokens": length,
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"expected": answer,
|
||||
}
|
||||
try:
|
||||
encoded = make_inputs(tokenizer, prompt, device)
|
||||
query = tokenizer(prompt, add_special_tokens=False, return_tensors="pt")
|
||||
query_ids = query["input_ids"].to(device)
|
||||
query_mask = query.get("attention_mask")
|
||||
if query_mask is None:
|
||||
query_mask = torch.ones_like(query_ids)
|
||||
query_mask = query_mask.to(device)
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize(device)
|
||||
started = time.perf_counter()
|
||||
with torch.inference_mode():
|
||||
kwargs: dict[str, Any] = {
|
||||
"max_new_tokens": max_new_tokens,
|
||||
"do_sample": False,
|
||||
"use_cache": True,
|
||||
"pad_token_id": tokenizer.pad_token_id,
|
||||
}
|
||||
if dynamic:
|
||||
kwargs.update(
|
||||
{
|
||||
"update_memory": False,
|
||||
"memory_query_input_ids": query_ids,
|
||||
"memory_query_attention_mask": query_mask,
|
||||
}
|
||||
)
|
||||
output = model.generate(**encoded, **kwargs)
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize(device)
|
||||
elapsed = time.perf_counter() - started
|
||||
response_ids = output[0, encoded["input_ids"].shape[1] :]
|
||||
response = tokenizer.decode(response_ids.detach().cpu().tolist(), skip_special_tokens=True).strip()
|
||||
row.update(
|
||||
{
|
||||
"status": "ok",
|
||||
"response": response,
|
||||
"passed": answer in response,
|
||||
"generated_tokens": int(response_ids.numel()),
|
||||
"seconds": elapsed,
|
||||
"tokens_per_second": int(response_ids.numel()) / max(elapsed, 1e-9),
|
||||
}
|
||||
)
|
||||
except (torch.cuda.OutOfMemoryError, RuntimeError) as exc:
|
||||
message = str(exc)
|
||||
if isinstance(exc, torch.cuda.OutOfMemoryError) or "out of memory" in message.lower():
|
||||
row.update({"status": "cuda_oom", "error": message[:500]})
|
||||
if device.type == "cuda":
|
||||
torch.cuda.empty_cache()
|
||||
else:
|
||||
raise
|
||||
if device.type == "cuda":
|
||||
row["peak_memory_allocated_gb"] = torch.cuda.max_memory_allocated(device) / 1024**3
|
||||
row["peak_memory_reserved_gb"] = torch.cuda.max_memory_reserved(device) / 1024**3
|
||||
rows.append(row)
|
||||
return {"rows": rows}
|
||||
|
||||
|
||||
def release(model: Any) -> None:
|
||||
del model
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
lengths = [int(value.strip()) for value in args.lengths.split(",") if value.strip()]
|
||||
use_4bit = not args.no_4bit
|
||||
tokenizer = load_tokenizer(args.base_model)
|
||||
report: dict[str, Any] = {
|
||||
"benchmark": "Natural Memory v1 practical long-context stress probe",
|
||||
"date": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"base_model": str(Path(args.base_model).resolve()),
|
||||
"memory_model": str(Path(args.memory_model).resolve()),
|
||||
"lengths": lengths,
|
||||
"quantization": "4bit_nf4" if use_4bit else "none",
|
||||
"max_new_tokens": args.max_new_tokens,
|
||||
}
|
||||
print("loading baseline")
|
||||
model = load_qwen_base(args.base_model, load_in_4bit=use_4bit)
|
||||
model.eval()
|
||||
report["baseline"] = probe_model(
|
||||
model, tokenizer, lengths, dynamic=False, max_new_tokens=args.max_new_tokens
|
||||
)
|
||||
release(model)
|
||||
print("loading Natural Memory v1")
|
||||
model = load_qwen_dynamic(args.memory_model, load_in_4bit=use_4bit)
|
||||
model.eval()
|
||||
report["natural_memory_v1"] = probe_model(
|
||||
model, tokenizer, lengths, dynamic=True, max_new_tokens=args.max_new_tokens
|
||||
)
|
||||
release(model)
|
||||
output = Path(args.output)
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(json.dumps(report, ensure_ascii=False, indent=2))
|
||||
print(f"saved={output}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,301 @@
|
||||
"""Comprehensive, hardware-independent evaluation for Natural Memory v2."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
if __package__ in {None, ""}:
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
from dynamic_memory_lab.memory_os_v2 import (
|
||||
KVBudgetManagerV2,
|
||||
MemoryRouterV2,
|
||||
PagedMemoryBankV2,
|
||||
STATUS_ACTIVE,
|
||||
STATUS_QUARANTINED,
|
||||
STATUS_RETRACTED,
|
||||
STATUS_SUPERSEDED,
|
||||
)
|
||||
from dynamic_memory_lab.train_memory_router_v2 import _latent_to_hidden, _make_basis, evaluate
|
||||
|
||||
|
||||
def _load_router(args: argparse.Namespace, device: torch.device) -> tuple[MemoryRouterV2, torch.Tensor, str]:
|
||||
router = MemoryRouterV2(
|
||||
args.hidden_size,
|
||||
router_dim=args.router_dim,
|
||||
num_heads=args.num_heads,
|
||||
max_hops=args.max_hops,
|
||||
).to(device)
|
||||
checkpoint = Path(args.router_checkpoint)
|
||||
basis_path = checkpoint.with_name("memory_router_v2_basis.pt")
|
||||
if checkpoint.exists() and basis_path.exists():
|
||||
state = torch.load(checkpoint, map_location=device, weights_only=True)
|
||||
router.load_state_dict(state, strict=True)
|
||||
basis = torch.load(basis_path, map_location=device, weights_only=True).to(device)
|
||||
return router, basis, "trained_checkpoint"
|
||||
basis = _make_basis(args.hidden_size, args.latent_size, device)
|
||||
return router, basis, "untrained_router"
|
||||
|
||||
|
||||
def _metric(value: bool) -> float:
|
||||
return 1.0 if value else 0.0
|
||||
|
||||
|
||||
def run(args: argparse.Namespace) -> dict[str, Any]:
|
||||
torch.manual_seed(args.seed)
|
||||
device = torch.device(args.device if args.device != "auto" else "cuda" if torch.cuda.is_available() else "cpu")
|
||||
router, basis, router_source = _load_router(args, device)
|
||||
router.eval()
|
||||
bank = PagedMemoryBankV2(
|
||||
args.hidden_size,
|
||||
router=router,
|
||||
page_capacity=args.page_capacity,
|
||||
max_pages=args.max_pages,
|
||||
hot_pages=args.hot_pages,
|
||||
top_k_pages=args.top_k_pages,
|
||||
top_k_records=args.top_k_records,
|
||||
max_hops=args.max_hops,
|
||||
coarse_index_bits=args.coarse_index_bits,
|
||||
)
|
||||
|
||||
# Route quality on held-out samples from the same latent factor space.
|
||||
route_metrics = evaluate(
|
||||
router,
|
||||
basis=basis,
|
||||
device=device,
|
||||
batches=args.router_eval_batches,
|
||||
batch_size=args.router_eval_batch_size,
|
||||
candidate_count=args.candidate_count,
|
||||
)
|
||||
|
||||
# Populate a large enough store to activate the coarse index. Keys are
|
||||
# low-entropy semantic points, while their positions are deliberately
|
||||
# unrelated to their topic ids.
|
||||
records: list[Any] = []
|
||||
started = time.perf_counter()
|
||||
for index in range(args.records):
|
||||
latent = torch.randn(basis.shape[1], device=device)
|
||||
key = _latent_to_hidden(latent.unsqueeze(0), basis, 0.04)[0]
|
||||
record, action = bank.write(
|
||||
text=f"长期事实 {index}",
|
||||
key=key,
|
||||
token_ids=torch.tensor([index % 997, 17, 23]),
|
||||
token_mask=torch.tensor([True, True, True]),
|
||||
importance=0.5 + 0.5 * (index % 7 == 0),
|
||||
confidence=0.95,
|
||||
source="synthetic_episode",
|
||||
)
|
||||
records.append(record)
|
||||
write_seconds = time.perf_counter() - started
|
||||
|
||||
recall_hits = 0
|
||||
page_hits = 0
|
||||
candidate_counts: list[int] = []
|
||||
query_count = min(args.query_count, len(records))
|
||||
for index in torch.randperm(len(records), device=device)[:query_count].tolist():
|
||||
target = records[index]
|
||||
query_key = target.key.to(device) if router_source == "untrained_router" else (
|
||||
target.key.to(device)
|
||||
)
|
||||
# ``target.key`` is already in compact address space. This tests the
|
||||
# same storage-space path used after a Qwen hidden state is projected.
|
||||
found, decision = bank.query(
|
||||
query_key=query_key,
|
||||
top_k_pages=args.top_k_pages,
|
||||
top_k_records=args.top_k_records,
|
||||
)
|
||||
found_ids = {record.record_id for record in found}
|
||||
recall_hits += int(target.record_id in found_ids)
|
||||
page_hits += int(target.page_id in decision.page_ids)
|
||||
candidate_counts.append(bank.stats()["last_coarse_candidates"])
|
||||
|
||||
# Conflict/versioning and explicit correction.
|
||||
conflict_key = torch.randn(args.hidden_size, device=device)
|
||||
first, _ = bank.write(
|
||||
text="用户当前工作地点是上海",
|
||||
key=conflict_key,
|
||||
entity="user",
|
||||
attribute="work_city",
|
||||
value="上海",
|
||||
confidence=0.90,
|
||||
)
|
||||
second, conflict_action = bank.write(
|
||||
text="用户当前工作地点是杭州",
|
||||
key=conflict_key,
|
||||
entity="user",
|
||||
attribute="work_city",
|
||||
value="杭州",
|
||||
confidence=0.98,
|
||||
)
|
||||
corrected, correction_action = bank.correct(
|
||||
text="纠正：用户当前工作地点是苏州",
|
||||
key=conflict_key,
|
||||
entity="user",
|
||||
attribute="work_city",
|
||||
value="苏州",
|
||||
confidence=1.0,
|
||||
)
|
||||
|
||||
# Pollution protection: untrusted write stays out of the searchable bank.
|
||||
quarantined, quarantine_action = bank.write(
|
||||
text="模型猜测的生日",
|
||||
key=torch.randn(args.hidden_size, device=device),
|
||||
confidence=0.05,
|
||||
trusted=False,
|
||||
)
|
||||
quarantine_before_approval = (
|
||||
quarantined.status == STATUS_QUARANTINED
|
||||
and quarantined.record_id not in bank.records
|
||||
and quarantine_action == "quarantined"
|
||||
)
|
||||
approved = bank.approve(quarantined.record_id)
|
||||
approved_active = approved.status == STATUS_ACTIVE
|
||||
|
||||
# Multi-hop: source page contains only the anchor; related evidence lives
|
||||
# in other pages. Restrict first-hop page selection to force expansion.
|
||||
hop_bank = PagedMemoryBankV2(
|
||||
args.hidden_size,
|
||||
router=router,
|
||||
page_capacity=1,
|
||||
max_pages=64,
|
||||
hot_pages=1,
|
||||
top_k_pages=1,
|
||||
top_k_records=3,
|
||||
max_hops=args.max_hops,
|
||||
coarse_index_bits=args.coarse_index_bits,
|
||||
)
|
||||
hop_b, _ = hop_bank.write(text="链路证据 B", key=torch.randn(args.hidden_size, device=device), slot_index=20001)
|
||||
hop_c, _ = hop_bank.write(text="链路证据 C", key=torch.randn(args.hidden_size, device=device), slot_index=20002)
|
||||
hop_a, _ = hop_bank.write(
|
||||
text="链路锚点 A",
|
||||
key=torch.randn(args.hidden_size, device=device),
|
||||
related_ids=[hop_b.record_id, hop_c.record_id],
|
||||
slot_index=20000,
|
||||
)
|
||||
hop_records, hop_decision = hop_bank.query(
|
||||
query_key=hop_a.key,
|
||||
top_k_pages=1,
|
||||
top_k_records=3,
|
||||
max_hops=args.max_hops,
|
||||
)
|
||||
hop_ids = {record.record_id for record in hop_records}
|
||||
multi_hop_success = hop_b.record_id in hop_ids or hop_c.record_id in hop_ids
|
||||
|
||||
# Idempotence, retraction and restart serialization.
|
||||
duplicate, duplicate_action = bank.write(
|
||||
text="长期事实 0",
|
||||
key=records[0].key,
|
||||
token_ids=records[0].token_ids,
|
||||
token_mask=records[0].token_mask,
|
||||
confidence=0.99,
|
||||
)
|
||||
bank.retract(approved.record_id)
|
||||
restart_payload = bank.export_payload()
|
||||
restored = PagedMemoryBankV2.from_payload(restart_payload, router=router)
|
||||
restored_records, restored_decision = restored.query(
|
||||
query_key=records[0].key,
|
||||
top_k_pages=args.top_k_pages,
|
||||
top_k_records=args.top_k_records,
|
||||
)
|
||||
|
||||
budget = KVBudgetManagerV2(
|
||||
max_tokens=args.kv_budget,
|
||||
hard_max_tokens=args.kv_hard_max,
|
||||
keep_recent_tokens=args.kv_keep_recent,
|
||||
)
|
||||
budget_checks = {
|
||||
"below_trigger": not budget.needs_compaction(int(args.kv_budget * 0.5)),
|
||||
"at_trigger": budget.needs_compaction(budget.trigger_tokens),
|
||||
"overflow": budget.overflow(args.kv_budget + 123),
|
||||
}
|
||||
|
||||
stats = bank.stats()
|
||||
summary = {
|
||||
"format_version": 2,
|
||||
"seed": args.seed,
|
||||
"device": str(device),
|
||||
"router_source": router_source,
|
||||
"router": route_metrics,
|
||||
"storage": {
|
||||
"records_requested": args.records,
|
||||
"records_stored_before_scenarios": len(records),
|
||||
"write_seconds": write_seconds,
|
||||
"pages": stats["pages"],
|
||||
"coarse_index_buckets": stats["coarse_index_buckets"],
|
||||
"coarse_candidate_mean": sum(candidate_counts) / max(1, len(candidate_counts)),
|
||||
"coarse_candidate_max": max(candidate_counts, default=0),
|
||||
"coarse_candidate_ratio": (
|
||||
sum(candidate_counts) / max(1, len(candidate_counts)) / max(1, stats["pages"])
|
||||
),
|
||||
},
|
||||
"retrieval": {
|
||||
"query_count": query_count,
|
||||
"record_recall_at_k": recall_hits / max(1, query_count),
|
||||
"page_recall_at_k": page_hits / max(1, query_count),
|
||||
"multi_hop_success": _metric(multi_hop_success),
|
||||
"multi_hop_hops": hop_decision.hop_count,
|
||||
"restart_record_recall": _metric(bool(restored_records)),
|
||||
"restart_page_count": restored.stats()["pages"],
|
||||
},
|
||||
"integrity": {
|
||||
"conflict_action": conflict_action,
|
||||
"correction_action": correction_action,
|
||||
"old_conflict_superseded": _metric(first.status == STATUS_SUPERSEDED),
|
||||
"latest_correction_active": _metric(corrected.status == STATUS_ACTIVE),
|
||||
"active_conflict_value": corrected.value,
|
||||
"quarantine_action": quarantine_action,
|
||||
"quarantine_isolation": _metric(quarantine_before_approval),
|
||||
"approved_active": _metric(approved_active),
|
||||
"retracted_status": bank.records[approved.record_id].status,
|
||||
"retraction_isolated": _metric(bank.records[approved.record_id].status == STATUS_RETRACTED),
|
||||
"duplicate_action": duplicate_action,
|
||||
"duplicate_idempotent": _metric(duplicate.record_id == records[0].record_id),
|
||||
},
|
||||
"kv_budget": budget_checks,
|
||||
"final_stats": stats,
|
||||
}
|
||||
output_path = Path(args.output)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_text(json.dumps(summary, indent=2, ensure_ascii=False), encoding="utf-8")
|
||||
return summary
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--output", default="W:/Flash/model/dynamic_memory_lab/natural_memory_v2_benchmark.json")
|
||||
parser.add_argument("--router-checkpoint", default="W:/Flash/model/dynamic_memory_lab/checkpoints/natural_memory_v2_router/memory_router_v2.pt")
|
||||
parser.add_argument("--device", default="auto")
|
||||
parser.add_argument("--hidden-size", type=int, default=2560)
|
||||
parser.add_argument("--router-dim", type=int, default=128)
|
||||
parser.add_argument("--num-heads", type=int, default=8)
|
||||
parser.add_argument("--max-hops", type=int, default=3)
|
||||
parser.add_argument("--latent-size", type=int, default=32)
|
||||
parser.add_argument("--candidate-count", type=int, default=32)
|
||||
parser.add_argument("--router-eval-batches", type=int, default=40)
|
||||
parser.add_argument("--router-eval-batch-size", type=int, default=64)
|
||||
parser.add_argument("--records", type=int, default=512)
|
||||
parser.add_argument("--query-count", type=int, default=128)
|
||||
parser.add_argument("--page-capacity", type=int, default=32)
|
||||
parser.add_argument("--max-pages", type=int, default=32768)
|
||||
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=20)
|
||||
parser.add_argument("--kv-budget", type=int, default=32768)
|
||||
parser.add_argument("--kv-hard-max", type=int, default=131072)
|
||||
parser.add_argument("--kv-keep-recent", type=int, default=8192)
|
||||
parser.add_argument("--seed", type=int, default=20260904)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(json.dumps(run(parse_args()), ensure_ascii=False, indent=2))
|
||||
@@ -0,0 +1,291 @@
|
||||
"""Paired teacher/student benchmark for Natural Memory versus full KV context."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable
|
||||
|
||||
import torch
|
||||
|
||||
from .qwen_integration import load_memory_config, load_qwen_base, load_qwen_dynamic, load_tokenizer
|
||||
from .stream_chat_qwen_memory import _chat_tensor, _write_turn
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
def _path(value: str | Path) -> Path:
|
||||
path = Path(value)
|
||||
if path.is_absolute() or path.exists():
|
||||
return path
|
||||
return PROJECT_ROOT / path
|
||||
|
||||
|
||||
def _normalize(text: str) -> str:
|
||||
return re.sub(r"[\s`*_#，。！？、；：,.!?;:'\"（）()\[\]{}]", "", str(text)).lower()
|
||||
|
||||
|
||||
def _contains(text: str, choices: Iterable[str]) -> bool:
|
||||
normalized = _normalize(text)
|
||||
return any(_normalize(choice) and _normalize(choice) in normalized for choice in choices)
|
||||
|
||||
|
||||
def _passed(response: str, case: dict[str, Any]) -> bool:
|
||||
return _contains(response, case.get("acceptable", [])) and not _contains(
|
||||
response, case.get("forbidden", [])
|
||||
)
|
||||
|
||||
|
||||
def _read_cases(path: Path, *, limit: int | None, offset: int, category: str | None) -> list[dict[str, Any]]:
|
||||
rows: list[dict[str, Any]] = []
|
||||
with path.open("r", encoding="utf-8") as handle:
|
||||
for raw in handle:
|
||||
if not raw.strip():
|
||||
continue
|
||||
case = json.loads(raw)
|
||||
if category and case.get("category") != category:
|
||||
continue
|
||||
if offset > 0:
|
||||
offset -= 1
|
||||
continue
|
||||
rows.append(case)
|
||||
if limit is not None and len(rows) >= limit:
|
||||
break
|
||||
if not rows:
|
||||
raise ValueError("no validation cases selected")
|
||||
return rows
|
||||
|
||||
|
||||
def _teacher_messages(case: dict[str, Any]) -> list[dict[str, str]]:
|
||||
messages: list[dict[str, str]] = []
|
||||
for fact in case["facts"]:
|
||||
messages.append({"role": "user", "content": str(fact["text"])})
|
||||
messages.append({"role": "assistant", "content": str(fact.get("assistant", "好的。"))})
|
||||
messages.append({"role": "user", "content": str(case["query"])})
|
||||
return messages
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _generate_teacher(model: Any, tokenizer: Any, case: dict[str, Any], max_new_tokens: int) -> str:
|
||||
device = model.get_input_embeddings().weight.device
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
_teacher_messages(case),
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
return_tensors="pt",
|
||||
return_dict=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
encoded = {key: value.to(device) for key, value in encoded.items() if isinstance(value, torch.Tensor)}
|
||||
output = model.generate(
|
||||
**encoded,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
return tokenizer.decode(output[0, encoded["input_ids"].shape[1] :], skip_special_tokens=True).strip()
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _generate_student(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
case: dict[str, Any],
|
||||
*,
|
||||
max_new_tokens: int,
|
||||
force_write: bool,
|
||||
) -> tuple[str, dict[str, Any]]:
|
||||
device = model._find_layer_device()
|
||||
model.reset_memory(batch_size=1, device=device)
|
||||
write_rows: list[dict[str, Any]] = []
|
||||
for fact in case["facts"]:
|
||||
changed = _write_turn(
|
||||
model,
|
||||
tokenizer,
|
||||
str(fact["text"]),
|
||||
device,
|
||||
force_write=force_write,
|
||||
)
|
||||
last_slot = model.runtime.text_last_written_slot
|
||||
write_rows.append(
|
||||
{
|
||||
"kind": fact.get("kind", "fact"),
|
||||
"should_write": bool(fact.get("should_write", True)),
|
||||
"changed": bool(changed),
|
||||
"slot": int(last_slot[0].item()) if isinstance(last_slot, torch.Tensor) else -1,
|
||||
}
|
||||
)
|
||||
encoded = {key: value.to(device) for key, value in _chat_tensor(
|
||||
tokenizer,
|
||||
str(case["query"]),
|
||||
).items()}
|
||||
query = tokenizer(str(case["query"]), add_special_tokens=False, return_tensors="pt")
|
||||
query_ids = query["input_ids"].to(device)
|
||||
query_mask = query.get("attention_mask")
|
||||
if query_mask is None:
|
||||
query_mask = torch.ones_like(query_ids)
|
||||
query_mask = query_mask.to(device)
|
||||
output = model.generate(
|
||||
**encoded,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
memory_query_input_ids=query_ids,
|
||||
memory_query_attention_mask=query_mask,
|
||||
memory_query_text=str(case["query"]),
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
response = tokenizer.decode(output[0, encoded["input_ids"].shape[1] :], skip_special_tokens=True).strip()
|
||||
return response, {
|
||||
"writes": write_rows,
|
||||
"valid_slots": int(model.runtime.text_slot_valid.sum().item())
|
||||
if isinstance(model.runtime.text_slot_valid, torch.Tensor)
|
||||
else 0,
|
||||
"v2": model.memory_v2_stats(),
|
||||
}
|
||||
|
||||
|
||||
def _release(model: Any) -> None:
|
||||
del model
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def _summarize(rows: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
categories: dict[str, list[float]] = defaultdict(list)
|
||||
for row in rows:
|
||||
categories[str(row["category"])].append(float(row["passed"]))
|
||||
return {
|
||||
"cases": len(rows),
|
||||
"accuracy": sum(float(row["passed"]) for row in rows) / max(1, len(rows)),
|
||||
"categories": {
|
||||
category: {"cases": len(values), "accuracy": sum(values) / len(values)}
|
||||
for category, values in sorted(categories.items())
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
if hasattr(sys.stdout, "reconfigure"):
|
||||
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-path", default="qwen3_5_4b_natural_memory_v2")
|
||||
parser.add_argument("--adapter", default=None, help="optional candidate adapter; omitted uses embedded policy")
|
||||
parser.add_argument("--data", default="data/mega_validation/memory_validation_100k.jsonl")
|
||||
parser.add_argument("--output", default="mega_memory_vs_full_kv_report.json")
|
||||
parser.add_argument("--limit", type=int, default=32)
|
||||
parser.add_argument("--offset", type=int, default=0)
|
||||
parser.add_argument("--category", default=None)
|
||||
parser.add_argument("--max-new-tokens", type=int, default=24)
|
||||
parser.add_argument("--force-write", action="store_true")
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
args = parser.parse_args()
|
||||
model_path = _path(args.model_path)
|
||||
data_path = _path(args.data)
|
||||
output_path = _path(args.output)
|
||||
cases = _read_cases(data_path, limit=args.limit, offset=args.offset, category=args.category)
|
||||
tokenizer = load_tokenizer(model_path)
|
||||
use_4bit = not args.no_4bit
|
||||
report: dict[str, Any] = {
|
||||
"format_version": 1,
|
||||
"model_path": str(model_path),
|
||||
"data": str(data_path),
|
||||
"selected_cases": len(cases),
|
||||
"force_write": bool(args.force_write),
|
||||
"quantization": "4bit_nf4" if use_4bit else "none",
|
||||
}
|
||||
|
||||
started = time.perf_counter()
|
||||
teacher = load_qwen_base(model_path, load_in_4bit=use_4bit)
|
||||
teacher.eval()
|
||||
teacher_rows: list[dict[str, Any]] = []
|
||||
for index, case in enumerate(cases, 1):
|
||||
response = _generate_teacher(teacher, tokenizer, case, args.max_new_tokens)
|
||||
teacher_rows.append(
|
||||
{
|
||||
"id": case["id"],
|
||||
"category": case["category"],
|
||||
"response": response,
|
||||
"passed": _passed(response, case),
|
||||
}
|
||||
)
|
||||
if index % 16 == 0:
|
||||
print(f"teacher {index}/{len(cases)}")
|
||||
report["teacher"] = _summarize(teacher_rows)
|
||||
_release(teacher)
|
||||
|
||||
adapter_path = _path(args.adapter) if args.adapter else None
|
||||
config_source = adapter_path or model_path
|
||||
config = load_memory_config(config_source)
|
||||
config.persistent_memory = True
|
||||
config.natural_language_memory = True
|
||||
config.automatic_memory = True
|
||||
student = load_qwen_dynamic(model_path, memory_config=config, load_in_4bit=use_4bit)
|
||||
if adapter_path is not None:
|
||||
student.load_memory_adapter(adapter_path, strict=True)
|
||||
student.eval()
|
||||
student_rows: list[dict[str, Any]] = []
|
||||
for index, case in enumerate(cases, 1):
|
||||
response, diagnostics = _generate_student(
|
||||
student,
|
||||
tokenizer,
|
||||
case,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
force_write=args.force_write,
|
||||
)
|
||||
student_rows.append(
|
||||
{
|
||||
"id": case["id"],
|
||||
"category": case["category"],
|
||||
"response": response,
|
||||
"passed": _passed(response, case),
|
||||
"diagnostics": diagnostics,
|
||||
}
|
||||
)
|
||||
if index % 16 == 0:
|
||||
print(f"student {index}/{len(cases)}")
|
||||
report["student"] = _summarize(student_rows)
|
||||
_release(student)
|
||||
|
||||
teacher_accuracy = float(report["teacher"]["accuracy"])
|
||||
student_accuracy = float(report["student"]["accuracy"])
|
||||
teacher_by_id = {row["id"]: row for row in teacher_rows}
|
||||
paired_teacher_pass = sum(bool(teacher_by_id[row["id"]]["passed"]) for row in student_rows)
|
||||
paired_student_pass = sum(bool(row["passed"]) and teacher_by_id[row["id"]]["passed"] for row in student_rows)
|
||||
category_gate: dict[str, Any] = {}
|
||||
for category in sorted({str(case["category"]) for case in cases}):
|
||||
teacher_cat = [row for row in teacher_rows if row["category"] == category]
|
||||
student_cat = [row for row in student_rows if row["category"] == category]
|
||||
t = sum(float(row["passed"]) for row in teacher_cat) / max(1, len(teacher_cat))
|
||||
s = sum(float(row["passed"]) for row in student_cat) / max(1, len(student_cat))
|
||||
category_gate[category] = {"teacher_accuracy": t, "student_accuracy": s, "ratio": s / max(t, 1e-9), "pass": s >= 0.95 * t}
|
||||
report["parity"] = {
|
||||
"teacher_accuracy": teacher_accuracy,
|
||||
"student_accuracy": student_accuracy,
|
||||
"student_to_teacher_ratio": student_accuracy / max(teacher_accuracy, 1e-9),
|
||||
"paired_teacher_pass": paired_teacher_pass,
|
||||
"paired_student_pass": paired_student_pass,
|
||||
"paired_ratio": paired_student_pass / max(1, paired_teacher_pass),
|
||||
"category_gate": category_gate,
|
||||
"required_ratio": 0.95,
|
||||
"pass": student_accuracy >= 0.95 * teacher_accuracy and all(item["pass"] for item in category_gate.values()),
|
||||
}
|
||||
report["elapsed_seconds"] = time.perf_counter() - started
|
||||
report["failures"] = [row for row in student_rows if not row["passed"]][:100]
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(json.dumps({"teacher": report["teacher"], "student": report["student"], "parity": report["parity"], "output": str(output_path)}, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,319 @@
|
||||
"""Production-oriented natural-language memory acceptance benchmark."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from .qwen_integration import (
|
||||
DEFAULT_MEMORY_RESET_TOKEN,
|
||||
load_memory_config,
|
||||
load_qwen_dynamic,
|
||||
load_tokenizer,
|
||||
resolve_memory_reset_token,
|
||||
)
|
||||
|
||||
|
||||
def _chat_tensor(tokenizer, messages, *, add_generation_prompt: bool):
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
add_generation_prompt=add_generation_prompt,
|
||||
return_tensors="pt",
|
||||
return_dict=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
return {
|
||||
key: value
|
||||
for key, value in encoded.items()
|
||||
if isinstance(value, torch.Tensor)
|
||||
}
|
||||
|
||||
|
||||
def _memory_system_prefix(tokenizer, content: str):
|
||||
full = tokenizer.apply_chat_template(
|
||||
[
|
||||
{"role": "system", "content": content},
|
||||
{"role": "user", "content": "__memory_query_boundary__"},
|
||||
],
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
return_tensors="pt",
|
||||
return_dict=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
input_ids = full["input_ids"]
|
||||
im_start = tokenizer.convert_tokens_to_ids("<|im_start|>")
|
||||
positions = (input_ids[0] == int(im_start)).nonzero(as_tuple=False).flatten()
|
||||
if positions.numel() < 2:
|
||||
raise RuntimeError("could not locate the system/user memory boundary")
|
||||
end = int(positions[1].item())
|
||||
return {
|
||||
"input_ids": input_ids[:, :end],
|
||||
"attention_mask": torch.ones((1, end), dtype=torch.long),
|
||||
}
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _generate(model, tokenizer, user_text: str, max_new_tokens: int) -> str:
|
||||
encoded = _chat_tensor(
|
||||
tokenizer,
|
||||
[{"role": "user", "content": user_text}],
|
||||
add_generation_prompt=True,
|
||||
)
|
||||
device = model._find_layer_device()
|
||||
encoded = {key: value.to(device) for key, value in encoded.items()}
|
||||
query = tokenizer(user_text, add_special_tokens=False, return_tensors="pt")
|
||||
query_ids = query["input_ids"].to(device)
|
||||
query_mask = query.get("attention_mask")
|
||||
if query_mask is None:
|
||||
query_mask = torch.ones_like(query_ids)
|
||||
query_mask = query_mask.to(device)
|
||||
output = model.generate(
|
||||
**encoded,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
memory_query_input_ids=query_ids,
|
||||
memory_query_attention_mask=query_mask,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
response_ids = output[0, encoded["input_ids"].shape[1] :]
|
||||
return tokenizer.decode(response_ids.detach().cpu().tolist(), skip_special_tokens=True).strip()
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _write(model, tokenizer, fact: str, acknowledgement: str) -> dict:
|
||||
device = model._find_layer_device()
|
||||
dialogue = [
|
||||
{"role": "user", "content": fact},
|
||||
{"role": "assistant", "content": acknowledgement},
|
||||
]
|
||||
encoded = _chat_tensor(tokenizer, dialogue, add_generation_prompt=False)
|
||||
encoded = {key: value.to(device) for key, value in encoded.items()}
|
||||
text_prefix = _memory_system_prefix(
|
||||
tokenizer,
|
||||
"以下是与当前用户相关的已保存长期记忆。仅在问题相关时使用，不要编造：\n" + fact,
|
||||
)
|
||||
text_ids = text_prefix["input_ids"].to(device)
|
||||
text_mask = text_prefix["attention_mask"].to(device)
|
||||
key = tokenizer(fact, add_special_tokens=False, return_tensors="pt")
|
||||
key_ids = key["input_ids"].to(device)
|
||||
key_mask = key.get("attention_mask")
|
||||
if key_mask is None:
|
||||
key_mask = torch.ones_like(key_ids)
|
||||
key_mask = key_mask.to(device)
|
||||
storage = tokenizer(fact, add_special_tokens=False, return_tensors="pt")
|
||||
storage_ids = storage["input_ids"].to(device)
|
||||
storage_mask = storage.get("attention_mask")
|
||||
if storage_mask is None:
|
||||
storage_mask = torch.ones_like(storage_ids)
|
||||
storage_mask = storage_mask.to(device)
|
||||
model(
|
||||
**encoded,
|
||||
read_memory=False,
|
||||
update_memory=True,
|
||||
return_memory=True,
|
||||
use_cache=False,
|
||||
memory_text_input_ids=text_ids,
|
||||
memory_text_attention_mask=text_mask,
|
||||
memory_key_input_ids=key_ids,
|
||||
memory_key_attention_mask=key_mask,
|
||||
memory_storage_input_ids=storage_ids,
|
||||
memory_storage_attention_mask=storage_mask,
|
||||
)
|
||||
address = model.memory.last_write_address
|
||||
probability = model.memory.last_write_probability
|
||||
stored_slot = model.runtime.text_last_written_slot
|
||||
return {
|
||||
"fact": fact,
|
||||
"write_probability": float(probability.detach().mean()) if probability is not None else None,
|
||||
"selected_slot": int(address.argmax(dim=-1)[0].item()) if address is not None else None,
|
||||
"stored_slot": int(stored_slot[0].item()) if stored_slot is not None else None,
|
||||
"valid_slots_after_write": int(model.runtime.text_slot_valid.sum())
|
||||
if model.runtime.text_slot_valid is not None
|
||||
else 0,
|
||||
}
|
||||
|
||||
|
||||
def _load_persistent_checkpoint(model_path: str, adapter_path: str, *, no_4bit: bool):
|
||||
config = load_memory_config(adapter_path)
|
||||
restarted = load_qwen_dynamic(
|
||||
model_path,
|
||||
memory_config=config,
|
||||
load_in_4bit=not no_4bit,
|
||||
)
|
||||
restarted.load_memory_adapter(adapter_path)
|
||||
restarted.eval()
|
||||
return restarted, config
|
||||
|
||||
|
||||
def main() -> None:
|
||||
if hasattr(sys.stdout, "reconfigure"):
|
||||
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-path", default=".")
|
||||
parser.add_argument("--adapter", default="dynamic_memory_lab/qwen_memory_adapter_native_v3")
|
||||
parser.add_argument(
|
||||
"--output-adapter",
|
||||
default="dynamic_memory_lab/qwen_memory_adapter_natural_production_v1",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--report",
|
||||
default="dynamic_memory_lab/benchmark_natural_language_memory.json",
|
||||
)
|
||||
parser.add_argument("--max-new-tokens", type=int, default=48)
|
||||
parser.add_argument("--text-memory-threshold", type=float, default=0.35)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
tokenizer = load_tokenizer(args.model_path)
|
||||
config = load_memory_config(args.adapter)
|
||||
config.persistent_memory = True
|
||||
config.natural_language_memory = True
|
||||
config.text_memory_threshold = args.text_memory_threshold
|
||||
config.reset_token_id = resolve_memory_reset_token(tokenizer, DEFAULT_MEMORY_RESET_TOKEN)
|
||||
model = load_qwen_dynamic(
|
||||
args.model_path,
|
||||
memory_config=config,
|
||||
load_in_4bit=not args.no_4bit,
|
||||
)
|
||||
model.load_memory_adapter(args.adapter)
|
||||
model.reset_memory()
|
||||
|
||||
writes = [
|
||||
_write(model, tokenizer, "请记住：我的工作地点代号是R7。", "好的，我会记住。"),
|
||||
_write(model, tokenizer, "请记住：我最喜欢的水果是红富士苹果。", "好的，我会记住。"),
|
||||
_write(model, tokenizer, "更新一下：我的工作地点代号改为K9。", "好的，已更新。"),
|
||||
]
|
||||
model.save_persistent_memory_checkpoint(args.output_adapter)
|
||||
saved_valid_slots = int(model.runtime.text_slot_valid.sum())
|
||||
saved_norm = float(model.runtime.state.detach().float().norm())
|
||||
|
||||
del model
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
restarted, restart_config = _load_persistent_checkpoint(
|
||||
args.model_path,
|
||||
args.output_adapter,
|
||||
no_4bit=args.no_4bit,
|
||||
)
|
||||
loaded_norm = float(restarted.runtime.state.detach().float().norm())
|
||||
|
||||
queries = [
|
||||
{"name": "replaced_work_code", "query": "我的工作地点代号是什么？", "expected": "K9"},
|
||||
{"name": "favorite_fruit", "query": "我最喜欢吃什么水果？", "expected": "红富士苹果"},
|
||||
{"name": "unknown_blood_type", "query": "我的血型是什么？如果没有记录，请明确说不知道。", "expected": "不知道"},
|
||||
]
|
||||
query_results = []
|
||||
for item in queries:
|
||||
response = _generate(restarted, tokenizer, item["query"], args.max_new_tokens)
|
||||
relevance = restarted.runtime.text_read_relevance
|
||||
overlap = restarted.runtime.text_read_overlap
|
||||
selected_slots = restarted.runtime.text_read_slots
|
||||
query_results.append(
|
||||
{
|
||||
**item,
|
||||
"response": response,
|
||||
"text_prefix_used": restarted.runtime.text_prefix_used,
|
||||
"retrieval_relevance": float(relevance[0].item()) if relevance is not None else None,
|
||||
"retrieval_overlap": overlap[0].detach().cpu().tolist()
|
||||
if overlap is not None
|
||||
else None,
|
||||
"retrieved_slots": selected_slots[0].detach().cpu().tolist()
|
||||
if selected_slots is not None
|
||||
else None,
|
||||
"expected_found": item["expected"] in response,
|
||||
"refused_unknown": item["name"] != "unknown_blood_type"
|
||||
or any(marker in response for marker in ("不知道", "没有记录", "无相关", "未找到", "不清楚")),
|
||||
}
|
||||
)
|
||||
|
||||
del restarted
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
restarted, second_restart_config = _load_persistent_checkpoint(
|
||||
args.model_path,
|
||||
args.output_adapter,
|
||||
no_4bit=args.no_4bit,
|
||||
)
|
||||
second_loaded_norm = float(restarted.runtime.state.detach().float().norm())
|
||||
second_restart_results = []
|
||||
for item in queries[:2]:
|
||||
response = _generate(restarted, tokenizer, item["query"], args.max_new_tokens)
|
||||
second_restart_results.append(
|
||||
{
|
||||
"name": item["name"],
|
||||
"response": response,
|
||||
"expected": item["expected"],
|
||||
"expected_found": item["expected"] in response,
|
||||
"text_prefix_used": restarted.runtime.text_prefix_used,
|
||||
}
|
||||
)
|
||||
|
||||
reset_inputs = _chat_tensor(
|
||||
tokenizer,
|
||||
[{"role": "user", "content": DEFAULT_MEMORY_RESET_TOKEN}],
|
||||
add_generation_prompt=True,
|
||||
)
|
||||
device = restarted._find_layer_device()
|
||||
reset_inputs = {key: value.to(device) for key, value in reset_inputs.items()}
|
||||
restarted.generate(
|
||||
**reset_inputs,
|
||||
max_new_tokens=1,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
use_cache=False,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
reset_norm = float(restarted.runtime.state.detach().float().norm())
|
||||
reset_valid_slots = int(restarted.runtime.text_slot_valid.sum())
|
||||
after_reset = _generate(restarted, tokenizer, "我的工作地点代号是什么？", args.max_new_tokens)
|
||||
|
||||
report = {
|
||||
"writes": writes,
|
||||
"history_passed_to_restart": False,
|
||||
"saved_valid_slots": saved_valid_slots,
|
||||
"saved_memory_norm": saved_norm,
|
||||
"loaded_memory_norm": loaded_norm,
|
||||
"restart_state_equal_norm": abs(saved_norm - loaded_norm) < 1e-5,
|
||||
"second_loaded_memory_norm": second_loaded_norm,
|
||||
"second_restart_state_equal_norm": abs(saved_norm - second_loaded_norm) < 1e-5,
|
||||
"queries": query_results,
|
||||
"second_restart_queries": second_restart_results,
|
||||
"second_restart_pass": all(row["expected_found"] for row in second_restart_results),
|
||||
"all_known_queries_pass": all(row["expected_found"] for row in query_results[:2]),
|
||||
"unknown_refusal_pass": query_results[2]["refused_unknown"],
|
||||
"reset_token": DEFAULT_MEMORY_RESET_TOKEN,
|
||||
"reset_token_id": restart_config.reset_token_id,
|
||||
"reset_memory_norm": reset_norm,
|
||||
"reset_valid_slots": reset_valid_slots,
|
||||
"reset_cleared_pass": reset_norm < 1e-5 and reset_valid_slots == 0,
|
||||
"response_after_reset": after_reset,
|
||||
}
|
||||
report["production_gate_pass"] = bool(
|
||||
report["history_passed_to_restart"] is False
|
||||
and report["restart_state_equal_norm"]
|
||||
and report["second_restart_state_equal_norm"]
|
||||
and report["all_known_queries_pass"]
|
||||
and report["second_restart_pass"]
|
||||
and report["unknown_refusal_pass"]
|
||||
and report["reset_cleared_pass"]
|
||||
)
|
||||
report_path = Path(args.report)
|
||||
report_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(json.dumps(report, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,725 @@
|
||||
"""Full local comparison between Qwen3.5-4B and Natural Memory v1.
|
||||
|
||||
This is an engineering benchmark, not a claim of state-of-the-art performance.
|
||||
It evaluates the same frozen Qwen3.5-4B backbone with and without the internal
|
||||
memory path, using deterministic greedy decoding and locally generated cases.
|
||||
The report includes general ability, extra math and reasoning cases, long
|
||||
context retrieval, throughput, latency, VRAM, automatic write decisions,
|
||||
conflict replacement, unknown-fact refusal, reset, and shard-backed restart.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
import statistics
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from .qwen_integration import (
|
||||
load_memory_config,
|
||||
load_qwen_base,
|
||||
load_qwen_dynamic,
|
||||
load_tokenizer,
|
||||
)
|
||||
from .stream_chat_qwen_memory import _chat_tensor, _memory_system_prefix, _write_turn
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--base-model", default=r"W:\Flash\model")
|
||||
parser.add_argument(
|
||||
"--memory-model",
|
||||
default=r"W:\Flash\model\dynamic_memory_lab\qwen3_5_4b_memory_merged_v13",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--data",
|
||||
default=r"W:\Flash\model\dynamic_memory_lab\data\comprehensive_general.jsonl",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
default=r"W:\Flash\model\dynamic_memory_lab\natural_memory_v1_full_benchmark.json",
|
||||
)
|
||||
parser.add_argument("--max-new-tokens", type=int, default=24)
|
||||
parser.add_argument("--perf-repeats", type=int, default=3)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def normalize(text: str) -> str:
|
||||
keep = str(text).lower()
|
||||
for char in " \t\r\n`*_#，。！？、；：,.!?;:'\"（）()[]{}<>|\\/：":
|
||||
keep = keep.replace(char, "")
|
||||
return keep
|
||||
|
||||
|
||||
def contains_answer(text: str, acceptable: list[str]) -> bool:
|
||||
normalized = normalize(text)
|
||||
for answer in acceptable:
|
||||
expected = normalize(str(answer))
|
||||
if not expected:
|
||||
continue
|
||||
if expected.isdigit() and len(expected) == 1:
|
||||
if any(
|
||||
normalized[index : index + 1] == expected
|
||||
and (index == 0 or not normalized[index - 1].isdigit())
|
||||
and (index + 1 == len(normalized) or not normalized[index + 1].isdigit())
|
||||
for index in range(len(normalized))
|
||||
):
|
||||
return True
|
||||
elif expected in normalized:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def load_jsonl(path: str | Path) -> list[dict[str, Any]]:
|
||||
rows = []
|
||||
for line in Path(path).read_text(encoding="utf-8").splitlines():
|
||||
if line.strip():
|
||||
rows.append(json.loads(line))
|
||||
return rows
|
||||
|
||||
|
||||
def prompt_inputs(tokenizer: Any, prompt: str, device: torch.device) -> dict[str, torch.Tensor]:
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
[{"role": "user", "content": prompt}],
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
return_tensors="pt",
|
||||
return_dict=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
return {
|
||||
key: value.to(device)
|
||||
for key, value in encoded.items()
|
||||
if isinstance(value, torch.Tensor)
|
||||
}
|
||||
|
||||
|
||||
def input_token_count(tokenizer: Any, prompt: str) -> int:
|
||||
encoded = tokenizer(prompt, add_special_tokens=False)
|
||||
return len(encoded["input_ids"])
|
||||
|
||||
|
||||
def generate_answer(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
prompt: str,
|
||||
*,
|
||||
dynamic: bool,
|
||||
max_new_tokens: int,
|
||||
) -> tuple[str, int, float]:
|
||||
device = model._find_layer_device() if dynamic else model.get_input_embeddings().weight.device
|
||||
encoded = prompt_inputs(tokenizer, prompt, device)
|
||||
query = tokenizer(prompt, add_special_tokens=False, return_tensors="pt")
|
||||
query_ids = query["input_ids"].to(device)
|
||||
query_mask = query.get("attention_mask")
|
||||
if query_mask is None:
|
||||
query_mask = torch.ones_like(query_ids)
|
||||
query_mask = query_mask.to(device)
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize(device)
|
||||
started = time.perf_counter()
|
||||
with torch.inference_mode():
|
||||
kwargs: dict[str, Any] = {
|
||||
"max_new_tokens": max_new_tokens,
|
||||
"do_sample": False,
|
||||
"use_cache": True,
|
||||
"pad_token_id": tokenizer.pad_token_id,
|
||||
}
|
||||
if dynamic:
|
||||
kwargs.update(
|
||||
{
|
||||
"update_memory": False,
|
||||
"memory_query_input_ids": query_ids,
|
||||
"memory_query_attention_mask": query_mask,
|
||||
}
|
||||
)
|
||||
output = model.generate(**encoded, **kwargs)
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize(device)
|
||||
elapsed = time.perf_counter() - started
|
||||
prompt_len = int(encoded["input_ids"].shape[1])
|
||||
response_ids = output[0, prompt_len:]
|
||||
response = tokenizer.decode(response_ids.detach().cpu().tolist(), skip_special_tokens=True).strip()
|
||||
return response, int(response_ids.numel()), elapsed
|
||||
|
||||
|
||||
def make_math_cases() -> list[dict[str, Any]]:
|
||||
raw = [
|
||||
("m01", "只输出最终整数：38 + 47 = ?", ["85"]),
|
||||
("m02", "只输出最终整数：900 - 376 = ?", ["524"]),
|
||||
("m03", "只输出最终整数：24 × 17 = ?", ["408"]),
|
||||
("m04", "只输出最终整数：936 ÷ 18 = ?", ["52"]),
|
||||
("m05", "只输出结果：2.75 + 3.6 = ?", ["6.35"]),
|
||||
("m06", "只输出结果：7/8 - 1/4 = ?", ["5/8", "0.625"]),
|
||||
("m07", "只输出整数：15 和 28 的最小公倍数是多少？", ["420"]),
|
||||
("m08", "只输出百分数：480 的 12.5% 是多少？", ["60"]),
|
||||
("m09", "只输出百分数：80 增长到 100，增长率是多少？", ["25%", "25"]),
|
||||
("m10", "只输出结果：3 的 5 次方是多少？", ["243"]),
|
||||
("m11", "只输出 x：4x + 7 = 31。", ["6"]),
|
||||
("m12", "只输出 x：9x - 18 = 0。", ["2"]),
|
||||
("m13", "只输出 x：2(x + 5) = 18。", ["4"]),
|
||||
("m14", "只输出 x：x/3 + 4 = 9。", ["15"]),
|
||||
("m15", "只输出 x：5x - 2 = 3x + 10。", ["6"]),
|
||||
("m16", "只输出下一个数：5，10，20，40，？", ["80"]),
|
||||
("m17", "只输出下一个数：3，6，11，18，27，？", ["38"]),
|
||||
("m18", "只输出下一个数：1，4，9，16，？", ["25"]),
|
||||
("m19", "只输出下一个数：2，3，5，8，12，？", ["17"]),
|
||||
("m20", "只输出整数：阶乘 6! 等于多少？", ["720"]),
|
||||
("m21", "只输出面积：长 12、宽 7 的矩形面积是多少？", ["84"]),
|
||||
("m22", "只输出周长：边长为 9 的正方形周长是多少？", ["36"]),
|
||||
("m23", "只输出面积：底 10、高 6 的三角形面积是多少？", ["30"]),
|
||||
("m24", "只输出角度：一个三角形两个角是 35 度和 65 度，第三个角是多少？", ["80"]),
|
||||
("m25", "只输出数量：3 件不同衬衫和 2 条不同裤子可以组成多少套穿搭？", ["6"]),
|
||||
("m26", "只输出数量：从 5 个人中选 2 个人，有多少种选法？", ["10"]),
|
||||
("m27", "只输出余数：17 除以 5 的余数是多少？", ["2"]),
|
||||
("m28", "只输出结果：平均数 8、12、16、20 是多少？", ["14"]),
|
||||
("m29", "只输出结果：一个商品原价 240 元，打八折后多少钱？", ["192"]),
|
||||
("m30", "只输出结果：2.4 × 0.5 = ?", ["1.2"]),
|
||||
]
|
||||
return [
|
||||
{"id": case_id, "category": "math", "prompt": prompt, "acceptable": answers}
|
||||
for case_id, prompt, answers in raw
|
||||
]
|
||||
|
||||
|
||||
def make_reasoning_cases() -> list[dict[str, Any]]:
|
||||
raw = [
|
||||
("r01", "只输出名字：甲比乙早到，乙比丙早到，谁最后到？", ["丙"]),
|
||||
("r02", "只输出名字：小李在小王左边，小王在小张左边，谁最右边？", ["小张"]),
|
||||
("r03", "只输出结论：所有鸟都有翅膀，企鹅是鸟，所以企鹅有翅膀吗？", ["是"]),
|
||||
("r04", "只输出结论：所有猫都是哺乳动物，鲸鱼是哺乳动物，所以鲸鱼是猫吗？", ["不是", "否"]),
|
||||
("r05", "只输出结论：如果下雨就带伞。现在下雨了，要不要带伞？", ["要"]),
|
||||
("r06", "只输出结论：只有持票者才能入场。小林没有票，他能入场吗？", ["不能", "不可以"]),
|
||||
("r07", "只输出星期：今天是星期三，五天后是星期几？", ["星期一", "周一"]),
|
||||
("r08", "只输出方向：你面向北，右转后面向哪个方向？", ["东"]),
|
||||
("r09", "只输出方向：你面向东，左转后面向哪个方向？", ["北"]),
|
||||
("r10", "只输出数量：盒子里有 4 个红球和 3 个蓝球，不看颜色拿出一个，至少有几个球？", ["1"]),
|
||||
("r11", "只输出名字：甲不是第一，乙在甲前面，丙在乙后面，谁可能是第一？", ["乙"]),
|
||||
("r12", "只输出结论：有些学生会游泳，小周是学生，能确定小周会游泳吗？", ["不能", "无法"]),
|
||||
("r13", "只输出下一个数：1，2，4，7，11，？", ["16"]),
|
||||
("r14", "只输出下一个数：81，27，9，3，？", ["1"]),
|
||||
("r15", "只输出名字：红色比蓝色重，绿色比红色轻但比蓝色重，哪个最轻？", ["蓝色"]),
|
||||
("r16", "只输出答案：苹果不是蔬菜，胡萝卜是蔬菜，香蕉是水果，哪个不是水果？", ["胡萝卜"]),
|
||||
("r17", "只输出结论：如果 A 大于 B 且 B 大于 C，那么 A 大于 C 吗？", ["是"]),
|
||||
("r18", "只输出结论：如果一个数能被 2 整除，它一定是偶数。14 能被 2 整除，它是偶数吗？", ["是"]),
|
||||
("r19", "只输出名字：小赵比小钱高，小孙比小赵矮但比小钱高，谁最高？", ["小赵"]),
|
||||
("r20", "只输出数量：一周中有几天的名字包含‘星’字？", ["7"]),
|
||||
("r21", "只输出结论：没有鱼是鸟，金鱼是鱼，所以金鱼是鸟吗？", ["不是", "否"]),
|
||||
("r22", "只输出顺序：春、夏、秋、冬之后又回到哪个季节？", ["春"]),
|
||||
("r23", "只输出结论：所有密码都需要保密，这个字符串是密码，所以它需要保密吗？", ["是"]),
|
||||
("r24", "只输出答案：小明有两个兄弟，每个兄弟都有一个姐姐，小明有几个姐姐？", ["1"]),
|
||||
]
|
||||
return [
|
||||
{"id": case_id, "category": "reasoning", "prompt": prompt, "acceptable": answers}
|
||||
for case_id, prompt, answers in raw
|
||||
]
|
||||
|
||||
|
||||
def make_context_cases(tokenizer: Any) -> list[dict[str, Any]]:
|
||||
rng = random.Random(20260904)
|
||||
filler = (
|
||||
"这是一段与问题无关的背景说明。系统记录了版本号、构建时间、测试批次、"
|
||||
"设备温度、日志摘要和普通项目备注。这些文字只是干扰项，不包含目标答案。"
|
||||
)
|
||||
cases: list[dict[str, Any]] = []
|
||||
for target_tokens in (512, 2048, 4096, 8192):
|
||||
for position in ("early", "middle", "late"):
|
||||
answer = f"CTX{target_tokens}-{position.upper()}-{rng.randrange(100, 999)}"
|
||||
needle = f"目标记录：本次检索需要返回的项目编码是 {answer}。"
|
||||
chunks: list[str] = []
|
||||
while input_token_count(tokenizer, " ".join(chunks + [filler, needle])) < target_tokens:
|
||||
chunks.append(filler)
|
||||
if position == "early":
|
||||
material = " ".join([needle] + chunks)
|
||||
elif position == "middle":
|
||||
half = len(chunks) // 2
|
||||
material = " ".join(chunks[:half] + [needle] + chunks[half:])
|
||||
else:
|
||||
material = " ".join(chunks + [needle])
|
||||
prompt = (
|
||||
"请阅读下面的材料，只输出目标记录中的项目编码，不要解释。\n"
|
||||
"---材料开始---\n"
|
||||
f"{material}\n"
|
||||
"---材料结束---\n"
|
||||
"问题：目标记录中的项目编码是什么？"
|
||||
)
|
||||
cases.append(
|
||||
{
|
||||
"id": f"ctx-{target_tokens}-{position}",
|
||||
"category": f"context_{target_tokens}",
|
||||
"prompt": prompt,
|
||||
"acceptable": [answer],
|
||||
"target_tokens": target_tokens,
|
||||
"position": position,
|
||||
}
|
||||
)
|
||||
return cases
|
||||
|
||||
|
||||
def evaluate_cases(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
cases: list[dict[str, Any]],
|
||||
*,
|
||||
dynamic: bool,
|
||||
max_new_tokens: int,
|
||||
) -> dict[str, Any]:
|
||||
rows: list[dict[str, Any]] = []
|
||||
categories: dict[str, list[float]] = {}
|
||||
started = time.perf_counter()
|
||||
for case in cases:
|
||||
if dynamic:
|
||||
model.reset_memory(batch_size=1, device=model._find_layer_device())
|
||||
response, generated_tokens, elapsed = generate_answer(
|
||||
model,
|
||||
tokenizer,
|
||||
str(case["prompt"]),
|
||||
dynamic=dynamic,
|
||||
max_new_tokens=max_new_tokens,
|
||||
)
|
||||
passed = contains_answer(response, list(case["acceptable"]))
|
||||
category = str(case["category"])
|
||||
categories.setdefault(category, []).append(float(passed))
|
||||
rows.append(
|
||||
{
|
||||
"id": case["id"],
|
||||
"category": category,
|
||||
"prompt_tokens": input_token_count(tokenizer, str(case["prompt"])),
|
||||
"acceptable": case["acceptable"],
|
||||
"generated": response,
|
||||
"generated_tokens": generated_tokens,
|
||||
"seconds": elapsed,
|
||||
"passed": passed,
|
||||
}
|
||||
)
|
||||
total = sum(sum(values) for values in categories.values())
|
||||
return {
|
||||
"cases": len(rows),
|
||||
"elapsed_seconds": time.perf_counter() - started,
|
||||
"overall_score": total / max(1, len(rows)),
|
||||
"categories": {
|
||||
category: {
|
||||
"count": len(values),
|
||||
"score": sum(values) / max(1, len(values)),
|
||||
}
|
||||
for category, values in sorted(categories.items())
|
||||
},
|
||||
"rows": rows,
|
||||
}
|
||||
|
||||
|
||||
def device_snapshot(model: Any) -> dict[str, Any]:
|
||||
device = model._find_layer_device() if hasattr(model, "_find_layer_device") else model.get_input_embeddings().weight.device
|
||||
params = sum(parameter.numel() for parameter in model.parameters())
|
||||
result: dict[str, Any] = {
|
||||
"device": str(device),
|
||||
"parameter_count": int(params),
|
||||
"parameter_count_billion": params / 1e9,
|
||||
}
|
||||
if device.type == "cuda":
|
||||
properties = torch.cuda.get_device_properties(device)
|
||||
result.update(
|
||||
{
|
||||
"gpu_name": properties.name,
|
||||
"gpu_total_memory_gb": properties.total_memory / 1024**3,
|
||||
}
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def measure_performance(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
cases: list[dict[str, Any]],
|
||||
*,
|
||||
dynamic: bool,
|
||||
max_new_tokens: int,
|
||||
repeats: int,
|
||||
) -> dict[str, Any]:
|
||||
selected = [cases[0]]
|
||||
for wanted in (512, 2048, 4096):
|
||||
matching = [case for case in cases if case.get("target_tokens") == wanted]
|
||||
if matching:
|
||||
selected.append(matching[1])
|
||||
rows = []
|
||||
device = model._find_layer_device() if dynamic else model.get_input_embeddings().weight.device
|
||||
for case in selected:
|
||||
latencies = []
|
||||
generated = 0
|
||||
for _ in range(max(1, repeats)):
|
||||
if dynamic:
|
||||
model.reset_memory(batch_size=1, device=device)
|
||||
_, token_count, elapsed = generate_answer(
|
||||
model,
|
||||
tokenizer,
|
||||
str(case["prompt"]),
|
||||
dynamic=dynamic,
|
||||
max_new_tokens=max_new_tokens,
|
||||
)
|
||||
latencies.append(elapsed)
|
||||
generated += token_count
|
||||
rows.append(
|
||||
{
|
||||
"id": case["id"],
|
||||
"prompt_tokens": input_token_count(tokenizer, str(case["prompt"])),
|
||||
"median_seconds": statistics.median(latencies),
|
||||
"mean_seconds": statistics.mean(latencies),
|
||||
"tokens_per_second": generated / max(1e-9, sum(latencies)),
|
||||
"repeats": len(latencies),
|
||||
}
|
||||
)
|
||||
return {"rows": rows}
|
||||
|
||||
|
||||
def memory_payload(model: Any) -> dict[str, torch.Tensor]:
|
||||
payload: dict[str, torch.Tensor] = {
|
||||
"memory_state": model.runtime.state.detach().cpu().clone(),
|
||||
}
|
||||
if model.memory_config.natural_language_memory:
|
||||
for name in (
|
||||
"text_token_ids",
|
||||
"text_token_mask",
|
||||
"text_slot_valid",
|
||||
"text_slot_keys",
|
||||
"text_slot_age",
|
||||
"text_write_counter",
|
||||
"text_key_token_ids",
|
||||
"text_key_token_mask",
|
||||
):
|
||||
value = getattr(model.runtime, name)
|
||||
if not isinstance(value, torch.Tensor):
|
||||
raise RuntimeError(f"runtime memory field is unavailable: {name}")
|
||||
payload[name] = value.detach().cpu().clone()
|
||||
return payload
|
||||
|
||||
|
||||
def answer_memory_query(model: Any, tokenizer: Any, text: str, max_new_tokens: int) -> tuple[str, bool]:
|
||||
response, _, _ = generate_answer(
|
||||
model,
|
||||
tokenizer,
|
||||
text,
|
||||
dynamic=True,
|
||||
max_new_tokens=max_new_tokens,
|
||||
)
|
||||
return response, bool(model.runtime.text_prefix_used)
|
||||
|
||||
|
||||
def run_memory_benchmark(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
package_path: str | Path,
|
||||
*,
|
||||
max_new_tokens: int,
|
||||
no_4bit: bool,
|
||||
) -> dict[str, Any]:
|
||||
"""Run memory tests and restore the user's original embedded state."""
|
||||
|
||||
original = memory_payload(model)
|
||||
decision_rows = []
|
||||
positives = [
|
||||
"我叫林舟。",
|
||||
"我的常住城市是苏州。",
|
||||
"我最喜欢的水果是红富士苹果。",
|
||||
"我正在开发 Natural Memory v1 项目。",
|
||||
"以后请把代码默认写成 Python。",
|
||||
"我的常用时区是 Asia/Shanghai。",
|
||||
"这是我的长期偏好：使用简洁中文。",
|
||||
"这个项目的重要约束是不要修改原始 Qwen 权重。",
|
||||
]
|
||||
negatives = [
|
||||
"我叫什么？",
|
||||
"帮我解释向量数据库是什么。",
|
||||
"你觉得今天的天气怎么样？",
|
||||
"请把 memory 翻译成中文。",
|
||||
"计算一下 17 × 19。",
|
||||
"如果我选择 GPU，会发生什么？",
|
||||
"我之前有没有提到我的城市？",
|
||||
"给我一个自然语言记忆方案。",
|
||||
]
|
||||
try:
|
||||
for expected, items in ((True, positives), (False, negatives)):
|
||||
for text in items:
|
||||
model.reset_memory(batch_size=1, device=model._find_layer_device())
|
||||
changed = _write_turn(model, tokenizer, text, model._find_layer_device())
|
||||
decision_rows.append(
|
||||
{
|
||||
"text": text,
|
||||
"expected_write": expected,
|
||||
"actual_write": changed,
|
||||
"write_probability": float(model.runtime.auto_memory_probability.mean())
|
||||
if isinstance(model.runtime.auto_memory_probability, torch.Tensor)
|
||||
else None,
|
||||
}
|
||||
)
|
||||
tp = sum(row["expected_write"] and row["actual_write"] for row in decision_rows)
|
||||
tn = sum((not row["expected_write"]) and (not row["actual_write"]) for row in decision_rows)
|
||||
fp = sum((not row["expected_write"]) and row["actual_write"] for row in decision_rows)
|
||||
fn = sum(row["expected_write"] and (not row["actual_write"]) for row in decision_rows)
|
||||
|
||||
model.reset_memory(batch_size=1, device=model._find_layer_device())
|
||||
first = "我的工作地点代号是NM-R7。"
|
||||
second = "我最喜欢的水果是青提。"
|
||||
replacement = "我的工作地点代号改为NM-K9。"
|
||||
writes = [
|
||||
{"text": first, "changed": _write_turn(model, tokenizer, first, model._find_layer_device())},
|
||||
{"text": second, "changed": _write_turn(model, tokenizer, second, model._find_layer_device())},
|
||||
{
|
||||
"text": replacement,
|
||||
"changed": _write_turn(model, tokenizer, replacement, model._find_layer_device()),
|
||||
},
|
||||
]
|
||||
before_restart = {}
|
||||
for name, query, expected in (
|
||||
("work_code", "我的工作地点代号是什么？", "NM-K9"),
|
||||
("fruit", "我最喜欢的水果是什么？", "青提"),
|
||||
("unknown", "我的血型是什么？如果没有记录，请明确说不知道。", "不知道"),
|
||||
):
|
||||
response, prefix_used = answer_memory_query(model, tokenizer, query, max_new_tokens)
|
||||
before_restart[name] = {
|
||||
"query": query,
|
||||
"expected": expected,
|
||||
"response": response,
|
||||
"expected_found": expected in response,
|
||||
"prefix_used": prefix_used,
|
||||
}
|
||||
|
||||
# Persist through ordinary natural-language turns. The caller releases
|
||||
# this model before loading a fresh process/model for the restart test;
|
||||
# keeping that lifecycle outside this function avoids two 4-bit Qwen
|
||||
# backbones occupying the GPU at the same time.
|
||||
model.save_embedded_memory_weights(package_path)
|
||||
return {
|
||||
"automatic_write_decision": {
|
||||
"rows": decision_rows,
|
||||
"true_positive": int(tp),
|
||||
"true_negative": int(tn),
|
||||
"false_positive": int(fp),
|
||||
"false_negative": int(fn),
|
||||
"precision": tp / max(1, tp + fp),
|
||||
"recall": tp / max(1, tp + fn),
|
||||
"specificity": tn / max(1, tn + fp),
|
||||
},
|
||||
"natural_language_writes": writes,
|
||||
"before_restart": before_restart,
|
||||
"embedded_write_persisted": True,
|
||||
}
|
||||
except Exception:
|
||||
# Best-effort restoration if a test fails halfway through.
|
||||
model._load_persistent_memory_payload(original)
|
||||
model.save_embedded_memory_weights(package_path)
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
raise
|
||||
|
||||
|
||||
def release(model: Any) -> None:
|
||||
del model
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
tokenizer = load_tokenizer(args.base_model)
|
||||
memory_config_hint = load_memory_config(args.memory_model)
|
||||
memory_variant = (
|
||||
"natural_memory_v2"
|
||||
if memory_config_hint.memory_version >= 2 or memory_config_hint.hierarchical_memory
|
||||
else "natural_memory_v1"
|
||||
)
|
||||
general = load_jsonl(args.data)
|
||||
math_cases = make_math_cases()
|
||||
reasoning_cases = make_reasoning_cases()
|
||||
context_cases = make_context_cases(tokenizer)
|
||||
all_cases = general + math_cases + reasoning_cases + context_cases
|
||||
use_4bit = not args.no_4bit
|
||||
report: dict[str, Any] = {
|
||||
"benchmark": f"Natural Memory {memory_variant.rsplit('_', 1)[-1]} full local comparison",
|
||||
"date": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"seed": 20260904,
|
||||
"base_model": str(Path(args.base_model).resolve()),
|
||||
"memory_model": str(Path(args.memory_model).resolve()),
|
||||
"data": str(Path(args.data).resolve()),
|
||||
"quantization": "4bit_nf4" if use_4bit else "none",
|
||||
"decoding": {"do_sample": False, "max_new_tokens": args.max_new_tokens},
|
||||
"case_counts": {
|
||||
"general_existing": len(general),
|
||||
"math": len(math_cases),
|
||||
"reasoning": len(reasoning_cases),
|
||||
"context": len(context_cases),
|
||||
"total": len(all_cases),
|
||||
},
|
||||
"context_targets": sorted({case["target_tokens"] for case in context_cases}),
|
||||
}
|
||||
|
||||
print(f"cases={len(all_cases)} quantization={report['quantization']}")
|
||||
print("loading Qwen3.5-4B baseline")
|
||||
started = time.perf_counter()
|
||||
baseline = load_qwen_base(args.base_model, load_in_4bit=use_4bit)
|
||||
baseline.eval()
|
||||
report["baseline"] = {
|
||||
"load_seconds": time.perf_counter() - started,
|
||||
"hardware": device_snapshot(baseline),
|
||||
}
|
||||
baseline_device = baseline.get_input_embeddings().weight.device
|
||||
if baseline_device.type == "cuda":
|
||||
torch.cuda.reset_peak_memory_stats(baseline_device)
|
||||
report["baseline"]["quality"] = evaluate_cases(
|
||||
baseline,
|
||||
tokenizer,
|
||||
all_cases,
|
||||
dynamic=False,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
)
|
||||
report["baseline"]["performance"] = measure_performance(
|
||||
baseline,
|
||||
tokenizer,
|
||||
context_cases,
|
||||
dynamic=False,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
repeats=args.perf_repeats,
|
||||
)
|
||||
if baseline_device.type == "cuda":
|
||||
report["baseline"]["peak_memory_allocated_gb"] = torch.cuda.max_memory_allocated(baseline_device) / 1024**3
|
||||
report["baseline"]["peak_memory_reserved_gb"] = torch.cuda.max_memory_reserved(baseline_device) / 1024**3
|
||||
release(baseline)
|
||||
|
||||
print(f"loading {memory_variant} embedded package")
|
||||
started = time.perf_counter()
|
||||
dynamic = load_qwen_dynamic(args.memory_model, load_in_4bit=use_4bit)
|
||||
dynamic.eval()
|
||||
report[memory_variant] = {
|
||||
"load_seconds": time.perf_counter() - started,
|
||||
"hardware": device_snapshot(dynamic),
|
||||
}
|
||||
dynamic_device = dynamic._find_layer_device()
|
||||
if dynamic_device.type == "cuda":
|
||||
torch.cuda.reset_peak_memory_stats(dynamic_device)
|
||||
report[memory_variant]["quality"] = evaluate_cases(
|
||||
dynamic,
|
||||
tokenizer,
|
||||
all_cases,
|
||||
dynamic=True,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
)
|
||||
report[memory_variant]["performance"] = measure_performance(
|
||||
dynamic,
|
||||
tokenizer,
|
||||
context_cases,
|
||||
dynamic=True,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
repeats=args.perf_repeats,
|
||||
)
|
||||
if dynamic_device.type == "cuda":
|
||||
report[memory_variant]["peak_memory_allocated_gb"] = torch.cuda.max_memory_allocated(dynamic_device) / 1024**3
|
||||
report[memory_variant]["peak_memory_reserved_gb"] = torch.cuda.max_memory_reserved(dynamic_device) / 1024**3
|
||||
|
||||
original_embedded_payload = memory_payload(dynamic)
|
||||
memory_report = run_memory_benchmark(
|
||||
dynamic,
|
||||
tokenizer,
|
||||
args.memory_model,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
no_4bit=args.no_4bit,
|
||||
)
|
||||
# Release the first Qwen backbone before constructing the fresh model used
|
||||
# by the shard-backed restart check. This is important on a 12 GB GPU.
|
||||
release(dynamic)
|
||||
dynamic = None
|
||||
restarted = None
|
||||
try:
|
||||
restarted = load_qwen_dynamic(args.memory_model, load_in_4bit=use_4bit)
|
||||
restarted.eval()
|
||||
after_restart = {}
|
||||
for name, query, expected in (
|
||||
("work_code", "我的工作地点代号是什么？", "NM-K9"),
|
||||
("fruit", "我最喜欢的水果是什么？", "青提"),
|
||||
):
|
||||
response, prefix_used = answer_memory_query(restarted, tokenizer, query, args.max_new_tokens)
|
||||
after_restart[name] = {
|
||||
"query": query,
|
||||
"expected": expected,
|
||||
"response": response,
|
||||
"expected_found": expected in response,
|
||||
"prefix_used": prefix_used,
|
||||
}
|
||||
restarted.reset_memory(batch_size=1, device=restarted._find_layer_device())
|
||||
reset_response, reset_prefix = answer_memory_query(
|
||||
restarted,
|
||||
tokenizer,
|
||||
"我的工作地点代号是什么？",
|
||||
args.max_new_tokens,
|
||||
)
|
||||
reset_slots = int(restarted.runtime.text_slot_valid.sum().item())
|
||||
memory_report.update(
|
||||
{
|
||||
"after_restart_without_history_or_pt": after_restart,
|
||||
"restart_pass": all(item["expected_found"] for item in after_restart.values()),
|
||||
"reset": {
|
||||
"response": reset_response,
|
||||
"prefix_used": reset_prefix,
|
||||
"valid_slots": reset_slots,
|
||||
"cleared": reset_slots == 0 and not reset_prefix,
|
||||
},
|
||||
}
|
||||
)
|
||||
finally:
|
||||
if restarted is not None:
|
||||
# Restore the user's pre-benchmark state, so the benchmark itself
|
||||
# does not overwrite the active embedded memory snapshot.
|
||||
restarted._load_persistent_memory_payload(original_embedded_payload)
|
||||
restarted.save_embedded_memory_weights(args.memory_model)
|
||||
memory_report["state_restored"] = int(restarted.runtime.text_slot_valid.sum().item()) == int(
|
||||
original_embedded_payload["text_slot_valid"].sum().item()
|
||||
)
|
||||
release(restarted)
|
||||
report["memory"] = memory_report
|
||||
|
||||
baseline_quality = report["baseline"]["quality"]
|
||||
dynamic_quality = report[memory_variant]["quality"]
|
||||
categories = sorted(
|
||||
set(baseline_quality["categories"]) & set(dynamic_quality["categories"])
|
||||
)
|
||||
category_deltas = {
|
||||
category: dynamic_quality["categories"][category]["score"]
|
||||
- baseline_quality["categories"][category]["score"]
|
||||
for category in categories
|
||||
}
|
||||
report["comparison"] = {
|
||||
"overall_delta": dynamic_quality["overall_score"] - baseline_quality["overall_score"],
|
||||
"category_deltas": category_deltas,
|
||||
"peak_memory_allocated_delta_gb": report[memory_variant].get("peak_memory_allocated_gb", 0.0)
|
||||
- report["baseline"].get("peak_memory_allocated_gb", 0.0),
|
||||
"load_seconds_delta": report[memory_variant]["load_seconds"]
|
||||
- report["baseline"]["load_seconds"],
|
||||
}
|
||||
|
||||
output = Path(args.output)
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
compact = {
|
||||
"baseline_score": baseline_quality["overall_score"],
|
||||
f"{memory_variant}_score": dynamic_quality["overall_score"],
|
||||
"overall_delta": report["comparison"]["overall_delta"],
|
||||
"memory_restart_pass": report["memory"]["restart_pass"],
|
||||
"memory_reset_pass": report["memory"]["reset"]["cleared"],
|
||||
"cases": len(all_cases),
|
||||
"output": str(output),
|
||||
}
|
||||
print(json.dumps(compact, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,423 @@
|
||||
"""Compare the unmodified local Qwen checkpoint with a memory-surgery adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import math
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
|
||||
import torch
|
||||
|
||||
from .qwen_integration import QwenMemoryConfig, load_memory_config, load_qwen_base, load_qwen_dynamic, load_tokenizer
|
||||
from .train_qwen_memory import encode_messages, load_records, pad_batch
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
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_surgery_smoke",
|
||||
help="dynamic-memory adapter directory; its memory_config.json selects the surgery mode",
|
||||
)
|
||||
parser.add_argument("--output", default="dynamic_memory_lab/benchmark_qwen.json")
|
||||
parser.add_argument("--max-length", type=int, default=512)
|
||||
parser.add_argument("--repeats", type=int, default=3)
|
||||
parser.add_argument("--warmup", type=int, default=1)
|
||||
parser.add_argument("--max-new-tokens", type=int, default=8)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _adapter_config(adapter_dir: str | Path) -> QwenMemoryConfig:
|
||||
return load_memory_config(adapter_dir)
|
||||
|
||||
|
||||
def _score_output(output: Any, labels: torch.Tensor) -> tuple[float, int, int, int, bool]:
|
||||
shifted_labels = labels[..., 1:]
|
||||
predictions = output.logits[..., :-1, :].argmax(dim=-1)
|
||||
target_positions = shifted_labels != -100
|
||||
token_count = int(target_positions.sum().item())
|
||||
if output.loss is None or token_count == 0:
|
||||
raise RuntimeError("benchmark example has no supervised target tokens")
|
||||
correct_tokens = int((predictions[target_positions] == shifted_labels[target_positions]).sum().item())
|
||||
first_target = target_positions.nonzero(as_tuple=False)[0]
|
||||
first_token_correct = int(
|
||||
predictions[first_target[0], first_target[1]] == shifted_labels[first_target[0], first_target[1]]
|
||||
)
|
||||
sequence_ok = correct_tokens == token_count
|
||||
return float(output.loss.detach().item()), token_count, correct_tokens, first_token_correct, sequence_ok
|
||||
|
||||
|
||||
def _evaluate_base(model: Any, tokenizer: Any, records: list[dict[str, Any]], max_length: int) -> dict[str, float]:
|
||||
device = model.get_input_embeddings().weight.device
|
||||
total_nll = 0.0
|
||||
total_tokens = 0
|
||||
correct_tokens = 0
|
||||
correct_first_tokens = 0
|
||||
correct_sequences = 0
|
||||
for record in records:
|
||||
query = encode_messages(tokenizer, record["query"], max_length)
|
||||
input_ids, attention_mask, labels = pad_batch([query], int(tokenizer.pad_token_id))
|
||||
with torch.inference_mode():
|
||||
output = model(
|
||||
input_ids=input_ids.to(device),
|
||||
attention_mask=attention_mask.to(device),
|
||||
labels=labels.to(device),
|
||||
use_cache=False,
|
||||
)
|
||||
loss, tokens, tokens_correct, first_token_correct, sequence_ok = _score_output(output, labels.to(device))
|
||||
total_nll += loss * tokens
|
||||
total_tokens += tokens
|
||||
correct_tokens += tokens_correct
|
||||
correct_first_tokens += first_token_correct
|
||||
correct_sequences += int(sequence_ok)
|
||||
mean_loss = total_nll / total_tokens
|
||||
return {
|
||||
"loss": mean_loss,
|
||||
"perplexity": math.exp(mean_loss),
|
||||
"token_accuracy": correct_tokens / total_tokens,
|
||||
"first_target_token_accuracy": correct_first_tokens / len(records),
|
||||
"exact_sequence_accuracy": correct_sequences / len(records),
|
||||
"supervised_tokens": total_tokens,
|
||||
}
|
||||
|
||||
|
||||
def _evaluate_dynamic(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
records: list[dict[str, Any]],
|
||||
max_length: int,
|
||||
) -> dict[str, float]:
|
||||
device = model._find_layer_device()
|
||||
pad_id = int(tokenizer.pad_token_id)
|
||||
total_nll = 0.0
|
||||
total_tokens = 0
|
||||
correct_tokens = 0
|
||||
correct_first_tokens = 0
|
||||
correct_sequences = 0
|
||||
for record in records:
|
||||
model.reset_memory()
|
||||
memory = encode_messages(tokenizer, record["memory"], max_length)
|
||||
query = encode_messages(tokenizer, record["query"], max_length)
|
||||
memory_input, memory_mask, _ = pad_batch([memory], pad_id)
|
||||
query_input, query_mask, query_labels = pad_batch([query], pad_id)
|
||||
with torch.inference_mode():
|
||||
memory_output = 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,
|
||||
)
|
||||
output = model(
|
||||
input_ids=query_input.to(device),
|
||||
attention_mask=query_mask.to(device),
|
||||
labels=query_labels.to(device),
|
||||
memory_state=memory_output.memory,
|
||||
read_memory=True,
|
||||
update_memory=False,
|
||||
return_memory=True,
|
||||
use_cache=False,
|
||||
)
|
||||
loss, tokens, tokens_correct, first_token_correct, sequence_ok = _score_output(
|
||||
output, query_labels.to(device)
|
||||
)
|
||||
total_nll += loss * tokens
|
||||
total_tokens += tokens
|
||||
correct_tokens += tokens_correct
|
||||
correct_first_tokens += first_token_correct
|
||||
correct_sequences += int(sequence_ok)
|
||||
model.reset_memory()
|
||||
mean_loss = total_nll / total_tokens
|
||||
return {
|
||||
"loss": mean_loss,
|
||||
"perplexity": math.exp(mean_loss),
|
||||
"token_accuracy": correct_tokens / total_tokens,
|
||||
"first_target_token_accuracy": correct_first_tokens / len(records),
|
||||
"exact_sequence_accuracy": correct_sequences / len(records),
|
||||
"supervised_tokens": total_tokens,
|
||||
}
|
||||
|
||||
|
||||
def _sync(device: torch.device) -> None:
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize(device)
|
||||
|
||||
|
||||
def _measure(name: str, fn: Callable[[], int], repeats: int, warmup: int, device: torch.device) -> dict[str, float]:
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
_sync(device)
|
||||
started = time.perf_counter()
|
||||
processed = 0
|
||||
for _ in range(repeats):
|
||||
processed += fn()
|
||||
_sync(device)
|
||||
elapsed = time.perf_counter() - started
|
||||
return {
|
||||
"seconds": elapsed / repeats,
|
||||
"tokens_per_second": processed / elapsed,
|
||||
}
|
||||
|
||||
|
||||
def _generation_prompt(tokenizer: Any, messages: list[dict[str, Any]], device: torch.device) -> dict[str, torch.Tensor]:
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
input_ids = encoded["input_ids"] if hasattr(encoded, "__getitem__") and "input_ids" in encoded else encoded
|
||||
if isinstance(input_ids, torch.Tensor):
|
||||
input_ids = input_ids.tolist()
|
||||
if input_ids and isinstance(input_ids[0], list):
|
||||
input_ids = input_ids[0]
|
||||
ids = torch.tensor([input_ids], dtype=torch.long, device=device)
|
||||
return {"input_ids": ids, "attention_mask": torch.ones_like(ids)}
|
||||
|
||||
|
||||
def _measure_base_generation(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
record: dict[str, Any],
|
||||
repeats: int,
|
||||
warmup: int,
|
||||
max_new_tokens: int,
|
||||
) -> dict[str, float]:
|
||||
device = model.get_input_embeddings().weight.device
|
||||
prompt = _generation_prompt(tokenizer, record["query"][:-1], device)
|
||||
|
||||
def run() -> int:
|
||||
output = model.generate(
|
||||
**prompt,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
return int(output.shape[1] - prompt["input_ids"].shape[1])
|
||||
|
||||
return _measure("baseline_generation", run, repeats, warmup, device)
|
||||
|
||||
|
||||
def _measure_dynamic_generation(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
record: dict[str, Any],
|
||||
max_length: int,
|
||||
repeats: int,
|
||||
warmup: int,
|
||||
max_new_tokens: int,
|
||||
) -> dict[str, float]:
|
||||
device = model._find_layer_device()
|
||||
pad_id = int(tokenizer.pad_token_id)
|
||||
memory = encode_messages(tokenizer, record["memory"], max_length)
|
||||
memory_input, memory_mask, _ = pad_batch([memory], pad_id)
|
||||
prompt = _generation_prompt(tokenizer, record["query"][:-1], device)
|
||||
|
||||
def run() -> int:
|
||||
model.reset_memory()
|
||||
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,
|
||||
)
|
||||
output = model.generate(
|
||||
**prompt,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
return int(output.shape[1] - prompt["input_ids"].shape[1])
|
||||
|
||||
result = _measure("dynamic_generation", run, repeats, warmup, device)
|
||||
model.reset_memory()
|
||||
return result
|
||||
|
||||
|
||||
def _clean_generated(text: str) -> str:
|
||||
return text.replace(" ", "").replace("\r", "").replace("\n", "").strip()
|
||||
|
||||
|
||||
def _generation_quality_base(model: Any, tokenizer: Any, records: list[dict[str, Any]], max_new_tokens: int) -> dict[str, Any]:
|
||||
device = model.get_input_embeddings().weight.device
|
||||
contains = 0
|
||||
prefixes = 0
|
||||
examples = []
|
||||
for record in records:
|
||||
prompt = _generation_prompt(tokenizer, record["query"][:-1], device)
|
||||
with torch.inference_mode():
|
||||
output = model.generate(
|
||||
**prompt,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
generated = _clean_generated(tokenizer.decode(output[0, prompt["input_ids"].shape[1] :], skip_special_tokens=True))
|
||||
expected = str(record["answer"])
|
||||
contains += int(expected in generated)
|
||||
prefixes += int(generated.startswith(expected))
|
||||
if len(examples) < 3:
|
||||
examples.append({"expected": expected, "generated": generated})
|
||||
return {
|
||||
"answer_contains_accuracy": contains / len(records),
|
||||
"answer_prefix_accuracy": prefixes / len(records),
|
||||
"examples": examples,
|
||||
}
|
||||
|
||||
|
||||
def _generation_quality_dynamic(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
records: list[dict[str, Any]],
|
||||
max_length: int,
|
||||
max_new_tokens: int,
|
||||
) -> dict[str, Any]:
|
||||
device = model._find_layer_device()
|
||||
pad_id = int(tokenizer.pad_token_id)
|
||||
contains = 0
|
||||
prefixes = 0
|
||||
examples = []
|
||||
for record in records:
|
||||
model.reset_memory()
|
||||
memory = encode_messages(tokenizer, record["memory"], max_length)
|
||||
memory_input, memory_mask, _ = pad_batch([memory], pad_id)
|
||||
prompt = _generation_prompt(tokenizer, record["query"][:-1], device)
|
||||
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,
|
||||
)
|
||||
output = model.generate(
|
||||
**prompt,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
generated = _clean_generated(tokenizer.decode(output[0, prompt["input_ids"].shape[1] :], skip_special_tokens=True))
|
||||
expected = str(record["answer"])
|
||||
contains += int(expected in generated)
|
||||
prefixes += int(generated.startswith(expected))
|
||||
if len(examples) < 3:
|
||||
examples.append({"expected": expected, "generated": generated})
|
||||
model.reset_memory()
|
||||
return {
|
||||
"answer_contains_accuracy": contains / len(records),
|
||||
"answer_prefix_accuracy": prefixes / len(records),
|
||||
"examples": examples,
|
||||
}
|
||||
|
||||
|
||||
def _release(model: Any) -> None:
|
||||
del model
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if args.repeats < 1 or args.warmup < 0:
|
||||
raise ValueError("repeats must be >= 1 and warmup must be >= 0")
|
||||
torch.manual_seed(args.seed)
|
||||
tokenizer = load_tokenizer(args.model_path)
|
||||
records = load_records(args.data)
|
||||
use_4bit = not args.no_4bit
|
||||
results: dict[str, Any] = {
|
||||
"model_path": str(Path(args.model_path).resolve()),
|
||||
"data": str(Path(args.data).resolve()),
|
||||
"adapter": str(Path(args.adapter).resolve()),
|
||||
"records": len(records),
|
||||
"max_length": args.max_length,
|
||||
"quantization": "4bit_nf4" if use_4bit else "none",
|
||||
}
|
||||
|
||||
print("loading baseline")
|
||||
baseline = load_qwen_base(args.model_path, load_in_4bit=use_4bit)
|
||||
baseline.eval()
|
||||
baseline_device = baseline.get_input_embeddings().weight.device
|
||||
if baseline_device.type == "cuda":
|
||||
torch.cuda.reset_peak_memory_stats(baseline_device)
|
||||
results["baseline"] = {
|
||||
"device": str(baseline_device),
|
||||
"scores": _evaluate_base(baseline, tokenizer, records, args.max_length),
|
||||
"generation_quality": _generation_quality_base(baseline, tokenizer, records, args.max_new_tokens),
|
||||
"generation": _measure_base_generation(
|
||||
baseline, tokenizer, records[0], args.repeats, args.warmup, args.max_new_tokens
|
||||
),
|
||||
}
|
||||
if baseline_device.type == "cuda":
|
||||
results["baseline"]["peak_memory_gb"] = torch.cuda.max_memory_allocated(baseline_device) / 1024**3
|
||||
_release(baseline)
|
||||
baseline = None
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
memory_config = _adapter_config(args.adapter)
|
||||
print(f"loading dynamic mode={memory_config.mode} layers={memory_config.layer_indices}")
|
||||
dynamic = load_qwen_dynamic(
|
||||
args.model_path,
|
||||
memory_config=memory_config,
|
||||
load_in_4bit=use_4bit,
|
||||
)
|
||||
dynamic.load_memory_adapter(args.adapter)
|
||||
dynamic.eval()
|
||||
dynamic_device = dynamic._find_layer_device()
|
||||
if dynamic_device.type == "cuda":
|
||||
torch.cuda.reset_peak_memory_stats(dynamic_device)
|
||||
results["dynamic_memory"] = {
|
||||
"device": str(dynamic_device),
|
||||
"mode": memory_config.mode,
|
||||
"layers": list(dynamic.layer_indices),
|
||||
"scores": _evaluate_dynamic(dynamic, tokenizer, records, args.max_length),
|
||||
"generation_quality": _generation_quality_dynamic(
|
||||
dynamic, tokenizer, records, args.max_length, args.max_new_tokens
|
||||
),
|
||||
"generation": _measure_dynamic_generation(
|
||||
dynamic,
|
||||
tokenizer,
|
||||
records[0],
|
||||
args.max_length,
|
||||
args.repeats,
|
||||
args.warmup,
|
||||
args.max_new_tokens,
|
||||
),
|
||||
"trainable_parameters": sum(parameter.numel() for parameter in dynamic.trainable_parameters),
|
||||
}
|
||||
if dynamic_device.type == "cuda":
|
||||
results["dynamic_memory"]["peak_memory_gb"] = torch.cuda.max_memory_allocated(dynamic_device) / 1024**3
|
||||
_release(dynamic)
|
||||
dynamic = None
|
||||
|
||||
output_path = Path(args.output)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_text(json.dumps(results, indent=2), encoding="utf-8")
|
||||
|
||||
print(json.dumps(results, indent=2, ensure_ascii=False))
|
||||
print(f"saved={output_path.resolve()}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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))
|
||||
@@ -0,0 +1,250 @@
|
||||
"""Build a larger hard-negative dataset for the automatic memory controller.
|
||||
|
||||
The normal bootstrap corpus contains mostly short, obvious examples. This
|
||||
hard set adds realistic negations, questions, replacement requests, and long
|
||||
noise clauses so the write/forget heads are evaluated on decisions that are
|
||||
easy to confuse with durable facts.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import random
|
||||
import string
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent
|
||||
ATTRIBUTES = (
|
||||
"常用时区",
|
||||
"最喜欢的水果",
|
||||
"项目代号",
|
||||
"备用联系人",
|
||||
"默认输出风格",
|
||||
"工作区域",
|
||||
"提醒时间",
|
||||
"档案代号",
|
||||
)
|
||||
|
||||
TRAIN_FACT_TEMPLATES = (
|
||||
"请记住：我的{attribute}是{value}。",
|
||||
"以后涉及{attribute}时，请使用{value}这个值。",
|
||||
"个人资料更新——我的{attribute}等于{value}，以后可能会问到。",
|
||||
"请把这条个人资料保存下来：我的{attribute}为{value}。",
|
||||
"我的{attribute}是{value}，这是需要长期保留的信息。",
|
||||
)
|
||||
EVAL_FACT_TEMPLATES = (
|
||||
"登记一下，我的{attribute}：{value}。",
|
||||
"将我的{attribute}记为{value}，后续请按这个资料回答。",
|
||||
"长期资料里新增一项：{attribute}={value}。",
|
||||
"请把我的{attribute}保存成{value}。",
|
||||
)
|
||||
TRAIN_NOISE_TEMPLATES = (
|
||||
"这是普通对话噪声：我暂时提到一个无关编号{value}，不需要长期保存。",
|
||||
"随口一提，编号{value}只是临时信息，请不要记住。",
|
||||
"请不要把这句话写入长期记忆：今天看到的临时编号是{value}。",
|
||||
"这只是一次性测试值{value}，不用保存，也不要据此推断个人资料。",
|
||||
"聊天中的无关内容：{value}；它不是我的个人事实。",
|
||||
)
|
||||
EVAL_NOISE_TEMPLATES = (
|
||||
"临时提到{value}，这不是需要保存的资料。",
|
||||
"忽略这个一次性编号{value}，不要将它写入记忆。",
|
||||
"普通闲聊内容：{value}，没有长期价值。",
|
||||
"不要记住{value}，它只是当前消息里的干扰项。",
|
||||
)
|
||||
TRAIN_QUERY_TEMPLATES = (
|
||||
"我的{attribute}是什么？",
|
||||
"只根据已经保存的资料，告诉我{attribute}。",
|
||||
"不要猜测，请读取记忆回答：我的{attribute}为？",
|
||||
"记忆中是否有我的{attribute}？",
|
||||
)
|
||||
EVAL_QUERY_TEMPLATES = (
|
||||
"跨对话后，我登记的{attribute}是哪一个？",
|
||||
"请从长期资料中查找我的{attribute}。",
|
||||
"之前保存的{attribute}内容是什么？",
|
||||
)
|
||||
TRAIN_FORGET_TEMPLATES = (
|
||||
"请删除关于我的{attribute}的记忆，不要再保留{value}。",
|
||||
"忘掉我的{attribute}，这条资料已经失效。",
|
||||
"撤销之前保存的{attribute}，以后不要再使用它。",
|
||||
"清除我的{attribute}记录；{value}不再有效。",
|
||||
)
|
||||
EVAL_FORGET_TEMPLATES = (
|
||||
"请移除长期记忆中的{attribute}，不要继续记住它。",
|
||||
"我的{attribute}已经作废，请忘记这项资料。",
|
||||
"撤回关于{attribute}的个人信息，不要再保留。",
|
||||
)
|
||||
|
||||
|
||||
def _value(rng: random.Random, prefix: str) -> str:
|
||||
alphabet = string.ascii_uppercase + string.digits
|
||||
return f"{prefix}-" + "".join(rng.choice(alphabet) for _ in range(8))
|
||||
|
||||
|
||||
def _row(
|
||||
*,
|
||||
group: str,
|
||||
index: int,
|
||||
text: str,
|
||||
write: float,
|
||||
forget: float,
|
||||
kind: str,
|
||||
attribute: str,
|
||||
value: str,
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"id": f"{group}:{index}",
|
||||
"group_id": group,
|
||||
"text": text,
|
||||
"messages": [{"role": "user", "content": text}],
|
||||
"write_label": write,
|
||||
"forget_label": forget,
|
||||
"kind": kind,
|
||||
"source": "synthetic_memory_policy_hardset",
|
||||
"subject": "验证用户",
|
||||
"attribute": attribute,
|
||||
"value": value,
|
||||
"answer": "",
|
||||
"answerable": None,
|
||||
}
|
||||
|
||||
|
||||
def _build_split(
|
||||
*,
|
||||
count: int,
|
||||
split: str,
|
||||
seed: int,
|
||||
fact_templates: tuple[str, ...],
|
||||
noise_templates: tuple[str, ...],
|
||||
query_templates: tuple[str, ...],
|
||||
forget_templates: tuple[str, ...],
|
||||
) -> list[dict[str, Any]]:
|
||||
rng = random.Random(seed)
|
||||
rows: list[dict[str, Any]] = []
|
||||
for group_index in range(count):
|
||||
group = f"hard-{split}-{group_index:06d}"
|
||||
attribute = ATTRIBUTES[group_index % len(ATTRIBUTES)]
|
||||
fact_value = _value(rng, "FACT")
|
||||
noise_value = _value(rng, "NOISE")
|
||||
replacement_value = _value(rng, "NEW")
|
||||
rows.extend(
|
||||
(
|
||||
_row(
|
||||
group=group,
|
||||
index=0,
|
||||
text=rng.choice(fact_templates).format(attribute=attribute, value=fact_value),
|
||||
write=1.0,
|
||||
forget=0.0,
|
||||
kind="fact",
|
||||
attribute=attribute,
|
||||
value=fact_value,
|
||||
),
|
||||
_row(
|
||||
group=group,
|
||||
index=1,
|
||||
text=rng.choice(noise_templates).format(value=noise_value),
|
||||
write=0.0,
|
||||
forget=0.0,
|
||||
kind="noise",
|
||||
attribute="",
|
||||
value=noise_value,
|
||||
),
|
||||
_row(
|
||||
group=group,
|
||||
index=2,
|
||||
text=rng.choice(query_templates).format(attribute=attribute),
|
||||
write=0.0,
|
||||
forget=0.0,
|
||||
kind="query",
|
||||
attribute=attribute,
|
||||
value="",
|
||||
),
|
||||
_row(
|
||||
group=group,
|
||||
index=3,
|
||||
text=rng.choice(forget_templates).format(attribute=attribute, value=fact_value),
|
||||
write=0.0,
|
||||
forget=1.0,
|
||||
kind="forget",
|
||||
attribute=attribute,
|
||||
value=fact_value,
|
||||
),
|
||||
_row(
|
||||
group=group,
|
||||
index=4,
|
||||
text=rng.choice(fact_templates).format(attribute=attribute, value=replacement_value),
|
||||
write=1.0,
|
||||
# A replacement is a write/update, not a delete. The
|
||||
# runtime retires the matched old version and keeps the
|
||||
# new fragment active. Only an explicit forget request
|
||||
# receives forget_label=1.
|
||||
forget=0.0,
|
||||
kind="replacement",
|
||||
attribute=attribute,
|
||||
value=replacement_value,
|
||||
),
|
||||
)
|
||||
)
|
||||
rng.shuffle(rows)
|
||||
return rows
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--output-dir", default="data/production_memory_hard")
|
||||
parser.add_argument("--train-groups", type=int, default=4000)
|
||||
parser.add_argument("--eval-groups", type=int, default=1000)
|
||||
parser.add_argument("--seed", type=int, default=20260905)
|
||||
args = parser.parse_args()
|
||||
if args.train_groups < 1 or args.eval_groups < 1:
|
||||
raise SystemExit("group counts must be positive")
|
||||
output_dir = Path(args.output_dir)
|
||||
if not output_dir.is_absolute() and not output_dir.exists():
|
||||
output_dir = PROJECT_ROOT / output_dir
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
train_rows = _build_split(
|
||||
count=args.train_groups,
|
||||
split="train",
|
||||
seed=args.seed,
|
||||
fact_templates=TRAIN_FACT_TEMPLATES,
|
||||
noise_templates=TRAIN_NOISE_TEMPLATES,
|
||||
query_templates=TRAIN_QUERY_TEMPLATES,
|
||||
forget_templates=TRAIN_FORGET_TEMPLATES,
|
||||
)
|
||||
eval_rows = _build_split(
|
||||
count=args.eval_groups,
|
||||
split="eval",
|
||||
seed=args.seed + 1,
|
||||
fact_templates=EVAL_FACT_TEMPLATES,
|
||||
noise_templates=EVAL_NOISE_TEMPLATES,
|
||||
query_templates=EVAL_QUERY_TEMPLATES,
|
||||
forget_templates=EVAL_FORGET_TEMPLATES,
|
||||
)
|
||||
for name, rows in (("train", train_rows), ("eval", eval_rows)):
|
||||
with (output_dir / f"{name}.jsonl").open("w", encoding="utf-8") as handle:
|
||||
for row in rows:
|
||||
handle.write(json.dumps(row, ensure_ascii=False) + "\n")
|
||||
manifest = {
|
||||
"format_version": 1,
|
||||
"generator": "build_memory_policy_hardset.py",
|
||||
"seed": args.seed,
|
||||
"train_groups": args.train_groups,
|
||||
"eval_groups": args.eval_groups,
|
||||
"train_examples": len(train_rows),
|
||||
"eval_examples": len(eval_rows),
|
||||
"labels": {
|
||||
"write_positive": sum(row["write_label"] >= 0.5 for row in train_rows),
|
||||
"forget_positive": sum(row["forget_label"] >= 0.5 for row in train_rows),
|
||||
},
|
||||
"warning": "Synthetic hard negatives; combine with redacted real conversations before production deployment.",
|
||||
}
|
||||
(output_dir / "manifest.json").write_text(
|
||||
json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8"
|
||||
)
|
||||
print(json.dumps(manifest, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,293 @@
|
||||
"""Build an embedded Natural Memory v2 package from the current v1 package.
|
||||
|
||||
Unchanged Qwen shards are hard-linked when the filesystem permits it. The
|
||||
custom memory shard is rewritten once to include the trained V2 router and a
|
||||
compact V2 page payload, while the official model shards remain untouched.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
from safetensors.torch import save_file
|
||||
|
||||
if __package__ in {None, ""}:
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
from dynamic_memory_lab.memory_os_v2 import MemoryOSV2, MemoryRouterV2
|
||||
|
||||
|
||||
def _link_or_copy(source: Path, target: Path, *, allow_copy: bool) -> None:
|
||||
try:
|
||||
os.link(source, target)
|
||||
except OSError:
|
||||
if not allow_copy:
|
||||
raise RuntimeError(
|
||||
"the destination filesystem does not support hard links; "
|
||||
"refusing to duplicate multi-gigabyte Qwen shards. "
|
||||
"Re-run with --allow-copy-base only when enough disk space "
|
||||
"has been explicitly reserved"
|
||||
)
|
||||
shutil.copy2(source, target)
|
||||
|
||||
|
||||
def _pack_v2_payload(payload: dict[str, Any]) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
|
||||
tensors: dict[str, torch.Tensor] = {}
|
||||
metadata = dict(payload)
|
||||
|
||||
def pack_record(item: dict[str, Any], prefix: str) -> dict[str, Any]:
|
||||
item = dict(item)
|
||||
for field in ("key", "summary", "token_ids", "token_mask"):
|
||||
value = item.pop(field, None)
|
||||
if isinstance(value, torch.Tensor):
|
||||
name = f"dynamic_memory.v2.{prefix}.{field}"
|
||||
tensors[name] = value.detach().cpu().contiguous()
|
||||
item[f"{field}_ref"] = name
|
||||
return item
|
||||
|
||||
metadata["records"] = [
|
||||
pack_record(item, f"records.{index}")
|
||||
for index, item in enumerate(payload.get("records", []))
|
||||
]
|
||||
metadata["quarantine"] = [
|
||||
pack_record(item, f"quarantine.{index}")
|
||||
for index, item in enumerate(payload.get("quarantine", []))
|
||||
]
|
||||
metadata["pages"] = []
|
||||
for index, item in enumerate(payload.get("pages", [])):
|
||||
item = dict(item)
|
||||
for field in ("key", "summary"):
|
||||
value = item.pop(field, None)
|
||||
if isinstance(value, torch.Tensor):
|
||||
name = f"dynamic_memory.v2.pages.{index}.{field}"
|
||||
tensors[name] = value.detach().cpu().contiguous()
|
||||
item[f"{field}_ref"] = name
|
||||
metadata["pages"].append(item)
|
||||
return tensors, metadata
|
||||
|
||||
|
||||
def build(args: argparse.Namespace) -> dict[str, Any]:
|
||||
base_dir = Path(args.base_package)
|
||||
output_dir = Path(args.output_dir)
|
||||
if output_dir.exists():
|
||||
raise FileExistsError(f"refusing to overwrite existing output: {output_dir}")
|
||||
output_dir.mkdir(parents=True)
|
||||
|
||||
base_manifest = json.loads((base_dir / "memory_merge.json").read_text(encoding="utf-8"))
|
||||
base_config = json.loads((base_dir / "memory_config.json").read_text(encoding="utf-8"))
|
||||
memory_name = str(base_manifest.get("memory_weights", "model.safetensors-00003-of-00003.safetensors"))
|
||||
source_memory = base_dir / memory_name
|
||||
if not source_memory.exists():
|
||||
raise FileNotFoundError(source_memory)
|
||||
|
||||
sources = [
|
||||
source
|
||||
for source in base_dir.iterdir()
|
||||
if source.is_file()
|
||||
and source.name not in {
|
||||
memory_name,
|
||||
"memory_config.json",
|
||||
"memory_merge.json",
|
||||
"model.safetensors.index.json",
|
||||
}
|
||||
]
|
||||
if not args.allow_copy_base:
|
||||
probe_source = next(
|
||||
(source for source in sources if source.name.endswith(".safetensors")),
|
||||
None,
|
||||
)
|
||||
if probe_source is not None:
|
||||
probe_target = output_dir / ".hardlink-probe"
|
||||
try:
|
||||
os.link(probe_source, probe_target)
|
||||
except OSError as error:
|
||||
raise RuntimeError(
|
||||
"base package cannot be assembled without copying its large "
|
||||
"shards on this filesystem; pass --allow-copy-base after "
|
||||
"checking free space"
|
||||
) from error
|
||||
finally:
|
||||
if probe_target.exists():
|
||||
probe_target.unlink()
|
||||
|
||||
for source in sources:
|
||||
_link_or_copy(source, output_dir / source.name, allow_copy=args.allow_copy_base)
|
||||
|
||||
with safe_open(str(source_memory), framework="pt", device="cpu") as handle:
|
||||
base_tensors = {key: handle.get_tensor(key) for key in handle.keys()}
|
||||
|
||||
router = MemoryRouterV2(
|
||||
int(base_config["hidden_size"]),
|
||||
router_dim=args.router_dim,
|
||||
num_heads=args.num_heads,
|
||||
max_hops=args.max_hops,
|
||||
)
|
||||
router_state = torch.load(args.router_checkpoint, map_location="cpu", weights_only=True)
|
||||
router.load_state_dict(router_state, strict=True)
|
||||
for key, value in router.state_dict().items():
|
||||
base_tensors[f"dynamic_memory.memory_router_v2.{key}"] = value.detach().cpu().contiguous()
|
||||
|
||||
memory_os = MemoryOSV2(
|
||||
int(base_config["hidden_size"]),
|
||||
router=router,
|
||||
)
|
||||
# Migrate the existing model-owned hot facts into V2 address pages. The
|
||||
# old bank remains intact; this is only a compatibility seed for the new
|
||||
# hierarchical route.
|
||||
legacy_prefix = "dynamic_memory.persistent."
|
||||
token_ids = base_tensors.get(f"{legacy_prefix}text_token_ids")
|
||||
token_mask = base_tensors.get(f"{legacy_prefix}text_token_mask")
|
||||
slot_valid = base_tensors.get(f"{legacy_prefix}text_slot_valid")
|
||||
slot_keys = base_tensors.get(f"{legacy_prefix}text_slot_keys")
|
||||
if all(isinstance(value, torch.Tensor) for value in (token_ids, token_mask, slot_valid, slot_keys)):
|
||||
for batch_index in range(slot_valid.shape[0]):
|
||||
for slot in range(slot_valid.shape[1]):
|
||||
if not bool(slot_valid[batch_index, slot].item()):
|
||||
continue
|
||||
ids = token_ids[batch_index, slot][token_mask[batch_index, slot]]
|
||||
memory_os.write(
|
||||
text=f"legacy_hot_slot:{batch_index}:{slot}",
|
||||
key=slot_keys[batch_index, slot],
|
||||
summary=slot_keys[batch_index, slot],
|
||||
memory_type="legacy_hot_text",
|
||||
importance=0.95,
|
||||
confidence=0.95,
|
||||
source="v1_migration",
|
||||
slot_index=slot,
|
||||
token_ids=ids,
|
||||
token_mask=torch.ones_like(ids, dtype=torch.bool),
|
||||
trusted=True,
|
||||
)
|
||||
|
||||
v2_tensors, v2_metadata = _pack_v2_payload(memory_os.export_payload())
|
||||
base_tensors.update(v2_tensors)
|
||||
save_file(
|
||||
base_tensors,
|
||||
str(output_dir / memory_name),
|
||||
metadata={
|
||||
"format": "qwen_dynamic_memory_embedded_v2",
|
||||
"memory_os_v2_payload": json.dumps(v2_metadata, ensure_ascii=False, separators=(",", ":")),
|
||||
},
|
||||
)
|
||||
|
||||
memory_config = dict(base_config)
|
||||
saved = dict(memory_config.get("memory_config", {}))
|
||||
saved.update(
|
||||
{
|
||||
"memory_version": 2,
|
||||
"hierarchical_memory": True,
|
||||
"memory_router_dim": args.router_dim,
|
||||
"memory_router_heads": args.num_heads,
|
||||
"memory_page_capacity": args.page_capacity,
|
||||
"memory_max_pages": args.max_pages,
|
||||
"memory_hot_pages": args.hot_pages,
|
||||
"memory_top_k_pages": args.top_k_pages,
|
||||
"memory_top_k_records": args.top_k_records,
|
||||
"memory_max_hops": args.max_hops,
|
||||
"memory_coarse_index_bits": args.coarse_index_bits,
|
||||
"memory_v2_read_threshold": args.read_threshold,
|
||||
"memory_v2_write_threshold": args.write_threshold,
|
||||
"memory_storage_mode": args.memory_storage_mode,
|
||||
"memory_storage_path": args.memory_storage_path,
|
||||
"memory_resident_pages": args.memory_resident_pages,
|
||||
"memory_gpu_cache_records": args.memory_gpu_cache_records,
|
||||
"memory_gpu_cache_tokens": args.memory_gpu_cache_tokens,
|
||||
"memory_gpu_cache_reserve_mb": args.memory_gpu_cache_reserve_mb,
|
||||
"memory_gpu_cache_adaptive": args.memory_gpu_cache_adaptive,
|
||||
"kv_budget_tokens": args.kv_budget,
|
||||
"kv_hard_max_tokens": args.kv_hard_max,
|
||||
"kv_compaction_trigger": args.kv_trigger,
|
||||
"kv_keep_recent_tokens": args.kv_keep_recent,
|
||||
"persistent_memory": True,
|
||||
}
|
||||
)
|
||||
memory_config["memory_config"] = saved
|
||||
memory_config["router_v2_ready"] = True
|
||||
memory_config["checkpoint_contains_user_memory"] = True
|
||||
(output_dir / "memory_config.json").write_text(
|
||||
json.dumps(memory_config, indent=2, ensure_ascii=False), encoding="utf-8"
|
||||
)
|
||||
|
||||
index = json.loads((base_dir / "model.safetensors.index.json").read_text(encoding="utf-8"))
|
||||
weight_map = index.setdefault("weight_map", {})
|
||||
for key in base_tensors:
|
||||
weight_map[key] = memory_name
|
||||
index.setdefault("metadata", {})["total_size"] = int(
|
||||
sum(value.numel() * value.element_size() for value in base_tensors.values())
|
||||
)
|
||||
(output_dir / "model.safetensors.index.json").write_text(
|
||||
json.dumps(index, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
|
||||
)
|
||||
manifest = dict(base_manifest)
|
||||
manifest.update(
|
||||
{
|
||||
"format_version": 2,
|
||||
"format": "qwen_dynamic_memory_v2_embedded",
|
||||
"base_model": str(base_dir),
|
||||
"source_router_checkpoint": str(args.router_checkpoint),
|
||||
"checkpoint_contains_user_memory": True,
|
||||
"memory_config": saved,
|
||||
}
|
||||
)
|
||||
(output_dir / "memory_merge.json").write_text(
|
||||
json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8"
|
||||
)
|
||||
return {
|
||||
"output_dir": str(output_dir),
|
||||
"memory_shard": str(output_dir / memory_name),
|
||||
"memory_tensor_count": len(base_tensors),
|
||||
"v2_records": memory_os.stats()["records"],
|
||||
"v2_pages": memory_os.stats()["pages"],
|
||||
}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--base-package", default="W:/Flash/model/dynamic_memory_lab/qwen3_5_4b_memory_merged_v13")
|
||||
parser.add_argument("--output-dir", default="W:/Flash/model/dynamic_memory_lab/qwen3_5_4b_memory_merged_v2")
|
||||
parser.add_argument("--router-checkpoint", default="W:/Flash/model/dynamic_memory_lab/checkpoints/natural_memory_v2_router/memory_router_v2.pt")
|
||||
parser.add_argument("--router-dim", type=int, default=128)
|
||||
parser.add_argument("--num-heads", type=int, default=8)
|
||||
parser.add_argument("--max-hops", type=int, default=3)
|
||||
parser.add_argument("--page-capacity", type=int, default=32)
|
||||
parser.add_argument("--max-pages", type=int, default=32768)
|
||||
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=20)
|
||||
parser.add_argument("--read-threshold", type=float, default=0.65)
|
||||
parser.add_argument("--write-threshold", type=float, default=0.50)
|
||||
parser.add_argument("--memory-storage-mode", choices=("embedded", "tiered"), default="embedded")
|
||||
parser.add_argument("--memory-storage-path", default=None)
|
||||
parser.add_argument("--memory-resident-pages", type=int, default=256)
|
||||
parser.add_argument("--memory-gpu-cache-records", type=int, default=256)
|
||||
parser.add_argument("--memory-gpu-cache-tokens", type=int, default=131072)
|
||||
parser.add_argument("--memory-gpu-cache-reserve-mb", type=int, default=2048)
|
||||
parser.add_argument(
|
||||
"--memory-gpu-cache-adaptive",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="keep the VRAM cache below the reserve line and fall back to system RAM when needed",
|
||||
)
|
||||
parser.add_argument("--kv-budget", type=int, default=32768)
|
||||
parser.add_argument("--kv-hard-max", type=int, default=131072)
|
||||
parser.add_argument("--kv-trigger", type=float, default=0.90)
|
||||
parser.add_argument("--kv-keep-recent", type=int, default=8192)
|
||||
parser.add_argument(
|
||||
"--allow-copy-base",
|
||||
action="store_true",
|
||||
help="allow copying the frozen Qwen shards when hard links are unavailable; requires substantial free disk space",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(json.dumps(build(parse_args()), ensure_ascii=False, indent=2))
|
||||
@@ -0,0 +1,307 @@
|
||||
"""Normalize conversation logs into a leak-resistant memory-policy dataset.
|
||||
|
||||
The runtime accepts many local data shapes because real users rarely keep
|
||||
their chat exports in one format. This command converts them to a small,
|
||||
auditable JSONL schema without inventing labels. It understands the current
|
||||
``native_memory`` episode format, the streaming demo format, and a generic
|
||||
format documented in the output manifest.
|
||||
|
||||
The bundled fallback files are bootstrap data for smoke tests. A real user
|
||||
corpus can be supplied with ``--source``/``--eval-source`` and receives the
|
||||
same normalization and group-level split guarantees.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
from collections import Counter, defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable
|
||||
|
||||
|
||||
DEFAULT_SOURCES = (
|
||||
"data/native_memory/train.jsonl",
|
||||
"data/native_memory/eval.jsonl",
|
||||
"data/demo_stream.jsonl",
|
||||
)
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
def _project_path(value: str | Path) -> Path:
|
||||
path = Path(value)
|
||||
if path.is_absolute() or path.exists():
|
||||
return path
|
||||
return PROJECT_ROOT / path
|
||||
|
||||
|
||||
def _read_jsonl(path: Path) -> Iterable[tuple[int, dict[str, Any]]]:
|
||||
with path.open("r", encoding="utf-8") as handle:
|
||||
for line_number, raw in enumerate(handle, 1):
|
||||
raw = raw.strip()
|
||||
if not raw:
|
||||
continue
|
||||
value = json.loads(raw)
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError(f"{path}:{line_number} must contain a JSON object")
|
||||
yield line_number, value
|
||||
|
||||
|
||||
def _message_text(messages: Any) -> str:
|
||||
if isinstance(messages, str):
|
||||
return messages.strip()
|
||||
if not isinstance(messages, list):
|
||||
return ""
|
||||
parts: list[str] = []
|
||||
for message in messages:
|
||||
if not isinstance(message, dict):
|
||||
continue
|
||||
content = message.get("content", "")
|
||||
if isinstance(content, str) and content.strip():
|
||||
role = str(message.get("role", "user"))
|
||||
parts.append(f"[{role}] {content.strip()}")
|
||||
return "\n".join(parts).strip()
|
||||
|
||||
|
||||
def _user_text(messages: Any) -> str:
|
||||
if isinstance(messages, str):
|
||||
return messages.strip()
|
||||
if not isinstance(messages, list):
|
||||
return ""
|
||||
for message in reversed(messages):
|
||||
if isinstance(message, dict) and message.get("role") == "user":
|
||||
content = message.get("content", "")
|
||||
if isinstance(content, str):
|
||||
return content.strip()
|
||||
return _message_text(messages)
|
||||
|
||||
|
||||
def _explicit_split(path: Path, *, forced: str | None) -> str | None:
|
||||
if forced in {"train", "eval"}:
|
||||
return forced
|
||||
name = path.name.lower()
|
||||
if any(mark in name for mark in ("eval", "valid", "test")):
|
||||
return "eval"
|
||||
if "train" in name:
|
||||
return "train"
|
||||
return None
|
||||
|
||||
|
||||
def _make_example(
|
||||
*,
|
||||
group_id: str,
|
||||
example_id: str,
|
||||
text: str,
|
||||
write_label: float,
|
||||
forget_label: float = 0.0,
|
||||
kind: str = "conversation",
|
||||
source: str,
|
||||
subject: str = "",
|
||||
attribute: str = "",
|
||||
value: Any = None,
|
||||
answer: str = "",
|
||||
answerable: bool | None = None,
|
||||
messages: Any = None,
|
||||
) -> dict[str, Any] | None:
|
||||
text = str(text or "").strip()
|
||||
if not text:
|
||||
return None
|
||||
return {
|
||||
"id": example_id,
|
||||
"group_id": group_id,
|
||||
"text": text,
|
||||
"messages": messages if isinstance(messages, list) else [{"role": "user", "content": text}],
|
||||
"write_label": float(max(0.0, min(1.0, write_label))),
|
||||
"forget_label": float(max(0.0, min(1.0, forget_label))),
|
||||
"kind": kind,
|
||||
"source": source,
|
||||
"subject": str(subject or ""),
|
||||
"attribute": str(attribute or ""),
|
||||
"value": "" if value is None else str(value),
|
||||
"answer": str(answer or ""),
|
||||
"answerable": answerable,
|
||||
}
|
||||
|
||||
|
||||
def normalize_record(record: dict[str, Any], *, source: str, line_number: int) -> list[dict[str, Any]]:
|
||||
"""Convert one source record into labeled write/query decisions."""
|
||||
|
||||
raw_id = str(record.get("id") or record.get("conversation_id") or f"line-{line_number}")
|
||||
group_id = f"{source}:{raw_id}"
|
||||
output: list[dict[str, Any]] = []
|
||||
|
||||
chunks = record.get("memory_chunks")
|
||||
if isinstance(chunks, list):
|
||||
for index, chunk in enumerate(chunks):
|
||||
if not isinstance(chunk, dict):
|
||||
continue
|
||||
messages = chunk.get("messages", [])
|
||||
item = _make_example(
|
||||
group_id=group_id,
|
||||
example_id=f"{raw_id}:memory:{index}",
|
||||
text=_user_text(messages) or str(chunk.get("text", "")),
|
||||
write_label=float(chunk.get("write_label", 1.0)),
|
||||
forget_label=float(chunk.get("forget_label", 0.0)),
|
||||
kind=str(chunk.get("kind", "fact")),
|
||||
source=source,
|
||||
subject=record.get("subject", ""),
|
||||
attribute=record.get("attribute", ""),
|
||||
value=chunk.get("value", record.get("value", "")),
|
||||
messages=messages,
|
||||
)
|
||||
if item is not None:
|
||||
output.append(item)
|
||||
query = record.get("query")
|
||||
query_text = _user_text(query)
|
||||
item = _make_example(
|
||||
group_id=group_id,
|
||||
example_id=f"{raw_id}:query",
|
||||
text=query_text,
|
||||
write_label=0.0,
|
||||
kind="query",
|
||||
source=source,
|
||||
subject=record.get("subject", ""),
|
||||
attribute=record.get("attribute", ""),
|
||||
answer=record.get("answer", ""),
|
||||
answerable=record.get("answerable"),
|
||||
messages=query if isinstance(query, list) else None,
|
||||
)
|
||||
if item is not None:
|
||||
output.append(item)
|
||||
return output
|
||||
|
||||
memory = record.get("memory")
|
||||
if isinstance(memory, list):
|
||||
for index, item_messages in enumerate(memory):
|
||||
item = _make_example(
|
||||
group_id=group_id,
|
||||
example_id=f"{raw_id}:memory:{index}",
|
||||
text=_user_text(item_messages),
|
||||
write_label=1.0,
|
||||
kind="fact",
|
||||
source=source,
|
||||
messages=item_messages if isinstance(item_messages, list) else None,
|
||||
)
|
||||
if item is not None:
|
||||
output.append(item)
|
||||
query = record.get("query")
|
||||
if query is not None:
|
||||
item = _make_example(
|
||||
group_id=group_id,
|
||||
example_id=f"{raw_id}:query",
|
||||
text=_user_text(query),
|
||||
write_label=0.0,
|
||||
kind="query",
|
||||
source=source,
|
||||
answer=record.get("answer", ""),
|
||||
answerable=record.get("answerable"),
|
||||
messages=query if isinstance(query, list) else None,
|
||||
)
|
||||
if item is not None:
|
||||
output.append(item)
|
||||
|
||||
event = record.get("memory_event")
|
||||
if not output and (record.get("text") is not None or record.get("messages") is not None):
|
||||
event = event if isinstance(event, dict) else {}
|
||||
item = _make_example(
|
||||
group_id=group_id,
|
||||
example_id=f"{raw_id}:turn",
|
||||
text=_user_text(record.get("messages")) or str(record.get("text", "")),
|
||||
write_label=float(event.get("write_label", event.get("write", record.get("write_label", 0.0)))),
|
||||
forget_label=float(event.get("forget_label", event.get("forget", record.get("forget_label", 0.0)))),
|
||||
kind=str(event.get("kind", record.get("kind", "conversation"))),
|
||||
source=source,
|
||||
subject=record.get("subject", ""),
|
||||
attribute=record.get("attribute", ""),
|
||||
value=record.get("value", ""),
|
||||
answer=record.get("answer", ""),
|
||||
answerable=record.get("answerable"),
|
||||
messages=record.get("messages"),
|
||||
)
|
||||
if item is not None:
|
||||
output.append(item)
|
||||
return output
|
||||
|
||||
|
||||
def _split_for_group(group_id: str, explicit: str | None, *, eval_ratio: float) -> str:
|
||||
if explicit is not None:
|
||||
return explicit
|
||||
digest = hashlib.sha1(group_id.encode("utf-8")).hexdigest()
|
||||
value = int(digest[:8], 16) / 0xFFFFFFFF
|
||||
return "eval" if value < eval_ratio else "train"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--source", action="append", help="input JSONL; may be repeated")
|
||||
parser.add_argument("--eval-source", action="append", default=[], help="input JSONL forced into eval")
|
||||
parser.add_argument("--output-dir", default="data/production_memory")
|
||||
parser.add_argument("--eval-ratio", type=float, default=0.2)
|
||||
args = parser.parse_args()
|
||||
if not 0.0 < args.eval_ratio < 1.0:
|
||||
raise SystemExit("--eval-ratio must be between 0 and 1")
|
||||
|
||||
source_paths = [_project_path(item) for item in (args.source or DEFAULT_SOURCES)]
|
||||
eval_paths = [_project_path(item) for item in args.eval_source]
|
||||
all_inputs = [(path, None) for path in source_paths] + [(path, "eval") for path in eval_paths]
|
||||
examples: list[tuple[str, dict[str, Any]]] = []
|
||||
source_stats: dict[str, Counter[str]] = defaultdict(Counter)
|
||||
seen: set[tuple[str, str, float, float, str]] = set()
|
||||
for path, forced_split in all_inputs:
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(path)
|
||||
source = str(path)
|
||||
name_split = _explicit_split(path, forced=forced_split)
|
||||
for line_number, record in _read_jsonl(path):
|
||||
normalized = normalize_record(record, source=source, line_number=line_number)
|
||||
for item in normalized:
|
||||
dedupe_key = (
|
||||
item["group_id"],
|
||||
item["text"],
|
||||
item["write_label"],
|
||||
item["forget_label"],
|
||||
item["kind"],
|
||||
)
|
||||
if dedupe_key in seen:
|
||||
source_stats[source]["deduplicated"] += 1
|
||||
continue
|
||||
seen.add(dedupe_key)
|
||||
split = _split_for_group(item["group_id"], name_split, eval_ratio=args.eval_ratio)
|
||||
examples.append((split, item))
|
||||
source_stats[source][split] += 1
|
||||
|
||||
output_dir = _project_path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
split_counts: Counter[str] = Counter()
|
||||
for split in ("train", "eval"):
|
||||
path = output_dir / f"{split}.jsonl"
|
||||
with path.open("w", encoding="utf-8") as handle:
|
||||
for item_split, item in examples:
|
||||
if item_split == split:
|
||||
handle.write(json.dumps(item, ensure_ascii=False) + "\n")
|
||||
split_counts[split] += 1
|
||||
manifest = {
|
||||
"format_version": 1,
|
||||
"schema": {
|
||||
"text": "current turn presented to the write policy",
|
||||
"messages": "optional original chat messages",
|
||||
"write_label": "1 durable memory, 0 ordinary query/casual turn",
|
||||
"forget_label": "1 explicit correction/forget request",
|
||||
"group_id": "conversation/episode identity; never split across train and eval",
|
||||
},
|
||||
"bootstrap_data_warning": "Default files are local bootstrap/synthetic data; pass real exports with --source for production training.",
|
||||
"inputs": [str(path) for path, _ in all_inputs],
|
||||
"counts": dict(split_counts),
|
||||
"source_stats": {key: dict(value) for key, value in source_stats.items()},
|
||||
"dedupe_count": sum(value.get("deduplicated", 0) for value in source_stats.values()),
|
||||
}
|
||||
(output_dir / "manifest.json").write_text(
|
||||
json.dumps(manifest, ensure_ascii=False, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
print(json.dumps(manifest, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,297 @@
|
||||
"""Minimal interactive chat using the persistent Qwen dynamic memory."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from .qwen_integration import (
|
||||
DEFAULT_MEMORY_RESET_TOKEN,
|
||||
QwenMemoryConfig,
|
||||
load_memory_config,
|
||||
load_qwen_dynamic,
|
||||
load_tokenizer,
|
||||
resolve_memory_reset_token,
|
||||
split_memory_candidates,
|
||||
)
|
||||
|
||||
|
||||
def _memory_system_prefix(tokenizer, content: str):
|
||||
"""Encode a valid system-message prefix without adding a fake query."""
|
||||
|
||||
full = tokenizer.apply_chat_template(
|
||||
[
|
||||
{"role": "system", "content": content},
|
||||
{"role": "user", "content": "__memory_query_boundary__"},
|
||||
],
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
return_tensors="pt",
|
||||
return_dict=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
input_ids = full["input_ids"]
|
||||
im_start = tokenizer.convert_tokens_to_ids("<|im_start|>")
|
||||
positions = (input_ids[0] == int(im_start)).nonzero(as_tuple=False).flatten()
|
||||
if positions.numel() < 2:
|
||||
raise RuntimeError("could not locate the system/user memory boundary")
|
||||
end = int(positions[1].item())
|
||||
return {
|
||||
"input_ids": input_ids[:, :end],
|
||||
"attention_mask": torch.ones((1, end), dtype=torch.long),
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
if hasattr(sys.stdout, "reconfigure"):
|
||||
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model-path", default=".")
|
||||
parser.add_argument("--adapter", default=None)
|
||||
parser.add_argument(
|
||||
"--natural-language-memory",
|
||||
action="store_true",
|
||||
help="enable the model-owned exact text memory bank and internal retrieval prefix",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--memory-state",
|
||||
default=None,
|
||||
help="user runtime memory file; it is loaded at startup and saved after each turn",
|
||||
)
|
||||
parser.add_argument("--max-new-tokens", type=int, default=128)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
parser.add_argument(
|
||||
"--persistent-memory",
|
||||
action="store_true",
|
||||
help="keep native memory inside the model instance across turns",
|
||||
)
|
||||
parser.add_argument("--reset-token", default=None)
|
||||
parser.add_argument("--reset-token-id", type=int, default=None)
|
||||
parser.add_argument(
|
||||
"--persist-in-adapter",
|
||||
action="store_true",
|
||||
help="also checkpoint current user memory into the adapter package",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
tokenizer = load_tokenizer(args.model_path)
|
||||
memory_config = load_memory_config(args.adapter) if args.adapter else None
|
||||
if args.natural_language_memory and memory_config is None:
|
||||
memory_config = QwenMemoryConfig(natural_language_memory=True)
|
||||
if memory_config is not None and args.natural_language_memory:
|
||||
memory_config.natural_language_memory = True
|
||||
if memory_config is not None and memory_config.native_mode and args.persistent_memory:
|
||||
memory_config.persistent_memory = True
|
||||
if memory_config is not None:
|
||||
if args.reset_token_id is not None:
|
||||
memory_config.reset_token_id = args.reset_token_id
|
||||
elif args.reset_token is not None:
|
||||
memory_config.reset_token_id = resolve_memory_reset_token(tokenizer, args.reset_token)
|
||||
elif memory_config.native_mode and memory_config.reset_token_id is None:
|
||||
memory_config.reset_token_id = resolve_memory_reset_token(tokenizer)
|
||||
model = load_qwen_dynamic(
|
||||
args.model_path,
|
||||
memory_config=memory_config,
|
||||
load_in_4bit=not args.no_4bit,
|
||||
)
|
||||
if args.adapter:
|
||||
model.load_memory_adapter(args.adapter)
|
||||
model.eval()
|
||||
device = model._find_layer_device()
|
||||
state_path = Path(args.memory_state) if args.memory_state else None
|
||||
if state_path is not None and state_path.exists():
|
||||
model.load_runtime_memory(state_path, device=device)
|
||||
print(f"已加载用户 memory_state：{state_path}")
|
||||
print("普通消息会自动判断并保存重要信息；/remember <事实> 强制写入，/reset 清空，/save 保存，/quit 退出。")
|
||||
if memory_config is not None and memory_config.reset_token_id is not None:
|
||||
print(f"也可在用户消息中发送重置 token：{args.reset_token or DEFAULT_MEMORY_RESET_TOKEN}")
|
||||
|
||||
def save_state() -> None:
|
||||
if state_path is not None and model.runtime.state is not None:
|
||||
model.save_runtime_memory(state_path)
|
||||
print(f"已保存：{state_path}")
|
||||
if args.persist_in_adapter and args.adapter and model.runtime.state is not None:
|
||||
model.save_persistent_memory_checkpoint(args.adapter)
|
||||
print(f"已将当前用户记忆写入模型适配器：{args.adapter}")
|
||||
|
||||
try:
|
||||
while True:
|
||||
user_text = input("你> ").strip()
|
||||
if user_text == "/quit":
|
||||
break
|
||||
if user_text == "/reset":
|
||||
model.reset_memory(batch_size=1, device=device)
|
||||
save_state()
|
||||
print("已清空动态记忆。")
|
||||
continue
|
||||
if user_text == "/save":
|
||||
save_state()
|
||||
continue
|
||||
if user_text.startswith("/remember "):
|
||||
fact = user_text[len("/remember ") :].strip()
|
||||
if not fact:
|
||||
continue
|
||||
messages = [
|
||||
{"role": "user", "content": fact},
|
||||
{"role": "assistant", "content": "好的，我会记住这件事。"},
|
||||
]
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
add_generation_prompt=False,
|
||||
return_tensors="pt",
|
||||
return_dict=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
encoded = {
|
||||
key: value.to(device)
|
||||
for key, value in encoded.items()
|
||||
if isinstance(value, torch.Tensor)
|
||||
}
|
||||
memory_text = _memory_system_prefix(
|
||||
tokenizer,
|
||||
"以下是与当前用户相关的已保存长期记忆。仅在问题相关时使用，不要编造：\n" + fact,
|
||||
)
|
||||
memory_text_input_ids = memory_text["input_ids"].to(device)
|
||||
memory_text_attention_mask = memory_text.get("attention_mask")
|
||||
if memory_text_attention_mask is None:
|
||||
memory_text_attention_mask = torch.ones_like(memory_text_input_ids)
|
||||
memory_text_attention_mask = memory_text_attention_mask.to(device)
|
||||
memory_key = tokenizer(fact, add_special_tokens=False, return_tensors="pt")
|
||||
memory_key_input_ids = memory_key["input_ids"].to(device)
|
||||
memory_key_attention_mask = memory_key.get("attention_mask")
|
||||
if memory_key_attention_mask is None:
|
||||
memory_key_attention_mask = torch.ones_like(memory_key_input_ids)
|
||||
memory_key_attention_mask = memory_key_attention_mask.to(device)
|
||||
memory_storage = tokenizer(
|
||||
fact,
|
||||
add_special_tokens=False,
|
||||
return_tensors="pt",
|
||||
)
|
||||
with torch.no_grad():
|
||||
model(
|
||||
**encoded,
|
||||
read_memory=False,
|
||||
update_memory=True,
|
||||
return_memory=True,
|
||||
use_cache=False,
|
||||
memory_text_input_ids=memory_text_input_ids,
|
||||
memory_text_attention_mask=memory_text_attention_mask,
|
||||
memory_key_input_ids=memory_key_input_ids,
|
||||
memory_key_attention_mask=memory_key_attention_mask,
|
||||
memory_storage_input_ids=memory_storage["input_ids"].to(device),
|
||||
memory_storage_attention_mask=memory_storage.get(
|
||||
"attention_mask",
|
||||
torch.ones_like(memory_storage["input_ids"]),
|
||||
).to(device),
|
||||
force_memory_write=True,
|
||||
)
|
||||
save_state()
|
||||
print("已写入动态记忆。")
|
||||
continue
|
||||
if not user_text:
|
||||
continue
|
||||
|
||||
messages = [{"role": "user", "content": user_text}]
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
return_tensors="pt",
|
||||
return_dict=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
encoded = {
|
||||
key: value.to(device)
|
||||
for key, value in encoded.items()
|
||||
if isinstance(value, torch.Tensor)
|
||||
}
|
||||
memory_query = tokenizer(
|
||||
user_text,
|
||||
add_special_tokens=False,
|
||||
return_tensors="pt",
|
||||
)
|
||||
memory_query_input_ids = memory_query["input_ids"].to(device)
|
||||
memory_query_attention_mask = memory_query.get("attention_mask")
|
||||
if memory_query_attention_mask is None:
|
||||
memory_query_attention_mask = torch.ones_like(memory_query_input_ids)
|
||||
memory_query_attention_mask = memory_query_attention_mask.to(device)
|
||||
with torch.no_grad():
|
||||
# Native mode first gives the prompt to the learned controller
|
||||
# so it can decide whether each fact-sized candidate is worth
|
||||
# storing. The generation itself is read-only, preventing the
|
||||
# model from accidentally memorizing its own answer text.
|
||||
if memory_config is not None and memory_config.native_mode:
|
||||
for candidate in split_memory_candidates(user_text):
|
||||
candidate_encoded = tokenizer.apply_chat_template(
|
||||
[{"role": "user", "content": candidate}],
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
return_tensors="pt",
|
||||
return_dict=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
candidate_encoded = {
|
||||
key: value.to(device)
|
||||
for key, value in candidate_encoded.items()
|
||||
if isinstance(value, torch.Tensor)
|
||||
}
|
||||
memory_text = _memory_system_prefix(
|
||||
tokenizer,
|
||||
"以下是与当前用户相关的已保存长期记忆。仅在问题相关时使用，不要编造：\n"
|
||||
+ candidate,
|
||||
)
|
||||
memory_key = tokenizer(
|
||||
candidate,
|
||||
add_special_tokens=False,
|
||||
return_tensors="pt",
|
||||
)
|
||||
memory_key_input_ids = memory_key["input_ids"].to(device)
|
||||
memory_key_attention_mask = memory_key.get("attention_mask")
|
||||
if memory_key_attention_mask is None:
|
||||
memory_key_attention_mask = torch.ones_like(memory_key_input_ids)
|
||||
memory_storage = tokenizer(
|
||||
candidate,
|
||||
add_special_tokens=False,
|
||||
return_tensors="pt",
|
||||
)
|
||||
model(
|
||||
**candidate_encoded,
|
||||
read_memory=False,
|
||||
update_memory=True,
|
||||
return_memory=True,
|
||||
use_cache=False,
|
||||
memory_text_input_ids=memory_text["input_ids"].to(device),
|
||||
memory_text_attention_mask=torch.ones_like(
|
||||
memory_text["input_ids"], device=device
|
||||
),
|
||||
memory_key_input_ids=memory_key_input_ids,
|
||||
memory_key_attention_mask=memory_key_attention_mask.to(device),
|
||||
memory_storage_input_ids=memory_storage["input_ids"].to(device),
|
||||
memory_storage_attention_mask=memory_storage.get(
|
||||
"attention_mask",
|
||||
torch.ones_like(memory_storage["input_ids"]),
|
||||
).to(device),
|
||||
)
|
||||
output_ids = model.generate(
|
||||
**encoded,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
memory_query_input_ids=memory_query_input_ids,
|
||||
memory_query_attention_mask=memory_query_attention_mask,
|
||||
)
|
||||
response_ids = output_ids[0, encoded["input_ids"].shape[1] :]
|
||||
print(f"AI> {tokenizer.decode(response_ids, skip_special_tokens=True)}")
|
||||
save_state()
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
print()
|
||||
finally:
|
||||
save_state()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,301 @@
|
||||
"""Compare original Qwen and Native Memory on general regression tasks."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from .qwen_integration import load_memory_config, load_qwen_base, load_qwen_dynamic, load_tokenizer
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-path", default=".")
|
||||
parser.add_argument("--adapter", default="dynamic_memory_lab/qwen_memory_adapter_native_v3")
|
||||
parser.add_argument("--data", default="dynamic_memory_lab/data/comprehensive_general.jsonl")
|
||||
parser.add_argument("--output", default="dynamic_memory_lab/comprehensive_benchmark_native_v3.json")
|
||||
parser.add_argument("--max-new-tokens", type=int, default=32)
|
||||
parser.add_argument("--perf-repeats", type=int, default=3)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
parser.add_argument("--category-drop-limit", type=float, default=0.10)
|
||||
parser.add_argument("--overall-drop-limit", type=float, default=0.05)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def load_cases(path: str | Path) -> list[dict[str, Any]]:
|
||||
cases = []
|
||||
for line in Path(path).read_text(encoding="utf-8").splitlines():
|
||||
if line.strip():
|
||||
cases.append(json.loads(line))
|
||||
if not cases:
|
||||
raise ValueError(f"no benchmark cases found in {path}")
|
||||
return cases
|
||||
|
||||
|
||||
def normalize(text: str) -> str:
|
||||
return re.sub(r"[\s`*_#，。！？、；：,.!?;:'\"（）()\[\]{}]", "", text).lower()
|
||||
|
||||
|
||||
def contains_answer(text: str, acceptable: list[str]) -> bool:
|
||||
normalized = normalize(text)
|
||||
for answer in acceptable:
|
||||
expected = normalize(str(answer))
|
||||
if not expected:
|
||||
continue
|
||||
if expected.isdigit() and len(expected) == 1:
|
||||
if re.search(rf"(?<!\d){re.escape(expected)}(?!\d)", normalized):
|
||||
return True
|
||||
elif expected in normalized:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def prompt_inputs(tokenizer: Any, prompt: str, device: torch.device) -> dict[str, torch.Tensor]:
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
[{"role": "user", "content": prompt}],
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
input_ids = encoded["input_ids"] if isinstance(encoded, dict) or hasattr(encoded, "__getitem__") else encoded
|
||||
if isinstance(input_ids, torch.Tensor):
|
||||
if input_ids.ndim == 1:
|
||||
input_ids = input_ids.unsqueeze(0)
|
||||
input_ids = input_ids.to(device)
|
||||
else:
|
||||
if input_ids and isinstance(input_ids[0], list):
|
||||
input_ids = input_ids[0]
|
||||
input_ids = torch.tensor([input_ids], dtype=torch.long, device=device)
|
||||
return {"input_ids": input_ids, "attention_mask": torch.ones_like(input_ids)}
|
||||
|
||||
|
||||
def _decode_generation(tokenizer: Any, output: torch.Tensor, prompt: dict[str, torch.Tensor]) -> str:
|
||||
start = prompt["input_ids"].shape[1]
|
||||
return tokenizer.decode(output[0, start:].detach().cpu().tolist(), skip_special_tokens=True).strip()
|
||||
|
||||
|
||||
def evaluate_model(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
cases: list[dict[str, Any]],
|
||||
*,
|
||||
adapted: bool,
|
||||
max_new_tokens: int,
|
||||
perf_repeats: int,
|
||||
) -> dict[str, Any]:
|
||||
device = model._find_layer_device() if adapted else model.get_input_embeddings().weight.device
|
||||
rows: list[dict[str, Any]] = []
|
||||
category_values: dict[str, list[float]] = defaultdict(list)
|
||||
category_examples: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
started = time.perf_counter()
|
||||
for case in cases:
|
||||
if adapted:
|
||||
model.reset_memory()
|
||||
prompt = prompt_inputs(tokenizer, str(case["prompt"]), device)
|
||||
with torch.inference_mode():
|
||||
if adapted:
|
||||
output = model.generate(
|
||||
**prompt,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
else:
|
||||
output = model.generate(
|
||||
**prompt,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
generated = _decode_generation(tokenizer, output, prompt)
|
||||
passed = contains_answer(generated, list(case["acceptable"]))
|
||||
category = str(case["category"])
|
||||
category_values[category].append(float(passed))
|
||||
if len(category_examples[category]) < 3:
|
||||
category_examples[category].append(
|
||||
{
|
||||
"id": case["id"],
|
||||
"prompt": case["prompt"],
|
||||
"acceptable": case["acceptable"],
|
||||
"generated": generated,
|
||||
"passed": passed,
|
||||
}
|
||||
)
|
||||
rows.append(
|
||||
{
|
||||
"id": case["id"],
|
||||
"category": category,
|
||||
"acceptable": case["acceptable"],
|
||||
"generated": generated,
|
||||
"passed": passed,
|
||||
}
|
||||
)
|
||||
elapsed = time.perf_counter() - started
|
||||
|
||||
perf_case = cases[0]
|
||||
if adapted:
|
||||
model.reset_memory()
|
||||
perf_prompt = prompt_inputs(tokenizer, str(perf_case["prompt"]), device)
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize(device)
|
||||
perf_start = time.perf_counter()
|
||||
generated_tokens = 0
|
||||
for _ in range(max(1, perf_repeats)):
|
||||
if adapted:
|
||||
model.reset_memory()
|
||||
with torch.inference_mode():
|
||||
if adapted:
|
||||
perf_output = model.generate(
|
||||
**perf_prompt,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
else:
|
||||
perf_output = model.generate(
|
||||
**perf_prompt,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
generated_tokens += int(perf_output.shape[1] - perf_prompt["input_ids"].shape[1])
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize(device)
|
||||
perf_elapsed = time.perf_counter() - perf_start
|
||||
score = sum(sum(values) for values in category_values.values()) / len(cases)
|
||||
return {
|
||||
"device": str(device),
|
||||
"cases": len(cases),
|
||||
"elapsed_seconds": elapsed,
|
||||
"overall_score": score,
|
||||
"categories": {
|
||||
category: {
|
||||
"count": len(values),
|
||||
"score": sum(values) / len(values),
|
||||
"examples": category_examples[category],
|
||||
}
|
||||
for category, values in sorted(category_values.items())
|
||||
},
|
||||
"performance": {
|
||||
"repeats": max(1, perf_repeats),
|
||||
"max_new_tokens": max_new_tokens,
|
||||
"tokens_per_second": generated_tokens / max(perf_elapsed, 1e-9),
|
||||
"seconds_per_run": perf_elapsed / max(1, perf_repeats),
|
||||
},
|
||||
"rows": rows,
|
||||
}
|
||||
|
||||
|
||||
def release(model: Any) -> None:
|
||||
del model
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
if hasattr(sys.stdout, "reconfigure"):
|
||||
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
|
||||
args = parse_args()
|
||||
cases = load_cases(args.data)
|
||||
tokenizer = load_tokenizer(args.model_path)
|
||||
use_4bit = not args.no_4bit
|
||||
report: dict[str, Any] = {
|
||||
"model_path": str(Path(args.model_path).resolve()),
|
||||
"adapter": str(Path(args.adapter).resolve()),
|
||||
"data": str(Path(args.data).resolve()),
|
||||
"quantization": "4bit_nf4" if use_4bit else "none",
|
||||
"cases": len(cases),
|
||||
}
|
||||
|
||||
print("loading baseline")
|
||||
baseline = load_qwen_base(args.model_path, load_in_4bit=use_4bit)
|
||||
baseline.eval()
|
||||
baseline_device = baseline.get_input_embeddings().weight.device
|
||||
if baseline_device.type == "cuda":
|
||||
torch.cuda.reset_peak_memory_stats(baseline_device)
|
||||
report["baseline"] = evaluate_model(
|
||||
baseline,
|
||||
tokenizer,
|
||||
cases,
|
||||
adapted=False,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
perf_repeats=args.perf_repeats,
|
||||
)
|
||||
if baseline_device.type == "cuda":
|
||||
report["baseline"]["peak_memory_gb"] = torch.cuda.max_memory_allocated(baseline_device) / 1024**3
|
||||
release(baseline)
|
||||
|
||||
config = load_memory_config(args.adapter)
|
||||
config.persistent_memory = False
|
||||
print(f"loading native adapter mode={config.mode} layers={config.layer_indices}")
|
||||
adapted = load_qwen_dynamic(
|
||||
args.model_path,
|
||||
memory_config=config,
|
||||
load_in_4bit=use_4bit,
|
||||
)
|
||||
adapted.load_memory_adapter(args.adapter)
|
||||
adapted.eval()
|
||||
adapted_device = adapted._find_layer_device()
|
||||
if adapted_device.type == "cuda":
|
||||
torch.cuda.reset_peak_memory_stats(adapted_device)
|
||||
report["native_memory"] = evaluate_model(
|
||||
adapted,
|
||||
tokenizer,
|
||||
cases,
|
||||
adapted=True,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
perf_repeats=args.perf_repeats,
|
||||
)
|
||||
if adapted_device.type == "cuda":
|
||||
report["native_memory"]["peak_memory_gb"] = torch.cuda.max_memory_allocated(adapted_device) / 1024**3
|
||||
|
||||
baseline_score = report["baseline"]["overall_score"]
|
||||
native_score = report["native_memory"]["overall_score"]
|
||||
baseline_categories = report["baseline"]["categories"]
|
||||
native_categories = report["native_memory"]["categories"]
|
||||
category_deltas = {
|
||||
category: native_categories[category]["score"] - baseline_categories[category]["score"]
|
||||
for category in baseline_categories.keys() & native_categories.keys()
|
||||
}
|
||||
report["regression"] = {
|
||||
"overall_delta": native_score - baseline_score,
|
||||
"category_deltas": category_deltas,
|
||||
"overall_drop_limit": args.overall_drop_limit,
|
||||
"category_drop_limit": args.category_drop_limit,
|
||||
"overall_regression_alert": native_score < baseline_score - args.overall_drop_limit,
|
||||
"category_regression_alerts": {
|
||||
category: delta < -args.category_drop_limit for category, delta in category_deltas.items()
|
||||
},
|
||||
"pass": native_score >= baseline_score - args.overall_drop_limit
|
||||
and all(delta >= -args.category_drop_limit for delta in category_deltas.values()),
|
||||
}
|
||||
legacy_path = Path(args.adapter).parent / "benchmark_qwen_native_v3_eval.json"
|
||||
if legacy_path.exists():
|
||||
report["same_memory_benchmark"] = json.loads(legacy_path.read_text(encoding="utf-8"))
|
||||
native_eval_path = Path(args.adapter) / "native_eval_report.json"
|
||||
if native_eval_path.exists():
|
||||
report["native_memory_holdout"] = json.loads(native_eval_path.read_text(encoding="utf-8"))
|
||||
output = Path(args.output)
|
||||
output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(json.dumps({key: report[key] for key in ("baseline", "native_memory", "regression")}, ensure_ascii=False, indent=2))
|
||||
print(f"saved={output}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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()
|
||||
@@ -0,0 +1,57 @@
|
||||
"""Debug the raw token pointer memory path on one benchmark record."""
|
||||
|
||||
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:
|
||||
config = QwenMemoryConfig(
|
||||
mode="blend",
|
||||
blend_init=0.1,
|
||||
write_token_offset=4,
|
||||
broadcast_write=True,
|
||||
raw_token_write=True,
|
||||
raw_logit_scale=30.0,
|
||||
)
|
||||
model = load_qwen_dynamic(".", memory_config=config, load_in_4bit=True)
|
||||
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))
|
||||
print("memory_ids", memory_input.tolist())
|
||||
print("memory_mask", memory_mask.tolist())
|
||||
prompt = _generation_prompt(tokenizer, record["query"][:-1], device)
|
||||
print("prompt_ids", prompt["input_ids"].tolist())
|
||||
with torch.inference_mode():
|
||||
output = 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,
|
||||
)
|
||||
print("raw_memory_shape", tuple(model.runtime.raw_memory.shape))
|
||||
print("raw_memory_norm", float(model.runtime.raw_memory.float().norm()))
|
||||
generated = model.generate(
|
||||
**prompt,
|
||||
max_new_tokens=1,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
print("generated", generated.tolist())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,49 @@
|
||||
"""Evaluate a saved dynamic-memory checkpoint."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
|
||||
import torch
|
||||
|
||||
from .model import DynamicMemoryConfig, DynamicMemoryLM
|
||||
from .tasks import sample_associative_batch
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--checkpoint", default="dynamic_memory_lab/checkpoints/latest.pt")
|
||||
parser.add_argument("--batches", type=int, default=100)
|
||||
parser.add_argument("--device", default="auto", choices=("auto", "cpu", "cuda"))
|
||||
args = parser.parse_args()
|
||||
|
||||
device = torch.device("cuda" if args.device == "auto" and torch.cuda.is_available() else ("cpu" if args.device == "auto" else args.device))
|
||||
checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=False)
|
||||
config = DynamicMemoryConfig(**checkpoint["config"])
|
||||
model = DynamicMemoryLM(config).to(device)
|
||||
model.load_state_dict(checkpoint["model"])
|
||||
model.eval()
|
||||
|
||||
for overwrite in (False, True):
|
||||
correct = 0
|
||||
total = 0
|
||||
with torch.no_grad():
|
||||
for _ in range(args.batches):
|
||||
batch = sample_associative_batch(
|
||||
batch_size=256,
|
||||
vocab_size=config.vocab_size,
|
||||
device=device,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
memory = None
|
||||
for chunk in batch.learn_chunks:
|
||||
memory = model(chunk, memory=memory, update_memory=True).memory
|
||||
output = model(batch.query_input, memory=memory, update_memory=False)
|
||||
prediction = output.logits[:, 0].argmax(dim=-1)
|
||||
correct += int((prediction == batch.expected).sum())
|
||||
total += batch.expected.numel()
|
||||
print(f"overwrite={overwrite} accuracy={correct / total:.3f}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,281 @@
|
||||
"""Evaluate the learned native memory controller on held-out streaming records."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import math
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from .qwen_integration import load_memory_config, load_qwen_dynamic, load_tokenizer
|
||||
from .train_qwen_memory import encode_messages, pad_batch
|
||||
from .train_native_memory import load_records
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-path", default=".")
|
||||
parser.add_argument("--adapter-dir", default="dynamic_memory_lab/qwen_memory_adapter_native")
|
||||
parser.add_argument("--data", default="dynamic_memory_lab/data/native_memory/eval.jsonl")
|
||||
parser.add_argument("--max-length", type=int, default=192)
|
||||
parser.add_argument("--max-new-tokens", type=int, default=8)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
parser.add_argument("--report", default=None)
|
||||
parser.add_argument("--restart-test", action="store_true")
|
||||
parser.add_argument("--limit", type=int, default=None)
|
||||
parser.add_argument(
|
||||
"--direct-logit-scale-override",
|
||||
type=float,
|
||||
default=None,
|
||||
help="temporarily override the adapter scale for generation diagnostics",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _mean(values: list[float]) -> float:
|
||||
return sum(values) / len(values) if values else float("nan")
|
||||
|
||||
|
||||
def _is_refusal(text: str) -> bool:
|
||||
# The base model may phrase an abstention as "没有相关记录" or "没有访问
|
||||
# 权限" rather than the exact training answer "不知道。". Count these
|
||||
# as safe abstentions; the report also keeps the raw generation.
|
||||
return any(marker in text for marker in ("不知道", "没有", "无相关", "未找到", "不清楚", "不确定", "无法", "不能"))
|
||||
|
||||
|
||||
def _encode_prompt(tokenizer: Any, messages: list[dict[str, Any]]) -> torch.Tensor:
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
enable_thinking=False,
|
||||
return_tensors="pt",
|
||||
)
|
||||
if hasattr(encoded, "input_ids"):
|
||||
return encoded.input_ids
|
||||
if isinstance(encoded, dict):
|
||||
return encoded["input_ids"]
|
||||
if encoded and isinstance(encoded[0], list):
|
||||
encoded = encoded[0]
|
||||
return torch.tensor([encoded], dtype=torch.long)
|
||||
|
||||
|
||||
def _query_forward(model: Any, tokenizer: Any, record: dict[str, Any], max_length: int) -> dict[str, float]:
|
||||
query_item = encode_messages(tokenizer, record["query"], max_length)
|
||||
query_input, query_mask, query_labels = pad_batch([query_item], int(tokenizer.pad_token_id))
|
||||
device = model._find_layer_device()
|
||||
output = model(
|
||||
input_ids=query_input.to(device),
|
||||
attention_mask=query_mask.to(device),
|
||||
labels=query_labels.to(device),
|
||||
read_memory=True,
|
||||
update_memory=False,
|
||||
return_memory=True,
|
||||
use_cache=False,
|
||||
)
|
||||
logits = output.logits[..., :-1, :]
|
||||
labels = query_labels.to(logits.device)[..., 1:]
|
||||
valid = labels.ne(-100)
|
||||
predictions = logits.argmax(dim=-1)
|
||||
token_accuracy = float((predictions[valid] == labels[valid]).float().mean()) if bool(valid.any()) else float("nan")
|
||||
return {"loss": float(output.loss.detach()), "token_accuracy": token_accuracy}
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def _generate_query(model: Any, tokenizer: Any, record: dict[str, Any], max_new_tokens: int) -> str:
|
||||
prompt = _encode_prompt(tokenizer, record["query"][:-1])
|
||||
device = model._find_layer_device()
|
||||
generated = model.generate(
|
||||
input_ids=prompt.to(device),
|
||||
attention_mask=torch.ones_like(prompt, device=device),
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
use_cache=False,
|
||||
update_memory=False,
|
||||
)
|
||||
new_tokens = generated[:, prompt.shape[1] :]
|
||||
return tokenizer.decode(new_tokens[0].detach().cpu().tolist(), skip_special_tokens=True).strip()
|
||||
|
||||
|
||||
def _controller_step(model: Any, tokenizer: Any, chunk: dict[str, Any], max_length: int) -> tuple[float, float]:
|
||||
item = encode_messages(tokenizer, chunk["messages"], max_length)
|
||||
inputs, mask, _ = pad_batch([item], int(tokenizer.pad_token_id))
|
||||
device = model._find_layer_device()
|
||||
model(
|
||||
input_ids=inputs.to(device),
|
||||
attention_mask=mask.to(device),
|
||||
read_memory=False,
|
||||
update_memory=True,
|
||||
return_memory=True,
|
||||
use_cache=False,
|
||||
)
|
||||
write = model.memory.last_write_probability
|
||||
forget = model.memory.last_forget_probability
|
||||
if write is None or forget is None:
|
||||
raise RuntimeError("adapter does not expose native controller probabilities")
|
||||
return float(write.detach().mean()), float(forget.detach().mean())
|
||||
|
||||
|
||||
def evaluate_records(model: Any, tokenizer: Any, records: list[dict[str, Any]], args: argparse.Namespace) -> dict[str, Any]:
|
||||
controller_rows: list[dict[str, Any]] = []
|
||||
query_rows: list[dict[str, Any]] = []
|
||||
generation_rows: list[dict[str, Any]] = []
|
||||
by_kind: dict[str, dict[str, list[float]]] = defaultdict(lambda: defaultdict(list))
|
||||
|
||||
for record in records:
|
||||
model.reset_memory()
|
||||
for chunk_index, chunk in enumerate(record["memory_chunks"]):
|
||||
write, forget = _controller_step(model, tokenizer, chunk, args.max_length)
|
||||
write_target = float(chunk.get("write_label", 1.0))
|
||||
forget_target = float(chunk.get("forget_label", 0.0))
|
||||
kind = str(chunk.get("kind", "unknown"))
|
||||
write_correct = float((write >= 0.5) == (write_target >= 0.5))
|
||||
forget_correct = float((forget >= 0.5) == (forget_target >= 0.5))
|
||||
row = {
|
||||
"record_id": record.get("id"),
|
||||
"chunk_index": chunk_index,
|
||||
"kind": kind,
|
||||
"write_probability": write,
|
||||
"write_target": write_target,
|
||||
"write_correct": write_correct,
|
||||
"forget_probability": forget,
|
||||
"forget_target": forget_target,
|
||||
"forget_correct": forget_correct,
|
||||
}
|
||||
controller_rows.append(row)
|
||||
by_kind[kind]["write_correct"].append(write_correct)
|
||||
by_kind[kind]["forget_correct"].append(forget_correct)
|
||||
by_kind[kind]["write_probability"].append(write)
|
||||
by_kind[kind]["forget_probability"].append(forget)
|
||||
|
||||
query = _query_forward(model, tokenizer, record, args.max_length)
|
||||
generated = _generate_query(model, tokenizer, record, args.max_new_tokens)
|
||||
answer = str(record.get("answer", ""))
|
||||
answerable = bool(record.get("answerable", False))
|
||||
contains_answer = bool(answer) and answer in generated if answerable else False
|
||||
says_unknown = _is_refusal(generated)
|
||||
generation_rows.append(
|
||||
{
|
||||
"record_id": record.get("id"),
|
||||
"answerable": answerable,
|
||||
"answer": answer,
|
||||
"generated": generated,
|
||||
"contains_answer": contains_answer,
|
||||
"says_unknown": says_unknown,
|
||||
}
|
||||
)
|
||||
query_rows.append({"record_id": record.get("id"), **query})
|
||||
|
||||
controller_metrics = {
|
||||
"write_accuracy": _mean([row["write_correct"] for row in controller_rows]),
|
||||
"forget_accuracy": _mean([row["forget_correct"] for row in controller_rows]),
|
||||
"write_bce_proxy": _mean([
|
||||
-(row["write_target"] * math.log(max(row["write_probability"], 1e-7))
|
||||
+ (1.0 - row["write_target"]) * math.log(max(1.0 - row["write_probability"], 1e-7))
|
||||
)
|
||||
for row in controller_rows
|
||||
]),
|
||||
"forget_bce_proxy": _mean([
|
||||
-(row["forget_target"] * math.log(max(row["forget_probability"], 1e-7))
|
||||
+ (1.0 - row["forget_target"]) * math.log(max(1.0 - row["forget_probability"], 1e-7))
|
||||
)
|
||||
for row in controller_rows
|
||||
]),
|
||||
"by_kind": {
|
||||
kind: {
|
||||
"count": len(values["write_correct"]),
|
||||
"write_accuracy": _mean(values["write_correct"]),
|
||||
"forget_accuracy": _mean(values["forget_correct"]),
|
||||
"write_probability": _mean(values["write_probability"]),
|
||||
"forget_probability": _mean(values["forget_probability"]),
|
||||
}
|
||||
for kind, values in by_kind.items()
|
||||
},
|
||||
}
|
||||
answerable_rows = [row for row in generation_rows if row["answerable"]]
|
||||
unknown_rows = [row for row in generation_rows if not row["answerable"]]
|
||||
query_metrics = {
|
||||
"mean_loss": _mean([row["loss"] for row in query_rows]),
|
||||
"token_accuracy": _mean([row["token_accuracy"] for row in query_rows]),
|
||||
}
|
||||
generation_metrics = {
|
||||
"answerable_count": len(answerable_rows),
|
||||
"answer_containment": _mean([float(row["contains_answer"]) for row in answerable_rows]),
|
||||
"unknown_count": len(unknown_rows),
|
||||
"unknown_refusal": _mean([float(row["says_unknown"]) for row in unknown_rows]),
|
||||
}
|
||||
return {
|
||||
"controller": controller_metrics,
|
||||
"query": query_metrics,
|
||||
"generation": generation_metrics,
|
||||
"controller_rows": controller_rows,
|
||||
"query_rows": query_rows,
|
||||
"generation_rows": generation_rows,
|
||||
}
|
||||
|
||||
|
||||
def restart_probe(model_path: str, adapter_dir: str, tokenizer: Any, record: dict[str, Any], args: argparse.Namespace) -> dict[str, Any]:
|
||||
"""Save state, recreate the model, and query without passing history."""
|
||||
config = load_memory_config(adapter_dir)
|
||||
config.persistent_memory = True
|
||||
if args.direct_logit_scale_override is not None:
|
||||
config.direct_logit_scale = args.direct_logit_scale_override
|
||||
first_model = load_qwen_dynamic(model_path, memory_config=config, load_in_4bit=not args.no_4bit)
|
||||
first_model.load_memory_adapter(adapter_dir)
|
||||
first_model.reset_memory()
|
||||
for chunk in record["memory_chunks"]:
|
||||
_controller_step(first_model, tokenizer, chunk, args.max_length)
|
||||
state_path = Path(adapter_dir) / "native_restart_probe_memory.pt"
|
||||
first_model.save_runtime_memory(state_path)
|
||||
del first_model
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
second_model = load_qwen_dynamic(model_path, memory_config=config, load_in_4bit=not args.no_4bit)
|
||||
second_model.load_memory_adapter(adapter_dir)
|
||||
second_model.load_runtime_memory(state_path)
|
||||
generated = _generate_query(second_model, tokenizer, record, args.max_new_tokens)
|
||||
answer = str(record.get("answer", ""))
|
||||
return {
|
||||
"record_id": record.get("id"),
|
||||
"generated_after_restart": generated,
|
||||
"answer": answer,
|
||||
"contains_answer": answer in generated if record.get("answerable") else False,
|
||||
"state_path": str(state_path),
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
records = load_records(args.data)
|
||||
if args.limit is not None:
|
||||
records = records[: args.limit]
|
||||
tokenizer = load_tokenizer(args.model_path)
|
||||
config = load_memory_config(args.adapter_dir)
|
||||
config.persistent_memory = True
|
||||
if args.direct_logit_scale_override is not None:
|
||||
config.direct_logit_scale = args.direct_logit_scale_override
|
||||
model = load_qwen_dynamic(args.model_path, memory_config=config, load_in_4bit=not args.no_4bit)
|
||||
model.load_memory_adapter(args.adapter_dir)
|
||||
model.eval()
|
||||
report = evaluate_records(model, tokenizer, records, args)
|
||||
if args.restart_test:
|
||||
probe_record = next(record for record in records if record.get("answerable"))
|
||||
report["restart_probe"] = restart_probe(args.model_path, args.adapter_dir, tokenizer, probe_record, args)
|
||||
report["model_path"] = str(args.model_path)
|
||||
report["adapter_dir"] = str(args.adapter_dir)
|
||||
report["data"] = str(args.data)
|
||||
report_path = Path(args.report) if args.report else Path(args.adapter_dir) / "native_eval_report.json"
|
||||
report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(json.dumps({key: report[key] for key in ("controller", "query", "generation", "restart_probe") if key in report}, ensure_ascii=False, indent=2))
|
||||
print(f"report={report_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -0,0 +1,264 @@
|
||||
"""Generate a large, deterministic validation set for KV-vs-memory parity.
|
||||
|
||||
Each case contains a complete fact episode, distractors, a query, and exact
|
||||
acceptance metadata. The default is 100,000 cases across ten categories.
|
||||
This is a data generator only; model quality must be measured by the paired
|
||||
teacher/student benchmark after generation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import random
|
||||
import string
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent
|
||||
CATEGORIES = (
|
||||
"single_fact",
|
||||
"distractor_32",
|
||||
"distractor_128",
|
||||
"conflict_update",
|
||||
"multi_hop",
|
||||
"unknown_abstention",
|
||||
"random_position",
|
||||
"long_context",
|
||||
"paraphrase",
|
||||
"forget_correction",
|
||||
)
|
||||
|
||||
ATTRIBUTES = (
|
||||
"常用语言",
|
||||
"工作地点",
|
||||
"最喜欢的水果",
|
||||
"项目代号",
|
||||
"常用时区",
|
||||
"提醒时间",
|
||||
"默认输出风格",
|
||||
"备用联系人",
|
||||
)
|
||||
UNKNOWN_MARKERS = ["不知道", "没有记录", "无法确定", "未找到相关信息"]
|
||||
|
||||
|
||||
def _path(value: str | Path) -> Path:
|
||||
path = Path(value)
|
||||
if path.is_absolute() or path.exists():
|
||||
return path
|
||||
return PROJECT_ROOT / path
|
||||
|
||||
|
||||
def _value(rng: random.Random, prefix: str) -> str:
|
||||
alphabet = string.ascii_uppercase + string.digits
|
||||
return f"{prefix}-" + "".join(rng.choice(alphabet) for _ in range(8))
|
||||
|
||||
|
||||
def _fact(
|
||||
text: str,
|
||||
*,
|
||||
entity: str,
|
||||
attribute: str,
|
||||
value: str,
|
||||
should_write: bool = True,
|
||||
kind: str = "fact",
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"text": text,
|
||||
"assistant": "好的，我会记录这条信息。",
|
||||
"entity": entity,
|
||||
"attribute": attribute,
|
||||
"value": value,
|
||||
"should_write": should_write,
|
||||
"kind": kind,
|
||||
}
|
||||
|
||||
|
||||
def _noise(rng: random.Random, subject: str, index: int) -> dict[str, Any]:
|
||||
attribute = f"干扰属性{index}"
|
||||
value = _value(rng, "N")
|
||||
return _fact(
|
||||
f"这是普通对话噪声：{subject}暂时提到一个无关编号{value}，不需要长期保存。",
|
||||
entity=subject,
|
||||
attribute=attribute,
|
||||
value=value,
|
||||
should_write=False,
|
||||
kind="noise",
|
||||
)
|
||||
|
||||
|
||||
def _base_case(category: str, index: int, rng: random.Random) -> dict[str, Any]:
|
||||
subject = f"验证用户-{index:07d}"
|
||||
attribute = rng.choice(ATTRIBUTES)
|
||||
value = _value(rng, "V")
|
||||
facts: list[dict[str, Any]] = []
|
||||
query = ""
|
||||
acceptable: list[str] = []
|
||||
forbidden: list[str] = []
|
||||
metadata: dict[str, Any] = {}
|
||||
|
||||
if category == "single_fact":
|
||||
facts = [_fact(f"请记住：我的{attribute}是{value}。", entity=subject, attribute=attribute, value=value)]
|
||||
query = f"跨对话后请回答：我的{attribute}是什么？"
|
||||
acceptable = [value]
|
||||
elif category in {"distractor_32", "distractor_128"}:
|
||||
count = 32 if category.endswith("32") else 128
|
||||
target = _fact(f"请记住：我的{attribute}是{value}。", entity=subject, attribute=attribute, value=value)
|
||||
facts = [_noise(rng, subject, item) for item in range(count)]
|
||||
facts.insert(rng.randrange(len(facts) + 1), target)
|
||||
query = f"在大量无关信息中，只读取我的{attribute}，它的值是什么？"
|
||||
acceptable = [value]
|
||||
metadata["distractor_count"] = count
|
||||
elif category == "conflict_update":
|
||||
old_value = _value(rng, "OLD")
|
||||
facts = [
|
||||
_fact(f"我的{attribute}是{old_value}。", entity=subject, attribute=attribute, value=old_value),
|
||||
_fact(
|
||||
f"更正一下：我的{attribute}已经改为{value}，旧值不要再使用。",
|
||||
entity=subject,
|
||||
attribute=attribute,
|
||||
value=value,
|
||||
kind="correction",
|
||||
),
|
||||
]
|
||||
query = f"我的{attribute}最新值是什么？"
|
||||
acceptable = [value]
|
||||
forbidden = [old_value]
|
||||
elif category == "multi_hop":
|
||||
project = f"项目-{_value(rng, 'P')}"
|
||||
person = f"成员-{_value(rng, 'M')}"
|
||||
code = _value(rng, "H")
|
||||
facts = [
|
||||
_fact(f"项目{project}的负责人是{person}。", entity=project, attribute="负责人", value=person),
|
||||
_fact(f"成员{person}的工作代号是{code}。", entity=person, attribute="工作代号", value=code),
|
||||
]
|
||||
query = f"请通过项目负责人关系，找出项目{project}负责人的工作代号。"
|
||||
acceptable = [code]
|
||||
metadata["hop_count"] = 2
|
||||
elif category == "unknown_abstention":
|
||||
known_attribute = rng.choice(ATTRIBUTES)
|
||||
missing_attribute = f"不存在的个人属性-{index:07d}"
|
||||
known_value = _value(rng, "KNOWN")
|
||||
facts = [_fact(f"我的{known_attribute}是{known_value}。", entity=subject, attribute=known_attribute, value=known_value)]
|
||||
query = f"我的{missing_attribute}是什么？如果没有记录，请明确说不知道。"
|
||||
acceptable = UNKNOWN_MARKERS
|
||||
forbidden = [known_value]
|
||||
metadata["answerable"] = False
|
||||
elif category == "random_position":
|
||||
count = 64
|
||||
facts = [_noise(rng, subject, item) for item in range(count)]
|
||||
target = _fact(f"请记住：我的{attribute}是{value}。", entity=subject, attribute=attribute, value=value)
|
||||
position = rng.randrange(count + 1)
|
||||
facts.insert(position, target)
|
||||
query = f"随机位置事实测试：我的{attribute}是什么？"
|
||||
acceptable = [value]
|
||||
metadata["target_position"] = position
|
||||
metadata["distractor_count"] = count
|
||||
elif category == "long_context":
|
||||
count = 96
|
||||
facts = [_noise(rng, subject, item) for item in range(count)]
|
||||
target = _fact(f"请记住：我的{attribute}是{value}。", entity=subject, attribute=attribute, value=value)
|
||||
position = rng.randrange(count + 1)
|
||||
facts.insert(position, target)
|
||||
query = f"在长上下文压缩之后，检索我的{attribute}并回答。"
|
||||
acceptable = [value]
|
||||
metadata["target_position"] = position
|
||||
metadata["distractor_count"] = count
|
||||
metadata["synthetic_padding_tokens"] = 2048 + (index % 5) * 512
|
||||
elif category == "paraphrase":
|
||||
fact_templates = (
|
||||
f"个人资料更新：{attribute}这一栏填写为{value}。",
|
||||
f"以后涉及{attribute}时，请使用{value}这个值。",
|
||||
f"记录一下，我的{attribute}偏好/设置是{value}。",
|
||||
)
|
||||
query_templates = (
|
||||
f"我之前登记的{attribute}内容是什么？",
|
||||
f"关于{attribute}，你保存的用户信息是哪一个？",
|
||||
f"不要猜，回忆一下我在{attribute}上的设置。",
|
||||
)
|
||||
facts = [_fact(rng.choice(fact_templates), entity=subject, attribute=attribute, value=value)]
|
||||
query = rng.choice(query_templates)
|
||||
acceptable = [value]
|
||||
elif category == "forget_correction":
|
||||
old_value = _value(rng, "FORGET")
|
||||
facts = [
|
||||
_fact(f"请记住：我的{attribute}是{old_value}。", entity=subject, attribute=attribute, value=old_value),
|
||||
_fact(
|
||||
f"请删除关于我的{attribute}的记忆，不要再保留这个信息。",
|
||||
entity=subject,
|
||||
attribute=attribute,
|
||||
value=old_value,
|
||||
should_write=False,
|
||||
kind="forget",
|
||||
),
|
||||
]
|
||||
query = f"我的{attribute}是什么？如果已经删除，请回答没有记录。"
|
||||
acceptable = UNKNOWN_MARKERS
|
||||
forbidden = [old_value]
|
||||
metadata["answerable"] = False
|
||||
else:
|
||||
raise ValueError(category)
|
||||
|
||||
return {
|
||||
"id": f"mega-{category}-{index:07d}",
|
||||
"category": category,
|
||||
"subject": subject,
|
||||
"facts": facts,
|
||||
"query": query,
|
||||
"acceptable": acceptable,
|
||||
"forbidden": forbidden,
|
||||
"metadata": metadata,
|
||||
"generator_version": 1,
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--output", default="data/mega_validation/memory_validation_100k.jsonl")
|
||||
parser.add_argument("--cases-per-category", type=int, default=10000)
|
||||
parser.add_argument("--seed", type=int, default=20260905)
|
||||
args = parser.parse_args()
|
||||
if args.cases_per_category < 1:
|
||||
raise SystemExit("--cases-per-category must be positive")
|
||||
output = _path(args.output)
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
counts: Counter[str] = Counter()
|
||||
digest = hashlib.sha256()
|
||||
total = 0
|
||||
with output.open("w", encoding="utf-8") as handle:
|
||||
for category_index, category in enumerate(CATEGORIES):
|
||||
for index in range(args.cases_per_category):
|
||||
case_seed = args.seed + category_index * 1_000_003 + index * 97
|
||||
case = _base_case(category, index, random.Random(case_seed))
|
||||
raw = (json.dumps(case, ensure_ascii=False, separators=(",", ":")) + "\n").encode("utf-8")
|
||||
handle.write(raw.decode("utf-8"))
|
||||
digest.update(raw)
|
||||
counts[category] += 1
|
||||
total += 1
|
||||
manifest = {
|
||||
"format_version": 1,
|
||||
"generator": "make_mega_memory_validation.py",
|
||||
"generator_version": 1,
|
||||
"seed": args.seed,
|
||||
"categories": list(CATEGORIES),
|
||||
"cases_per_category": args.cases_per_category,
|
||||
"total_cases": total,
|
||||
"counts": dict(counts),
|
||||
"sha256": digest.hexdigest(),
|
||||
"teacher_student_protocol": {
|
||||
"teacher": "complete fact episode plus query in one full KV context",
|
||||
"student": "facts presented one turn at a time, memory read for query, no history replay",
|
||||
"primary_gate": "student answer agreement >= 0.95 * teacher answer accuracy",
|
||||
},
|
||||
"warning": "Synthetic stress validation; add redacted real conversations before production certification.",
|
||||
}
|
||||
manifest_path = output.with_suffix(".manifest.json")
|
||||
manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(json.dumps({"output": str(output), "manifest": str(manifest_path), **manifest}, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,156 @@
|
||||
"""Create train/eval streams for learned write, forget, and no-hallucination tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import random
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
ATTRIBUTES = ("常用语言", "备用联系人", "工作区域", "档案代号", "提醒时间")
|
||||
VALUES = tuple("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789")
|
||||
|
||||
FACT_USER = (
|
||||
"请记住：{subject}的{attribute}是代号{value}。",
|
||||
"把这条个人资料写入记忆：{subject}的{attribute}为{value}。",
|
||||
"个人事实更新——{subject}的{attribute}等于{value}，以后可能会问到。",
|
||||
)
|
||||
FACT_ASSISTANT = (
|
||||
"好的，这条资料已经记录。",
|
||||
"收到，我会保留这条个人事实。",
|
||||
"已保存。",
|
||||
)
|
||||
NOISE_USER = (
|
||||
"今天天气不错，随便聊聊。",
|
||||
"请给我一个简短的问候。",
|
||||
"这是一条不需要长期记忆的闲聊。",
|
||||
)
|
||||
NOISE_ASSISTANT = (
|
||||
"好的。",
|
||||
"明白。",
|
||||
"收到。",
|
||||
)
|
||||
QUERY_KNOWN = (
|
||||
"只根据已经保存的个人资料，{subject}的{attribute}是什么？",
|
||||
"不要猜测，请读取记忆回答：{subject}的{attribute}为？",
|
||||
"跨对话查询：请问{subject}的{attribute}代号是什么？",
|
||||
)
|
||||
QUERY_UNKNOWN = (
|
||||
"记忆中是否有{subject}的{attribute}？如果没有，请明确说不知道。",
|
||||
"请查询个人记忆：{subject}的{attribute}是什么？没有记录时不要猜。",
|
||||
)
|
||||
|
||||
|
||||
def _messages(user: str, assistant: str) -> list[dict[str, str]]:
|
||||
return [{"role": "user", "content": user}, {"role": "assistant", "content": assistant}]
|
||||
|
||||
|
||||
def _fact_chunk(subject: str, attribute: str, value: str, rng: random.Random, *, forget: int = 0) -> dict:
|
||||
fields = {"subject": subject, "attribute": attribute, "value": value}
|
||||
return {
|
||||
"messages": _messages(rng.choice(FACT_USER).format(**fields), rng.choice(FACT_ASSISTANT)),
|
||||
"value": value,
|
||||
"write_label": 1.0,
|
||||
"forget_label": float(forget),
|
||||
"kind": "fact" if not forget else "replacement",
|
||||
}
|
||||
|
||||
|
||||
def _noise_chunk(rng: random.Random) -> dict:
|
||||
return {
|
||||
"messages": _messages(rng.choice(NOISE_USER), rng.choice(NOISE_ASSISTANT)),
|
||||
"value": None,
|
||||
"write_label": 0.0,
|
||||
"forget_label": 0.0,
|
||||
"kind": "noise",
|
||||
}
|
||||
|
||||
|
||||
def make_record(index: int, *, prefix: str, rng: random.Random) -> dict:
|
||||
subject = f"{prefix}{index:05d}"
|
||||
attribute = rng.choice(ATTRIBUTES)
|
||||
value = rng.choice(VALUES)
|
||||
mode = rng.random()
|
||||
chunks: list[dict] = []
|
||||
if mode < 0.20:
|
||||
chunks.append(_noise_chunk(rng))
|
||||
answer = "不知道。"
|
||||
query = rng.choice(QUERY_UNKNOWN).format(subject=subject, attribute=attribute)
|
||||
answerable = False
|
||||
elif mode < 0.45:
|
||||
old_value = rng.choice(tuple(item for item in VALUES if item != value))
|
||||
chunks.append(_fact_chunk(subject, attribute, old_value, rng))
|
||||
chunks.append(_noise_chunk(rng))
|
||||
chunks.append(_fact_chunk(subject, attribute, value, rng, forget=1))
|
||||
answer = value
|
||||
query = rng.choice(QUERY_KNOWN).format(subject=subject, attribute=attribute)
|
||||
answerable = True
|
||||
else:
|
||||
if rng.random() < 0.35:
|
||||
chunks.append(_noise_chunk(rng))
|
||||
chunks.append(_fact_chunk(subject, attribute, value, rng))
|
||||
if rng.random() < 0.35:
|
||||
chunks.append(_noise_chunk(rng))
|
||||
answer = value
|
||||
query = rng.choice(QUERY_KNOWN).format(subject=subject, attribute=attribute)
|
||||
answerable = True
|
||||
return {
|
||||
"id": f"{prefix.lower()}-{index:05d}",
|
||||
"memory_chunks": chunks,
|
||||
"query": _messages(query, answer),
|
||||
"subject": subject,
|
||||
"attribute": attribute,
|
||||
"answer": answer,
|
||||
"answerable": answerable,
|
||||
}
|
||||
|
||||
|
||||
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/native_memory")
|
||||
parser.add_argument("--train-count", type=int, default=512)
|
||||
parser.add_argument("--eval-count", type=int, default=128)
|
||||
parser.add_argument("--seed", type=int, default=20260904)
|
||||
args = parser.parse_args()
|
||||
train = [
|
||||
make_record(i, prefix="训练用户", rng=random.Random(args.seed + i * 17))
|
||||
for i in range(args.train_count)
|
||||
]
|
||||
evaluation = [
|
||||
make_record(i, prefix="评估用户", rng=random.Random(args.seed + 100000 + i * 17))
|
||||
for i in range(args.eval_count)
|
||||
]
|
||||
output_dir = Path(args.output_dir)
|
||||
write_jsonl(output_dir / "train.jsonl", train)
|
||||
write_jsonl(output_dir / "eval.jsonl", evaluation)
|
||||
(output_dir / "manifest.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"seed": args.seed,
|
||||
"train_count": len(train),
|
||||
"eval_count": len(evaluation),
|
||||
"task": "learned persistent memory with noise, unknowns, and replacement",
|
||||
"contains_write_labels": True,
|
||||
"contains_forget_labels": True,
|
||||
"answer_is_not_derived_from_subject": True,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
print(f"train={len(train)} path={output_dir / 'train.jsonl'}")
|
||||
print(f"eval={len(evaluation)} path={output_dir / 'eval.jsonl'}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+1922
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,307 @@
|
||||
"""Merge dynamic-memory tensors into a copy of a local safetensors shard.
|
||||
|
||||
The base Qwen shards are never modified. The output directory contains
|
||||
hardlinks for unchanged files and a newly written second shard with the
|
||||
memory tensors appended to its safetensors payload. The normal HF loader
|
||||
continues to see the original Qwen keys; ``qwen_integration`` loads the
|
||||
embedded ``dynamic_memory.*`` keys when it sees ``memory_merge.json``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import struct
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from safetensors.torch import save_file
|
||||
|
||||
|
||||
def _read_header(path: Path) -> tuple[dict[str, Any], int, int]:
|
||||
with path.open("rb") as handle:
|
||||
raw_length = handle.read(8)
|
||||
if len(raw_length) != 8:
|
||||
raise ValueError(f"not a safetensors file: {path}")
|
||||
header_length = struct.unpack("<Q", raw_length)[0]
|
||||
header_bytes = handle.read(header_length)
|
||||
if len(header_bytes) != header_length:
|
||||
raise ValueError(f"truncated safetensors header: {path}")
|
||||
return json.loads(header_bytes.decode("utf-8")), int(header_length), 8 + int(header_length)
|
||||
|
||||
|
||||
def _load_adapter_tensors(adapter_dir: Path, memory_state: Path | None) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
|
||||
tensors: dict[str, torch.Tensor] = {}
|
||||
memory = torch.load(adapter_dir / "memory.pt", map_location="cpu", weights_only=True)
|
||||
tensors.update({f"dynamic_memory.memory.{key}": value.detach().cpu().contiguous() for key, value in memory.items()})
|
||||
|
||||
retriever_path = adapter_dir / "text_retriever.pt"
|
||||
if retriever_path.exists():
|
||||
retriever = torch.load(retriever_path, map_location="cpu", weights_only=True)
|
||||
tensors.update(
|
||||
{f"dynamic_memory.text_retriever.{key}": value.detach().cpu().contiguous() for key, value in retriever.items()}
|
||||
)
|
||||
|
||||
policy_path = adapter_dir / "memory_policy.pt"
|
||||
if policy_path.exists():
|
||||
policy = torch.load(policy_path, map_location="cpu", weights_only=True)
|
||||
tensors.update(
|
||||
{f"dynamic_memory.memory_policy.{key}": value.detach().cpu().contiguous() for key, value in policy.items()}
|
||||
)
|
||||
|
||||
surgery_path = adapter_dir / "surgery.pt"
|
||||
if surgery_path.exists():
|
||||
surgery = torch.load(surgery_path, map_location="cpu", weights_only=True)
|
||||
for layer, value in surgery.get("blend_logits", {}).items():
|
||||
tensors[f"dynamic_memory.blend_logits.{layer}"] = value.detach().cpu().contiguous()
|
||||
|
||||
metadata = json.loads((adapter_dir / "memory_config.json").read_text(encoding="utf-8"))
|
||||
memory_config = dict(metadata.get("memory_config", {}))
|
||||
|
||||
persistent_source = memory_state
|
||||
if persistent_source is None and (adapter_dir / "persistent_memory.pt").exists():
|
||||
persistent_source = adapter_dir / "persistent_memory.pt"
|
||||
if persistent_source is not None:
|
||||
payload = torch.load(persistent_source, map_location="cpu", weights_only=True)
|
||||
if not isinstance(payload, dict) or not isinstance(payload.get("memory_state"), torch.Tensor):
|
||||
raise ValueError(f"memory state file has no tensor memory_state: {persistent_source}")
|
||||
for key, value in payload.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
tensors[f"dynamic_memory.persistent.{key}"] = value.detach().cpu().contiguous()
|
||||
memory_config["persistent_memory"] = True
|
||||
|
||||
if not tensors:
|
||||
raise ValueError("no tensors were found to merge")
|
||||
metadata["memory_config"] = memory_config
|
||||
metadata["checkpoint_contains_user_memory"] = persistent_source is not None
|
||||
return tensors, metadata
|
||||
|
||||
|
||||
def _write_merged_shard(base_shard: Path, extra_shard: Path, output_shard: Path) -> None:
|
||||
base_header, base_header_length, base_data_start = _read_header(base_shard)
|
||||
extra_header, _, extra_data_start = _read_header(extra_shard)
|
||||
if "__metadata__" in extra_header:
|
||||
extra_metadata = extra_header.pop("__metadata__")
|
||||
else:
|
||||
extra_metadata = {}
|
||||
if "__metadata__" in base_header:
|
||||
base_metadata = base_header.pop("__metadata__")
|
||||
else:
|
||||
base_metadata = {}
|
||||
collisions = set(base_header).intersection(extra_header)
|
||||
if collisions:
|
||||
raise ValueError(f"safetensors key collision while merging: {sorted(collisions)[:4]}")
|
||||
|
||||
base_payload_size = base_shard.stat().st_size - base_data_start
|
||||
merged_header: dict[str, Any] = {}
|
||||
new_metadata = dict(base_metadata)
|
||||
new_metadata.update({str(key): str(value) for key, value in extra_metadata.items()})
|
||||
merged_header["__metadata__"] = new_metadata
|
||||
|
||||
for key, entry in base_header.items():
|
||||
item = dict(entry)
|
||||
start, end = item["data_offsets"]
|
||||
item["data_offsets"] = [int(start), int(end)]
|
||||
merged_header[key] = item
|
||||
# Copy the entries: the offset-fixup loop mutates merged_header and must
|
||||
# not mutate extra_header, otherwise the second loop iteration applies the
|
||||
# base offset repeatedly.
|
||||
merged_header.update({key: dict(entry) for key, entry in extra_header.items()})
|
||||
|
||||
# Offsets are relative to the beginning of the data section, so changing
|
||||
# the header length does not move the base tensors in that coordinate
|
||||
# system. Recompute the header until its padded size is stable.
|
||||
while True:
|
||||
encoded = json.dumps(merged_header, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
|
||||
padded_length = (len(encoded) + 7) // 8 * 8
|
||||
for key, entry in base_header.items():
|
||||
start, end = entry["data_offsets"]
|
||||
merged_header[key]["data_offsets"] = [int(start), int(end)]
|
||||
# ``data_offsets`` are relative to the beginning of the data section,
|
||||
# not absolute file offsets. The extra payload follows the base
|
||||
# payload inside that merged data section.
|
||||
extra_start = base_payload_size
|
||||
for key, entry in extra_header.items():
|
||||
start, end = entry["data_offsets"]
|
||||
merged_header[key]["data_offsets"] = [int(extra_start + start), int(extra_start + end)]
|
||||
encoded_next = json.dumps(merged_header, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
|
||||
next_length = (len(encoded_next) + 7) // 8 * 8
|
||||
if next_length == padded_length:
|
||||
encoded = encoded_next + b" " * (next_length - len(encoded_next))
|
||||
break
|
||||
|
||||
output_shard.parent.mkdir(parents=True, exist_ok=True)
|
||||
with output_shard.open("wb") as output, base_shard.open("rb") as base, extra_shard.open("rb") as extra:
|
||||
output.write(struct.pack("<Q", len(encoded)))
|
||||
output.write(encoded)
|
||||
base.seek(base_data_start)
|
||||
shutil.copyfileobj(base, output, length=16 * 1024 * 1024)
|
||||
extra.seek(extra_data_start)
|
||||
shutil.copyfileobj(extra, output, length=16 * 1024 * 1024)
|
||||
|
||||
|
||||
def _hardlink_base_files(base_dir: Path, output_dir: Path, merged_shard_name: str) -> None:
|
||||
for source in base_dir.iterdir():
|
||||
if not source.is_file() or source.name in {"memory_config.json", "memory_merge.json", merged_shard_name}:
|
||||
continue
|
||||
destination = output_dir / source.name
|
||||
try:
|
||||
os.link(source, destination)
|
||||
except OSError:
|
||||
try:
|
||||
# Some Windows volumes reject hardlinks but allow symlinks;
|
||||
# this keeps the unchanged 5GB shard from being duplicated.
|
||||
os.symlink(source, destination)
|
||||
except OSError:
|
||||
# Last-resort fallback for filesystems that allow neither.
|
||||
# The large first shard is still copied only when necessary.
|
||||
shutil.copy2(source, destination)
|
||||
|
||||
|
||||
def merge_package(
|
||||
base_dir: Path,
|
||||
adapter_dir: Path,
|
||||
output_dir: Path,
|
||||
memory_state: Path | None,
|
||||
) -> dict[str, Any]:
|
||||
if output_dir.exists():
|
||||
raise FileExistsError(f"refusing to overwrite existing output: {output_dir}")
|
||||
base_index = base_dir / "model.safetensors.index.json"
|
||||
index = json.loads(base_index.read_text(encoding="utf-8"))
|
||||
shard_names = sorted(set(index["weight_map"].values()))
|
||||
if len(shard_names) != 2 or "model.safetensors-00002-of-00002.safetensors" not in shard_names:
|
||||
raise ValueError(f"expected the local two-shard Qwen layout, got {shard_names}")
|
||||
|
||||
output_dir.mkdir(parents=True)
|
||||
merged_shard_name = "model.safetensors-00002-of-00002.safetensors"
|
||||
_hardlink_base_files(base_dir, output_dir, merged_shard_name)
|
||||
|
||||
tensors, adapter_metadata = _load_adapter_tensors(adapter_dir, memory_state)
|
||||
extra_shard = output_dir / "memory_extra.safetensors"
|
||||
save_file(tensors, str(extra_shard), metadata={"format": "qwen_dynamic_memory_embedded_v1"})
|
||||
_write_merged_shard(
|
||||
base_dir / merged_shard_name,
|
||||
extra_shard,
|
||||
output_dir / merged_shard_name,
|
||||
)
|
||||
# This is a self-owned staging file; the merged shard is the only copy
|
||||
# that belongs in the output package.
|
||||
extra_shard.unlink()
|
||||
|
||||
(output_dir / "memory_config.json").write_text(
|
||||
json.dumps(adapter_metadata, indent=2, ensure_ascii=False),
|
||||
encoding="utf-8",
|
||||
)
|
||||
manifest = {
|
||||
"format_version": 1,
|
||||
"format": "qwen_dynamic_memory_merged_shard",
|
||||
"shard_mode": "combined",
|
||||
"memory_weights": merged_shard_name,
|
||||
"base_model": str(base_dir),
|
||||
"source_adapter": str(adapter_dir),
|
||||
"source_memory_state": str(memory_state) if memory_state is not None else None,
|
||||
"tensor_prefix": "dynamic_memory.",
|
||||
"checkpoint_contains_user_memory": bool(adapter_metadata.get("checkpoint_contains_user_memory")),
|
||||
"memory_config": adapter_metadata.get("memory_config", {}),
|
||||
}
|
||||
(output_dir / "memory_merge.json").write_text(
|
||||
json.dumps(manifest, indent=2, ensure_ascii=False),
|
||||
encoding="utf-8",
|
||||
)
|
||||
return manifest
|
||||
|
||||
|
||||
def add_memory_shard(
|
||||
base_dir: Path,
|
||||
adapter_dir: Path,
|
||||
shard_path: Path,
|
||||
memory_state: Path | None,
|
||||
) -> dict[str, Any]:
|
||||
"""Add a small extra-only memory shard beside an existing Qwen model."""
|
||||
|
||||
if shard_path.exists():
|
||||
raise FileExistsError(f"refusing to overwrite existing shard: {shard_path}")
|
||||
if not (base_dir / "config.json").exists() or not (base_dir / "model.safetensors.index.json").exists():
|
||||
raise ValueError(f"not a complete local Qwen model directory: {base_dir}")
|
||||
|
||||
tensors, adapter_metadata = _load_adapter_tensors(adapter_dir, memory_state)
|
||||
shard_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
save_file(
|
||||
tensors,
|
||||
str(shard_path),
|
||||
metadata={"format": "qwen_dynamic_memory_extra_shard_v1"},
|
||||
)
|
||||
index_path = base_dir / "model.safetensors.index.json"
|
||||
index = json.loads(index_path.read_text(encoding="utf-8"))
|
||||
weight_map = index.setdefault("weight_map", {})
|
||||
collisions = sorted(set(weight_map).intersection(tensors))
|
||||
if collisions:
|
||||
raise ValueError(f"weight-map key collision while adding memory shard: {collisions[:4]}")
|
||||
weight_map.update({key: shard_path.name for key in tensors})
|
||||
tensor_bytes = sum(value.numel() * value.element_size() for value in tensors.values())
|
||||
index.setdefault("metadata", {})["total_size"] = int(
|
||||
index.get("metadata", {}).get("total_size", 0) + tensor_bytes
|
||||
)
|
||||
index_path.write_text(json.dumps(index, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
|
||||
memory_config_path = base_dir / "memory_config.json"
|
||||
memory_config_path.write_text(
|
||||
json.dumps(adapter_metadata, indent=2, ensure_ascii=False),
|
||||
encoding="utf-8",
|
||||
)
|
||||
manifest = {
|
||||
"format_version": 1,
|
||||
"format": "qwen_dynamic_memory_extra_shard",
|
||||
"shard_mode": "extra_only",
|
||||
"memory_weights": shard_path.name,
|
||||
"base_model": str(base_dir),
|
||||
"source_adapter": str(adapter_dir),
|
||||
"source_memory_state": str(memory_state) if memory_state is not None else None,
|
||||
"tensor_prefix": "dynamic_memory.",
|
||||
"checkpoint_contains_user_memory": bool(adapter_metadata.get("checkpoint_contains_user_memory")),
|
||||
"memory_config": adapter_metadata.get("memory_config", {}),
|
||||
}
|
||||
(base_dir / "memory_merge.json").write_text(
|
||||
json.dumps(manifest, indent=2, ensure_ascii=False),
|
||||
encoding="utf-8",
|
||||
)
|
||||
return manifest
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--base-model", required=True)
|
||||
parser.add_argument("--adapter", required=True)
|
||||
parser.add_argument("--output", default=None)
|
||||
parser.add_argument("--add-shard", default=None, help="add a small extra-only shard beside the base model")
|
||||
parser.add_argument("--memory-state", default=None)
|
||||
args = parser.parse_args()
|
||||
if args.add_shard:
|
||||
manifest = add_memory_shard(
|
||||
Path(args.base_model),
|
||||
Path(args.adapter),
|
||||
Path(args.add_shard),
|
||||
Path(args.memory_state) if args.memory_state else None,
|
||||
)
|
||||
output_shard = Path(args.add_shard)
|
||||
print(f"merged_package={Path(args.base_model).resolve()}")
|
||||
else:
|
||||
if not args.output:
|
||||
parser.error("--output is required unless --add-shard is used")
|
||||
manifest = merge_package(
|
||||
Path(args.base_model),
|
||||
Path(args.adapter),
|
||||
Path(args.output),
|
||||
Path(args.memory_state) if args.memory_state else None,
|
||||
)
|
||||
output_shard = Path(args.output) / str(manifest["memory_weights"])
|
||||
print(f"merged_package={Path(args.output).resolve()}")
|
||||
print(f"merged_shard={output_shard.resolve()}")
|
||||
print(f"merged_size_gb={output_shard.stat().st_size / (1024**3):.3f}")
|
||||
print(f"contains_user_memory={manifest['checkpoint_contains_user_memory']}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,219 @@
|
||||
"""A small causal LM with an explicit writable memory state.
|
||||
|
||||
This is intentionally independent from the Qwen checkpoint in the parent
|
||||
directory. It is a research reference implementation for architecture work.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
from torch import Tensor, nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
@dataclass
|
||||
class DynamicMemoryConfig:
|
||||
vocab_size: int = 128
|
||||
max_seq_len: int = 64
|
||||
d_model: int = 128
|
||||
n_layers: int = 4
|
||||
n_heads: int = 4
|
||||
mlp_ratio: int = 4
|
||||
memory_slots: int = 8
|
||||
dropout: float = 0.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class DynamicMemoryOutput:
|
||||
logits: Tensor
|
||||
memory: Tensor
|
||||
loss: Optional[Tensor] = None
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, dim: int, eps: float = 1e-6) -> None:
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
variance = x.pow(2).mean(dim=-1, keepdim=True)
|
||||
return x * torch.rsqrt(variance + self.eps) * self.weight
|
||||
|
||||
|
||||
class CausalSelfAttention(nn.Module):
|
||||
def __init__(self, config: DynamicMemoryConfig) -> None:
|
||||
super().__init__()
|
||||
if config.d_model % config.n_heads != 0:
|
||||
raise ValueError("d_model must be divisible by n_heads")
|
||||
self.n_heads = config.n_heads
|
||||
self.head_dim = config.d_model // config.n_heads
|
||||
self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False)
|
||||
self.out = nn.Linear(config.d_model, config.d_model, bias=False)
|
||||
self.dropout = config.dropout
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
batch, seq_len, dim = x.shape
|
||||
qkv = self.qkv(x).view(batch, seq_len, 3, self.n_heads, self.head_dim)
|
||||
q, k, v = qkv.unbind(dim=2)
|
||||
q = q.transpose(1, 2)
|
||||
k = k.transpose(1, 2)
|
||||
v = v.transpose(1, 2)
|
||||
y = F.scaled_dot_product_attention(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
dropout_p=self.dropout if self.training else 0.0,
|
||||
is_causal=True,
|
||||
)
|
||||
y = y.transpose(1, 2).contiguous().view(batch, seq_len, dim)
|
||||
return self.out(y)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, config: DynamicMemoryConfig) -> None:
|
||||
super().__init__()
|
||||
hidden = config.d_model * config.mlp_ratio
|
||||
self.up = nn.Linear(config.d_model, hidden, bias=False)
|
||||
self.down = nn.Linear(hidden, config.d_model, bias=False)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return self.down(F.silu(self.up(x)))
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, config: DynamicMemoryConfig) -> None:
|
||||
super().__init__()
|
||||
self.norm1 = RMSNorm(config.d_model)
|
||||
self.attn = CausalSelfAttention(config)
|
||||
self.norm2 = RMSNorm(config.d_model)
|
||||
self.mlp = MLP(config)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
x = x + self.attn(self.norm1(x))
|
||||
x = x + self.mlp(self.norm2(x))
|
||||
return x
|
||||
|
||||
|
||||
class DynamicMemory(nn.Module):
|
||||
"""A differentiable key-value memory with explicit read/write behavior.
|
||||
|
||||
The memory tensor is returned to the caller and is not a model parameter.
|
||||
It can therefore change during inference without changing the backbone.
|
||||
"""
|
||||
|
||||
def __init__(self, config: DynamicMemoryConfig) -> None:
|
||||
super().__init__()
|
||||
self.slots = config.memory_slots
|
||||
self.dim = config.d_model
|
||||
self.read_q = nn.Linear(self.dim, self.dim, bias=False)
|
||||
self.read_k = nn.Linear(self.dim, self.dim, bias=False)
|
||||
self.read_v = nn.Linear(self.dim, self.dim, bias=False)
|
||||
self.read_out = nn.Linear(self.dim, self.dim, bias=False)
|
||||
self.read_gate = nn.Linear(self.dim, 1)
|
||||
|
||||
self.slot_keys = nn.Parameter(torch.randn(self.slots, self.dim) / self.dim**0.5)
|
||||
self.write_value = nn.Linear(self.dim, self.slots * self.dim, bias=False)
|
||||
self.write_gate = nn.Linear(self.dim, self.slots)
|
||||
|
||||
def initial_state(self, batch_size: int, *, device: torch.device, dtype: torch.dtype) -> Tensor:
|
||||
return torch.zeros(batch_size, self.slots, self.dim, device=device, dtype=dtype)
|
||||
|
||||
def read(self, x: Tensor, memory: Tensor) -> Tensor:
|
||||
q = self.read_q(x)
|
||||
k = self.read_k(memory)
|
||||
v = self.read_v(memory)
|
||||
scores = torch.matmul(q, k.transpose(-1, -2)) / self.dim**0.5
|
||||
retrieved = torch.matmul(scores.softmax(dim=-1), v)
|
||||
retrieved = self.read_out(retrieved)
|
||||
gate = torch.sigmoid(self.read_gate(x))
|
||||
return gate * retrieved
|
||||
|
||||
def update(self, x: Tensor, memory: Tensor) -> Tensor:
|
||||
# The final token is used as a compact summary of the newly observed
|
||||
# chunk. This makes chunk boundaries explicit and keeps the experiment
|
||||
# cheap enough to run repeatedly on a single consumer GPU.
|
||||
summary = x[:, -1]
|
||||
proposal = self.write_value(summary).view(-1, self.slots, self.dim)
|
||||
address = (summary @ self.slot_keys.t()).softmax(dim=-1)
|
||||
strength = torch.sigmoid(self.write_gate(summary)) * address
|
||||
strength = strength.unsqueeze(-1)
|
||||
return memory + strength * (proposal - memory)
|
||||
|
||||
|
||||
class DynamicMemoryLM(nn.Module):
|
||||
"""Decoder-only Transformer with a persistent, caller-owned memory state."""
|
||||
|
||||
def __init__(self, config: DynamicMemoryConfig) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.token_emb = nn.Embedding(config.vocab_size, config.d_model)
|
||||
self.pos_emb = nn.Embedding(config.max_seq_len, config.d_model)
|
||||
self.memory = DynamicMemory(config)
|
||||
self.blocks = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layers))
|
||||
self.norm = RMSNorm(config.d_model)
|
||||
self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
|
||||
self.lm_head.weight = self.token_emb.weight
|
||||
self.apply(self._init_weights)
|
||||
|
||||
@staticmethod
|
||||
def _init_weights(module: nn.Module) -> None:
|
||||
if isinstance(module, nn.Linear):
|
||||
nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
||||
if module.bias is not None:
|
||||
nn.init.zeros_(module.bias)
|
||||
elif isinstance(module, nn.Embedding):
|
||||
nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
||||
|
||||
def _ensure_memory(self, memory: Optional[Tensor], batch_size: int, device: torch.device, dtype: torch.dtype) -> Tensor:
|
||||
if memory is None:
|
||||
return self.memory.initial_state(batch_size, device=device, dtype=dtype)
|
||||
if memory.ndim != 3 or memory.shape[0] != batch_size:
|
||||
raise ValueError("memory must have shape [batch, memory_slots, d_model]")
|
||||
return memory
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Tensor,
|
||||
*,
|
||||
memory: Optional[Tensor] = None,
|
||||
update_memory: bool = True,
|
||||
labels: Optional[Tensor] = None,
|
||||
) -> DynamicMemoryOutput:
|
||||
if input_ids.ndim != 2:
|
||||
raise ValueError("input_ids must have shape [batch, seq]")
|
||||
batch_size, seq_len = input_ids.shape
|
||||
if seq_len > self.config.max_seq_len:
|
||||
raise ValueError(f"sequence length {seq_len} exceeds max_seq_len={self.config.max_seq_len}")
|
||||
|
||||
x = self.token_emb(input_ids)
|
||||
positions = torch.arange(seq_len, device=input_ids.device)
|
||||
x = x + self.pos_emb(positions)[None, :, :]
|
||||
memory = self._ensure_memory(memory, batch_size, input_ids.device, x.dtype)
|
||||
|
||||
# Read uses the state from before this chunk. The write happens only
|
||||
# after logits are computed, preventing target-token leakage.
|
||||
x = x + self.memory.read(x, memory)
|
||||
for block in self.blocks:
|
||||
x = block(x)
|
||||
x = self.norm(x)
|
||||
logits = self.lm_head(x)
|
||||
|
||||
new_memory = self.memory.update(x, memory) if update_memory else memory
|
||||
loss = None
|
||||
if labels is not None:
|
||||
if labels.shape != input_ids.shape:
|
||||
raise ValueError("labels must have the same shape as input_ids")
|
||||
loss = F.cross_entropy(
|
||||
logits[:, :-1].reshape(-1, logits.size(-1)),
|
||||
labels[:, 1:].reshape(-1),
|
||||
ignore_index=-100,
|
||||
)
|
||||
return DynamicMemoryOutput(logits=logits, memory=new_memory, loss=loss)
|
||||
|
||||
|
||||
def count_parameters(model: nn.Module) -> int:
|
||||
return sum(parameter.numel() for parameter in model.parameters())
|
||||
@@ -0,0 +1,54 @@
|
||||
"""Unified entry point for Natural Memory v2 local production workflows.
|
||||
|
||||
Examples:
|
||||
python -m dynamic_memory_lab.natural_memory_app chat --model-path ...
|
||||
python -m dynamic_memory_lab.natural_memory_app serve --port 8765
|
||||
python -m dynamic_memory_lab.natural_memory_app build-dataset
|
||||
python -m dynamic_memory_lab.natural_memory_app train-policy --steps 240
|
||||
python -m dynamic_memory_lab.natural_memory_app stress --rounds 40
|
||||
python -m dynamic_memory_lab.natural_memory_app make-mega-validation
|
||||
python -m dynamic_memory_lab.natural_memory_app benchmark-kv --limit 32
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"command",
|
||||
choices=(
|
||||
"chat",
|
||||
"serve",
|
||||
"build-dataset",
|
||||
"train-policy",
|
||||
"stress",
|
||||
"make-mega-validation",
|
||||
"benchmark-kv",
|
||||
),
|
||||
help="workflow to run; remaining arguments are passed to that workflow",
|
||||
)
|
||||
args, remaining = parser.parse_known_args()
|
||||
sys.argv = [sys.argv[0], *remaining]
|
||||
if args.command == "chat":
|
||||
from .stream_chat_qwen_memory import main as run
|
||||
elif args.command == "serve":
|
||||
from .natural_memory_service import main as run
|
||||
elif args.command == "build-dataset":
|
||||
from .build_production_memory_dataset import main as run
|
||||
elif args.command == "train-policy":
|
||||
from .train_production_memory_policy import main as run
|
||||
elif args.command == "stress":
|
||||
from .stress_test_natural_memory import main as run
|
||||
elif args.command == "make-mega-validation":
|
||||
from .make_mega_memory_validation import main as run
|
||||
else:
|
||||
from .benchmark_memory_vs_full_kv import main as run
|
||||
run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,418 @@
|
||||
"""Local production API for a single-user Natural Memory v2 instance.
|
||||
|
||||
The service binds to localhost by default, keeps one model lock so concurrent
|
||||
requests cannot corrupt the model-owned memory, and persists changed memory
|
||||
back into the embedded third safetensors shard when enabled.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import threading
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
import torch
|
||||
|
||||
from .qwen_integration import load_qwen_dynamic, load_tokenizer
|
||||
from .stream_chat_qwen_memory import _chat_tensor, _persist_memory, _write_turn
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
def _project_path(value: str | Path) -> Path:
|
||||
path = Path(value)
|
||||
if path.is_absolute() or path.exists():
|
||||
return path
|
||||
return PROJECT_ROOT / path
|
||||
|
||||
|
||||
class NaturalMemoryService:
|
||||
def __init__(
|
||||
self,
|
||||
model_path: str | Path,
|
||||
*,
|
||||
no_4bit: bool = False,
|
||||
auto_persist: bool = True,
|
||||
auth_token: str | None = None,
|
||||
) -> None:
|
||||
self.model_path = _project_path(model_path)
|
||||
self.auth_token = auth_token
|
||||
self.auto_persist = bool(auto_persist)
|
||||
self.lock = threading.RLock()
|
||||
self.tokenizer = load_tokenizer(self.model_path)
|
||||
self.model = load_qwen_dynamic(self.model_path, load_in_4bit=not no_4bit)
|
||||
self.model.eval()
|
||||
self.device = self.model._find_layer_device()
|
||||
if self.model.memory_os_v2 is None:
|
||||
raise RuntimeError("the selected package does not contain hierarchical memory")
|
||||
|
||||
def close(self) -> None:
|
||||
with self.lock:
|
||||
self.model.close_memory_storage()
|
||||
del self.model
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def _persist(self) -> None:
|
||||
if not self.auto_persist:
|
||||
return
|
||||
_persist_memory(self.model, embedded_dir=self.model_path, state_path=None)
|
||||
|
||||
def _encode_plain(self, text: str) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
encoded = self.tokenizer(text, add_special_tokens=False, return_tensors="pt")
|
||||
ids = encoded["input_ids"].to(self.device)
|
||||
mask = encoded.get("attention_mask")
|
||||
if mask is None:
|
||||
mask = torch.ones_like(ids)
|
||||
return ids, mask.to(self.device)
|
||||
|
||||
def health(self) -> dict[str, Any]:
|
||||
with self.lock:
|
||||
return {
|
||||
"status": "ok",
|
||||
"model_path": str(self.model_path),
|
||||
"device": str(self.device),
|
||||
"auto_persist": self.auto_persist,
|
||||
"memory": self.model.memory_v2_stats(),
|
||||
"audit": self.model.audit_memory(),
|
||||
"cuda": {
|
||||
"available": torch.cuda.is_available(),
|
||||
"allocated_mb": round(torch.cuda.memory_allocated() / (1024 * 1024), 2)
|
||||
if torch.cuda.is_available()
|
||||
else None,
|
||||
"reserved_mb": round(torch.cuda.memory_reserved() / (1024 * 1024), 2)
|
||||
if torch.cuda.is_available()
|
||||
else None,
|
||||
},
|
||||
}
|
||||
|
||||
def list_memory(self, params: dict[str, list[str]]) -> dict[str, Any]:
|
||||
with self.lock:
|
||||
query = params.get("query", [""])[0]
|
||||
status = params.get("status", ["active"])[0]
|
||||
limit = min(10000, max(1, int(params.get("limit", ["100"])[0])))
|
||||
offset = max(0, int(params.get("offset", ["0"])[0]))
|
||||
records = self.model.list_memory_records(
|
||||
query_text=query, status=status, limit=limit, offset=offset
|
||||
)
|
||||
return {"records": records, "returned": len(records), "offset": offset, "limit": limit}
|
||||
|
||||
def get_memory(self, record_id: str) -> dict[str, Any]:
|
||||
with self.lock:
|
||||
return self.model.get_memory_record(record_id)
|
||||
|
||||
def export_memory(self, params: dict[str, list[str]]) -> dict[str, Any]:
|
||||
with self.lock:
|
||||
limit = min(10000, max(1, int(params.get("limit", ["10000"])[0])))
|
||||
offset = max(0, int(params.get("offset", ["0"])[0]))
|
||||
return self.model.export_memory_records(limit=limit, offset=offset)
|
||||
|
||||
def audit(self) -> dict[str, Any]:
|
||||
with self.lock:
|
||||
return self.model.audit_memory()
|
||||
|
||||
def edit_memory(self, record_id: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
with self.lock, torch.inference_mode():
|
||||
text = payload.get("text")
|
||||
ids = mask = None
|
||||
if isinstance(text, str) and text.strip():
|
||||
ids, mask = self._encode_plain(text)
|
||||
result = self.model.edit_memory_record(
|
||||
record_id,
|
||||
text=text if isinstance(text, str) else None,
|
||||
entity=payload.get("entity"),
|
||||
attribute=payload.get("attribute"),
|
||||
value=payload.get("value"),
|
||||
importance=payload.get("importance"),
|
||||
confidence=payload.get("confidence"),
|
||||
evidence=payload.get("evidence") if isinstance(payload.get("evidence"), list) else None,
|
||||
token_ids=ids,
|
||||
token_mask=mask,
|
||||
)
|
||||
self._persist()
|
||||
return result
|
||||
|
||||
def retract_memory(self, record_id: str) -> dict[str, Any]:
|
||||
with self.lock:
|
||||
result = self.model.retract_memory_record(record_id)
|
||||
self._persist()
|
||||
return result
|
||||
|
||||
def reset_memory(self) -> dict[str, Any]:
|
||||
with self.lock:
|
||||
self.model.reset_memory(batch_size=1, device=self.device)
|
||||
self._persist()
|
||||
return {"reset": True, "memory": self.model.memory_v2_stats()}
|
||||
|
||||
def direct_write(self, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
text = str(payload.get("text", "")).strip()
|
||||
if not text:
|
||||
raise ValueError("text is required")
|
||||
with self.lock, torch.inference_mode():
|
||||
ids, mask = self._encode_plain(text)
|
||||
key = self.model._encode_model_key(ids, mask)[0]
|
||||
record, action = self.model.write_hierarchical_memory(
|
||||
text=text,
|
||||
key=key,
|
||||
summary=key,
|
||||
token_ids=ids[0].detach().cpu(),
|
||||
token_mask=mask[0].detach().cpu().bool(),
|
||||
entity=str(payload.get("entity", "")),
|
||||
attribute=str(payload.get("attribute", "")),
|
||||
value=str(payload.get("value", "")),
|
||||
importance=float(payload.get("importance", 0.9)),
|
||||
confidence=float(payload.get("confidence", 0.99)),
|
||||
source="api",
|
||||
trusted=True,
|
||||
force=bool(payload.get("force", True)),
|
||||
)
|
||||
self._persist()
|
||||
return {"action": action, "record": self.model.get_memory_record(record.record_id)}
|
||||
|
||||
def _prepare_chat(self, message: str) -> tuple[dict[str, torch.Tensor], torch.Tensor, torch.Tensor, bool]:
|
||||
encoded = {key: value.to(self.device) for key, value in _chat_tensor(self.tokenizer, message).items()}
|
||||
query_ids, query_mask = self._encode_plain(message)
|
||||
reset_id = self.model.memory_config.reset_token_id
|
||||
if reset_id is not None and bool((encoded["input_ids"] == reset_id).any()):
|
||||
self.model.reset_memory(batch_size=1, device=self.device)
|
||||
self._persist()
|
||||
return encoded, query_ids, query_mask, True
|
||||
changed = False
|
||||
if self.model.memory_config.native_mode:
|
||||
changed = _write_turn(self.model, self.tokenizer, message, self.device)
|
||||
if changed:
|
||||
self._persist()
|
||||
return encoded, query_ids, query_mask, changed
|
||||
|
||||
def chat(self, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
message = str(payload.get("message", "")).strip()
|
||||
if not message:
|
||||
raise ValueError("message is required")
|
||||
max_new_tokens = min(256, max(1, int(payload.get("max_new_tokens", 128))))
|
||||
with self.lock, torch.inference_mode():
|
||||
encoded, query_ids, query_mask, changed = self._prepare_chat(message)
|
||||
output = self.model.generate(
|
||||
**encoded,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
memory_query_input_ids=query_ids,
|
||||
memory_query_attention_mask=query_mask,
|
||||
memory_query_text=message,
|
||||
use_cache=True,
|
||||
pad_token_id=self.tokenizer.pad_token_id,
|
||||
)
|
||||
answer = self.tokenizer.decode(output[0, encoded["input_ids"].shape[1] :], skip_special_tokens=True)
|
||||
return {
|
||||
"answer": answer,
|
||||
"memory_changed": changed,
|
||||
"memory": self.model.memory_v2_stats(),
|
||||
}
|
||||
|
||||
def stream_chat(self, payload: dict[str, Any]):
|
||||
"""Yield answer fragments while holding the single-model lock."""
|
||||
|
||||
from transformers import TextIteratorStreamer
|
||||
|
||||
message = str(payload.get("message", "")).strip()
|
||||
if not message:
|
||||
raise ValueError("message is required")
|
||||
max_new_tokens = min(256, max(1, int(payload.get("max_new_tokens", 128))))
|
||||
self.lock.acquire()
|
||||
try:
|
||||
encoded, query_ids, query_mask, changed = self._prepare_chat(message)
|
||||
streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)
|
||||
errors: list[BaseException] = []
|
||||
|
||||
def worker() -> None:
|
||||
try:
|
||||
self.model.generate(
|
||||
**encoded,
|
||||
streamer=streamer,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
memory_query_input_ids=query_ids,
|
||||
memory_query_attention_mask=query_mask,
|
||||
memory_query_text=message,
|
||||
use_cache=True,
|
||||
pad_token_id=self.tokenizer.pad_token_id,
|
||||
)
|
||||
except BaseException as error:
|
||||
errors.append(error)
|
||||
streamer.on_finalized_text("", stream_end=True)
|
||||
|
||||
thread = threading.Thread(target=worker, name="natural-memory-api-generation", daemon=True)
|
||||
thread.start()
|
||||
yield {"type": "meta", "memory_changed": changed}
|
||||
for chunk in streamer:
|
||||
yield {"type": "token", "text": chunk}
|
||||
thread.join(timeout=10.0)
|
||||
if thread.is_alive():
|
||||
raise RuntimeError("generation thread did not stop")
|
||||
if errors:
|
||||
raise RuntimeError("stream generation failed") from errors[0]
|
||||
yield {"type": "done", "memory": self.model.memory_v2_stats()}
|
||||
finally:
|
||||
self.lock.release()
|
||||
|
||||
|
||||
class _Handler(BaseHTTPRequestHandler):
|
||||
server: "_Server"
|
||||
protocol_version = "HTTP/1.1"
|
||||
|
||||
def log_message(self, format: str, *args: Any) -> None:
|
||||
print(f"[natural-memory] {self.address_string()} {format % args}")
|
||||
|
||||
@property
|
||||
def service(self) -> NaturalMemoryService:
|
||||
return self.server.service
|
||||
|
||||
def _authorized(self) -> bool:
|
||||
expected = self.service.auth_token
|
||||
if not expected:
|
||||
return True
|
||||
return self.headers.get("Authorization", "") == f"Bearer {expected}"
|
||||
|
||||
def _send_json(self, payload: Any, status: int = 200) -> None:
|
||||
body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Type", "application/json; charset=utf-8")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.send_header("Cache-Control", "no-store")
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
def _send_error_json(self, error: BaseException, status: int = 400) -> None:
|
||||
self._send_json({"error": type(error).__name__, "message": str(error)}, status)
|
||||
|
||||
def _body(self) -> dict[str, Any]:
|
||||
length = int(self.headers.get("Content-Length", "0"))
|
||||
if length > 1024 * 1024:
|
||||
raise ValueError("request body exceeds 1 MiB")
|
||||
raw = self.rfile.read(length) if length else b"{}"
|
||||
value = json.loads(raw.decode("utf-8"))
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError("request body must be a JSON object")
|
||||
return value
|
||||
|
||||
def do_GET(self) -> None:
|
||||
if not self._authorized():
|
||||
self._send_json({"error": "Unauthorized"}, 401)
|
||||
return
|
||||
parsed = urlparse(self.path)
|
||||
try:
|
||||
if parsed.path == "/health":
|
||||
self._send_json(self.service.health())
|
||||
elif parsed.path == "/v1/memory":
|
||||
self._send_json(self.service.list_memory(parse_qs(parsed.query)))
|
||||
elif parsed.path.startswith("/v1/memory/export"):
|
||||
self._send_json(self.service.export_memory(parse_qs(parsed.query)))
|
||||
elif parsed.path == "/v1/memory/audit":
|
||||
self._send_json(self.service.audit())
|
||||
elif parsed.path.startswith("/v1/memory/"):
|
||||
self._send_json(self.service.get_memory(parsed.path.rsplit("/", 1)[1]))
|
||||
else:
|
||||
self._send_json({"error": "not_found"}, 404)
|
||||
except KeyError as error:
|
||||
self._send_error_json(error, 404)
|
||||
except Exception as error:
|
||||
self._send_error_json(error, 400)
|
||||
|
||||
def do_POST(self) -> None:
|
||||
if not self._authorized():
|
||||
self._send_json({"error": "Unauthorized"}, 401)
|
||||
return
|
||||
try:
|
||||
payload = self._body()
|
||||
parsed = urlparse(self.path)
|
||||
if parsed.path == "/v1/chat":
|
||||
if bool(payload.get("stream", False)):
|
||||
self._send_stream(self.service.stream_chat(payload))
|
||||
else:
|
||||
self._send_json(self.service.chat(payload))
|
||||
elif parsed.path == "/v1/memory":
|
||||
self._send_json(self.service.direct_write(payload), 201)
|
||||
elif parsed.path == "/v1/memory/reset":
|
||||
self._send_json(self.service.reset_memory())
|
||||
elif parsed.path.startswith("/v1/memory/"):
|
||||
self._send_json(self.service.edit_memory(parsed.path.rsplit("/", 1)[1], payload))
|
||||
else:
|
||||
self._send_json({"error": "not_found"}, 404)
|
||||
except KeyError as error:
|
||||
self._send_error_json(error, 404)
|
||||
except Exception as error:
|
||||
self._send_error_json(error, 400)
|
||||
|
||||
def do_DELETE(self) -> None:
|
||||
if not self._authorized():
|
||||
self._send_json({"error": "Unauthorized"}, 401)
|
||||
return
|
||||
try:
|
||||
path = urlparse(self.path).path
|
||||
if not path.startswith("/v1/memory/"):
|
||||
self._send_json({"error": "not_found"}, 404)
|
||||
return
|
||||
self._send_json(self.service.retract_memory(path.rsplit("/", 1)[1]))
|
||||
except KeyError as error:
|
||||
self._send_error_json(error, 404)
|
||||
except Exception as error:
|
||||
self._send_error_json(error, 400)
|
||||
|
||||
def _send_stream(self, events) -> None:
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
|
||||
self.send_header("Cache-Control", "no-cache")
|
||||
self.send_header("Connection", "close")
|
||||
self.end_headers()
|
||||
try:
|
||||
for event in events:
|
||||
body = json.dumps(event, ensure_ascii=False)
|
||||
self.wfile.write(f"data: {body}\n\n".encode("utf-8"))
|
||||
self.wfile.flush()
|
||||
except Exception as error:
|
||||
body = json.dumps({"type": "error", "message": str(error)}, ensure_ascii=False)
|
||||
self.wfile.write(f"data: {body}\n\n".encode("utf-8"))
|
||||
|
||||
|
||||
class _Server(ThreadingHTTPServer):
|
||||
def __init__(self, address, service: NaturalMemoryService):
|
||||
super().__init__(address, _Handler)
|
||||
self.service = service
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-path", default="qwen3_5_4b_natural_memory_v2")
|
||||
parser.add_argument("--host", default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, default=8765)
|
||||
parser.add_argument("--auth-token", default=None)
|
||||
parser.add_argument("--no-auto-persist", action="store_true")
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
args = parser.parse_args()
|
||||
service = NaturalMemoryService(
|
||||
args.model_path,
|
||||
no_4bit=args.no_4bit,
|
||||
auto_persist=not args.no_auto_persist,
|
||||
auth_token=args.auth_token,
|
||||
)
|
||||
server = _Server((args.host, args.port), service)
|
||||
print(f"Natural Memory API listening on http://{args.host}:{args.port}")
|
||||
try:
|
||||
server.serve_forever()
|
||||
except KeyboardInterrupt:
|
||||
print("\nStopping Natural Memory API")
|
||||
finally:
|
||||
server.server_close()
|
||||
service.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+3567
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,30 @@
|
||||
"""Fast environment and gradient smoke test."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
from .model import DynamicMemoryConfig, DynamicMemoryLM, count_parameters
|
||||
from .tasks import sample_associative_batch
|
||||
|
||||
|
||||
def main() -> None:
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
config = DynamicMemoryConfig(vocab_size=64, max_seq_len=16, d_model=64, n_layers=2, n_heads=4, memory_slots=4)
|
||||
model = DynamicMemoryLM(config).to(device)
|
||||
batch = sample_associative_batch(batch_size=8, vocab_size=config.vocab_size, device=device)
|
||||
memory = model(batch.learn_chunks[0]).memory
|
||||
output = model(batch.query_input, memory=memory, update_memory=False, labels=batch.query_labels)
|
||||
if output.loss is None or not torch.isfinite(output.loss):
|
||||
raise RuntimeError("non-finite loss")
|
||||
output.loss.backward()
|
||||
gradients = [p.grad for p in model.parameters() if p.grad is not None]
|
||||
if not gradients:
|
||||
raise RuntimeError("no gradients produced")
|
||||
print(f"smoke_ok device={device} parameters={count_parameters(model):,} loss={output.loss.detach().item():.4f}")
|
||||
if device.type == "cuda":
|
||||
print(f"gpu={torch.cuda.get_device_name(0)} memory_allocated_mb={torch.cuda.memory_allocated() / 1024**2:.1f}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,448 @@
|
||||
"""Streaming chat for natural-language memory with restart-safe autosave.
|
||||
|
||||
Every turn is encoded independently. The model's internal reader decides
|
||||
whether a saved memory prefix is relevant; this script never reconstructs
|
||||
conversation history. A user memory state is atomically saved before the
|
||||
streaming generation starts, so restarting the process is safe at any time.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from .qwen_integration import (
|
||||
DEFAULT_MEMORY_RESET_TOKEN,
|
||||
QwenMemoryConfig,
|
||||
load_memory_config,
|
||||
load_qwen_dynamic,
|
||||
load_tokenizer,
|
||||
resolve_memory_reset_token,
|
||||
split_memory_candidates,
|
||||
)
|
||||
|
||||
|
||||
def _memory_system_prefix(tokenizer, content: str) -> dict[str, torch.Tensor]:
|
||||
"""Encode a valid system prefix without adding a fake user question."""
|
||||
|
||||
full = tokenizer.apply_chat_template(
|
||||
[
|
||||
{"role": "system", "content": content},
|
||||
{"role": "user", "content": "__memory_query_boundary__"},
|
||||
],
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
return_tensors="pt",
|
||||
return_dict=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
input_ids = full["input_ids"]
|
||||
im_start = tokenizer.convert_tokens_to_ids("<|im_start|>")
|
||||
positions = (input_ids[0] == int(im_start)).nonzero(as_tuple=False).flatten()
|
||||
if positions.numel() < 2:
|
||||
raise RuntimeError("could not locate the system/user memory boundary")
|
||||
end = int(positions[1].item())
|
||||
return {
|
||||
"input_ids": input_ids[:, :end],
|
||||
"attention_mask": torch.ones((1, end), dtype=torch.long),
|
||||
}
|
||||
|
||||
|
||||
def _chat_tensor(tokenizer, user_text: str) -> dict[str, torch.Tensor]:
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
[{"role": "user", "content": user_text}],
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
return_tensors="pt",
|
||||
return_dict=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
return {
|
||||
key: value
|
||||
for key, value in encoded.items()
|
||||
if isinstance(value, torch.Tensor)
|
||||
}
|
||||
|
||||
|
||||
def _atomic_save(model, path: Path) -> None:
|
||||
"""Save one user's state without exposing a partially written file."""
|
||||
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary = path.with_name(path.name + ".tmp")
|
||||
try:
|
||||
model.save_runtime_memory(temporary)
|
||||
os.replace(temporary, path)
|
||||
finally:
|
||||
if temporary.exists():
|
||||
temporary.unlink()
|
||||
|
||||
|
||||
def _persist_memory(
|
||||
model,
|
||||
*,
|
||||
embedded_dir: Path | None,
|
||||
state_path: Path | None,
|
||||
) -> None:
|
||||
"""Persist either into the merged shard or into a normal runtime file."""
|
||||
|
||||
if embedded_dir is not None:
|
||||
if getattr(model.memory_config, "memory_storage_mode", "embedded") == "tiered":
|
||||
model.flush_memory_storage()
|
||||
return
|
||||
model.save_embedded_memory_weights(embedded_dir)
|
||||
return
|
||||
if state_path is None:
|
||||
raise ValueError("no persistence target is configured")
|
||||
_atomic_save(model, state_path)
|
||||
|
||||
|
||||
def _slot_count(model) -> int:
|
||||
valid = model.runtime.text_slot_valid
|
||||
return int(valid.sum().item()) if isinstance(valid, torch.Tensor) else 0
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _write_turn(
|
||||
model,
|
||||
tokenizer,
|
||||
text: str,
|
||||
device: torch.device,
|
||||
*,
|
||||
force_write: bool = False,
|
||||
) -> bool:
|
||||
"""Run the learned write controller for one user turn."""
|
||||
changed = False
|
||||
for candidate in split_memory_candidates(text):
|
||||
encoded = _chat_tensor(tokenizer, candidate)
|
||||
encoded = {key: value.to(device) for key, value in encoded.items()}
|
||||
memory_prefix = _memory_system_prefix(
|
||||
tokenizer,
|
||||
"以下是与当前用户相关的已保存长期记忆。仅在问题相关时使用，不要编造：\n" + candidate,
|
||||
)
|
||||
memory_text_ids = memory_prefix["input_ids"].to(device)
|
||||
memory_text_mask = memory_prefix["attention_mask"].to(device)
|
||||
memory_key = tokenizer(candidate, add_special_tokens=False, return_tensors="pt")
|
||||
memory_key_ids = memory_key["input_ids"].to(device)
|
||||
memory_key_mask = memory_key.get("attention_mask")
|
||||
if memory_key_mask is None:
|
||||
memory_key_mask = torch.ones_like(memory_key_ids)
|
||||
memory_key_mask = memory_key_mask.to(device)
|
||||
memory_storage = tokenizer(
|
||||
candidate,
|
||||
add_special_tokens=False,
|
||||
return_tensors="pt",
|
||||
)
|
||||
memory_storage_ids = memory_storage["input_ids"].to(device)
|
||||
memory_storage_mask = memory_storage.get("attention_mask")
|
||||
if memory_storage_mask is None:
|
||||
memory_storage_mask = torch.ones_like(memory_storage_ids)
|
||||
memory_storage_mask = memory_storage_mask.to(device)
|
||||
model(
|
||||
**encoded,
|
||||
# The write controller must see only the current user turn. If
|
||||
# it reads the already-retrieved prefix first, the policy can
|
||||
# mistake recalled facts for a new fact and write questions back.
|
||||
read_memory=False,
|
||||
update_memory=True,
|
||||
return_memory=True,
|
||||
use_cache=False,
|
||||
memory_text_input_ids=memory_text_ids,
|
||||
memory_text_attention_mask=memory_text_mask,
|
||||
memory_key_input_ids=memory_key_ids,
|
||||
memory_key_attention_mask=memory_key_mask,
|
||||
memory_storage_input_ids=memory_storage_ids,
|
||||
memory_storage_attention_mask=memory_storage_mask,
|
||||
force_memory_write=force_write,
|
||||
memory_text=candidate,
|
||||
)
|
||||
last_written = model.runtime.text_last_written_slot
|
||||
if isinstance(last_written, torch.Tensor):
|
||||
changed = changed or bool((last_written >= 0).any())
|
||||
return changed
|
||||
|
||||
|
||||
def _stream_answer(
|
||||
model,
|
||||
tokenizer,
|
||||
encoded: dict[str, torch.Tensor],
|
||||
max_new_tokens: int,
|
||||
memory_query_input_ids: torch.Tensor,
|
||||
memory_query_attention_mask: torch.Tensor,
|
||||
memory_query_text: str,
|
||||
) -> None:
|
||||
"""Stream one answer through TextIteratorStreamer in a worker thread."""
|
||||
|
||||
from transformers import StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
|
||||
|
||||
streamer = TextIteratorStreamer(
|
||||
tokenizer,
|
||||
skip_prompt=True,
|
||||
skip_special_tokens=True,
|
||||
)
|
||||
errors: list[BaseException] = []
|
||||
stop_event = threading.Event()
|
||||
|
||||
class StopOnEvent(StoppingCriteria):
|
||||
def __call__(self, input_ids, scores, **kwargs):
|
||||
return torch.full(
|
||||
(input_ids.shape[0],),
|
||||
stop_event.is_set(),
|
||||
dtype=torch.bool,
|
||||
device=input_ids.device,
|
||||
)
|
||||
|
||||
def worker() -> None:
|
||||
try:
|
||||
model.generate(
|
||||
**encoded,
|
||||
streamer=streamer,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
memory_query_input_ids=memory_query_input_ids,
|
||||
memory_query_attention_mask=memory_query_attention_mask,
|
||||
memory_query_text=memory_query_text,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
stopping_criteria=StoppingCriteriaList([StopOnEvent()]),
|
||||
)
|
||||
except BaseException as error: # propagate through the main thread
|
||||
errors.append(error)
|
||||
streamer.on_finalized_text("", stream_end=True)
|
||||
|
||||
thread = threading.Thread(target=worker, name="qwen-stream-generation", daemon=True)
|
||||
thread.start()
|
||||
interrupted = False
|
||||
try:
|
||||
for chunk in streamer:
|
||||
print(chunk, end="", flush=True)
|
||||
except KeyboardInterrupt:
|
||||
interrupted = True
|
||||
stop_event.set()
|
||||
print("\n[已中断生成；最近一次 memory state 已保存，可直接重启]", flush=True)
|
||||
finally:
|
||||
thread.join(timeout=10.0)
|
||||
if thread.is_alive():
|
||||
raise RuntimeError("generation did not stop after interrupt; refusing concurrent state save")
|
||||
if interrupted:
|
||||
raise KeyboardInterrupt
|
||||
if errors:
|
||||
raise RuntimeError("streaming generation failed") from errors[0]
|
||||
print()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
if hasattr(sys.stdout, "reconfigure"):
|
||||
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
|
||||
if hasattr(sys.stdin, "reconfigure"):
|
||||
sys.stdin.reconfigure(encoding="utf-8", errors="replace")
|
||||
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-path", default=".")
|
||||
parser.add_argument(
|
||||
"--adapter",
|
||||
default=None,
|
||||
help="external adapter; omitted means use an embedded merge package or v13",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--memory-state",
|
||||
default=None,
|
||||
help="optional external state file; omitted for merged models means write back to the memory shard",
|
||||
)
|
||||
parser.add_argument("--max-new-tokens", type=int, default=128)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
parser.add_argument("--natural-language-memory", action="store_true")
|
||||
parser.add_argument("--reset-token", default=None)
|
||||
parser.add_argument("--reset-token-id", type=int, default=None)
|
||||
parser.add_argument(
|
||||
"--kv-offload",
|
||||
action="store_true",
|
||||
help="keep generation KV on CPU when supported by the installed Transformers",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--kv-cache-implementation",
|
||||
default=None,
|
||||
help="optional Transformers cache implementation, for example offloaded",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-auto-compact",
|
||||
action="store_true",
|
||||
help="disable model-owned old-context archiving when the hot KV budget is exceeded",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tiered-memory-path",
|
||||
default=None,
|
||||
help="optional SQLite page store for warm/cold memory; relative paths use the model package",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--memory-resident-pages",
|
||||
type=int,
|
||||
default=None,
|
||||
help="maximum number of tiered pages kept resident",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.adapter is None and not (Path(args.model_path) / "memory_merge.json").exists():
|
||||
args.adapter = "dynamic_memory_lab/qwen_memory_adapter_natural_auto_v13"
|
||||
|
||||
embedded_dir = (
|
||||
Path(args.model_path)
|
||||
if (
|
||||
(Path(args.model_path) / "memory_merge.json").exists()
|
||||
and args.memory_state is None
|
||||
and args.adapter is None
|
||||
)
|
||||
else None
|
||||
)
|
||||
|
||||
tokenizer = load_tokenizer(args.model_path)
|
||||
memory_config = load_memory_config(args.adapter) if args.adapter else None
|
||||
if args.tiered_memory_path is not None and memory_config is None:
|
||||
memory_config = load_memory_config(args.model_path)
|
||||
if args.tiered_memory_path is not None:
|
||||
memory_config.memory_storage_mode = "tiered"
|
||||
memory_config.memory_storage_path = args.tiered_memory_path
|
||||
if args.memory_resident_pages is not None:
|
||||
if memory_config is None:
|
||||
memory_config = QwenMemoryConfig()
|
||||
memory_config.memory_resident_pages = args.memory_resident_pages
|
||||
if args.natural_language_memory and memory_config is None:
|
||||
memory_config = QwenMemoryConfig(natural_language_memory=True)
|
||||
if memory_config is not None and args.natural_language_memory:
|
||||
memory_config.natural_language_memory = True
|
||||
if memory_config is not None:
|
||||
if args.reset_token_id is not None:
|
||||
memory_config.reset_token_id = args.reset_token_id
|
||||
elif args.reset_token is not None:
|
||||
memory_config.reset_token_id = resolve_memory_reset_token(tokenizer, args.reset_token)
|
||||
elif memory_config.native_mode and memory_config.reset_token_id is None:
|
||||
memory_config.reset_token_id = resolve_memory_reset_token(tokenizer)
|
||||
|
||||
model = load_qwen_dynamic(
|
||||
args.model_path,
|
||||
memory_config=memory_config,
|
||||
load_in_4bit=not args.no_4bit,
|
||||
)
|
||||
if memory_config is None:
|
||||
# An embedded merge package supplies its architecture metadata during
|
||||
# load_qwen_dynamic(); keep the CLI state machine in sync with it.
|
||||
memory_config = model.memory_config
|
||||
if memory_config.native_mode and memory_config.reset_token_id is None:
|
||||
memory_config.reset_token_id = resolve_memory_reset_token(tokenizer)
|
||||
if args.kv_offload:
|
||||
model.memory_config.kv_offload = True
|
||||
if args.kv_cache_implementation is not None:
|
||||
model.memory_config.kv_cache_implementation = args.kv_cache_implementation
|
||||
if args.no_auto_compact:
|
||||
model.memory_config.auto_compact_context = False
|
||||
if args.adapter:
|
||||
model.load_memory_adapter(args.adapter)
|
||||
model.eval()
|
||||
device = model._find_layer_device()
|
||||
state_path = (
|
||||
Path(args.memory_state)
|
||||
if args.memory_state
|
||||
else (None if embedded_dir is not None else Path("dynamic_memory_lab/data/stream_user_memory.pt"))
|
||||
)
|
||||
if state_path is not None and state_path.exists():
|
||||
model.load_runtime_memory(state_path, device=device)
|
||||
print(f"已恢复 memory state：{state_path}，有效记忆槽 {_slot_count(model)} 个")
|
||||
elif model.memory_config.persistent_memory and model.runtime.state is not None:
|
||||
print(f"已使用合并权重内置的 memory state，有效记忆槽 {_slot_count(model)} 个")
|
||||
else:
|
||||
model.reset_memory(batch_size=1, device=device)
|
||||
if state_path is None:
|
||||
print("新建 memory state：合并权重回写模式")
|
||||
else:
|
||||
print(f"新建 memory state：{state_path}")
|
||||
|
||||
reset_token = args.reset_token or DEFAULT_MEMORY_RESET_TOKEN
|
||||
print("流式聊天已启动。不会发送历史聊天记录。")
|
||||
print("命令：/remember <事实>、/reset、/save、/quit；也可发送 reset token。")
|
||||
print(f"reset token：{reset_token}")
|
||||
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
user_text = input("你> ").strip()
|
||||
except EOFError:
|
||||
break
|
||||
if user_text == "/quit":
|
||||
break
|
||||
if user_text == "/save":
|
||||
_persist_memory(model, embedded_dir=embedded_dir, state_path=state_path)
|
||||
target = "主权重切片" if embedded_dir is not None else str(state_path)
|
||||
print(f"已保存到 {target}，当前有效记忆槽 {_slot_count(model)} 个")
|
||||
continue
|
||||
if user_text == "/reset":
|
||||
model.reset_memory(batch_size=1, device=device)
|
||||
_persist_memory(model, embedded_dir=embedded_dir, state_path=state_path)
|
||||
print("已清空并保存。")
|
||||
continue
|
||||
if user_text.startswith("/remember "):
|
||||
fact = user_text[len("/remember ") :].strip()
|
||||
if fact:
|
||||
_write_turn(model, tokenizer, fact, device, force_write=True)
|
||||
_persist_memory(model, embedded_dir=embedded_dir, state_path=state_path)
|
||||
target = "主权重切片" if embedded_dir is not None else str(state_path)
|
||||
print(f"已写入并保存到 {target}，当前有效记忆槽 {_slot_count(model)} 个")
|
||||
continue
|
||||
if not user_text:
|
||||
continue
|
||||
|
||||
encoded = _chat_tensor(tokenizer, user_text)
|
||||
encoded = {key: value.to(device) for key, value in encoded.items()}
|
||||
memory_query = tokenizer(
|
||||
user_text,
|
||||
add_special_tokens=False,
|
||||
return_tensors="pt",
|
||||
)
|
||||
memory_query_input_ids = memory_query["input_ids"].to(device)
|
||||
memory_query_attention_mask = memory_query.get("attention_mask")
|
||||
if memory_query_attention_mask is None:
|
||||
memory_query_attention_mask = torch.ones_like(memory_query_input_ids)
|
||||
memory_query_attention_mask = memory_query_attention_mask.to(device)
|
||||
if (
|
||||
memory_config is not None
|
||||
and memory_config.reset_token_id is not None
|
||||
and bool((encoded["input_ids"] == memory_config.reset_token_id).any())
|
||||
):
|
||||
# Persist the clear operation before generation, so an
|
||||
# interrupted response cannot resurrect the old memory.
|
||||
model.reset_memory(batch_size=1, device=device)
|
||||
_persist_memory(model, embedded_dir=embedded_dir, state_path=state_path)
|
||||
elif memory_config is not None and memory_config.native_mode:
|
||||
# The controller sees the current turn only. Save before
|
||||
# generation; generation itself is read-only.
|
||||
changed = _write_turn(model, tokenizer, user_text, device)
|
||||
if changed:
|
||||
_persist_memory(model, embedded_dir=embedded_dir, state_path=state_path)
|
||||
|
||||
print("AI> ", end="", flush=True)
|
||||
_stream_answer(
|
||||
model,
|
||||
tokenizer,
|
||||
encoded,
|
||||
args.max_new_tokens,
|
||||
memory_query_input_ids,
|
||||
memory_query_attention_mask,
|
||||
user_text,
|
||||
)
|
||||
except KeyboardInterrupt:
|
||||
print("\n[已退出；最近一次记忆已保存，可直接重启]", flush=True)
|
||||
finally:
|
||||
model.close_memory_storage()
|
||||
del model
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,198 @@
|
||||
"""Run a bounded, restart-like Natural Memory v2 endurance test.
|
||||
|
||||
The test deliberately clears only the in-process working copy. It never
|
||||
calls ``save_embedded_memory_weights`` and therefore does not modify the
|
||||
shipped model package. Defaults are conservative for a 12 GB GPU.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import statistics
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from .qwen_integration import load_qwen_dynamic, load_tokenizer
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
def _project_path(value: str | Path) -> Path:
|
||||
path = Path(value)
|
||||
if path.is_absolute() or path.exists():
|
||||
return path
|
||||
return PROJECT_ROOT / path
|
||||
|
||||
|
||||
def _gpu_stats() -> dict[str, float | int | None]:
|
||||
if not torch.cuda.is_available():
|
||||
return {"allocated_mb": None, "reserved_mb": None, "free_mb": None, "total_mb": None}
|
||||
allocated = torch.cuda.memory_allocated() / (1024 * 1024)
|
||||
reserved = torch.cuda.memory_reserved() / (1024 * 1024)
|
||||
free, total = torch.cuda.mem_get_info()
|
||||
return {
|
||||
"allocated_mb": round(allocated, 2),
|
||||
"reserved_mb": round(reserved, 2),
|
||||
"free_mb": round(free / (1024 * 1024), 2),
|
||||
"total_mb": round(total / (1024 * 1024), 2),
|
||||
}
|
||||
|
||||
|
||||
def _long_input(tokenizer, target_tokens: int) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
seed_text = (
|
||||
"这是 Natural Memory v2 的长上下文压力测试。它应当把旧上下文归档到有界的记忆记录，"
|
||||
"保留最近工作窗口，并且不让当前 token 对全部页面做注意力。"
|
||||
)
|
||||
seed = tokenizer(seed_text, add_special_tokens=False, return_tensors="pt")["input_ids"][0]
|
||||
repeats = max(1, (target_tokens + seed.numel() - 1) // seed.numel())
|
||||
ids = seed.repeat(repeats)[:target_tokens].unsqueeze(0)
|
||||
return ids, torch.ones_like(ids)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-path", default="qwen3_5_4b_natural_memory_v2")
|
||||
parser.add_argument("--rounds", type=int, default=40)
|
||||
parser.add_argument("--records-per-round", type=int, default=8)
|
||||
parser.add_argument("--long-context-every", type=int, default=10)
|
||||
parser.add_argument("--long-context-tokens", type=int, default=2048)
|
||||
parser.add_argument("--kv-budget", type=int, default=1024)
|
||||
parser.add_argument("--chunk-tokens", type=int, default=256)
|
||||
parser.add_argument("--duration-minutes", type=float, default=0.0)
|
||||
parser.add_argument("--output", default="natural_memory_v2_stress_report.json")
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
args = parser.parse_args()
|
||||
if args.rounds < 1 or args.records_per_round < 1:
|
||||
raise SystemExit("--rounds and --records-per-round must be positive")
|
||||
if args.kv_budget < 1 or args.long_context_tokens < args.kv_budget:
|
||||
raise SystemExit("--long-context-tokens must be >= --kv-budget >= 1")
|
||||
|
||||
model_path = _project_path(args.model_path)
|
||||
tokenizer = load_tokenizer(model_path)
|
||||
model = load_qwen_dynamic(model_path, load_in_4bit=not args.no_4bit)
|
||||
model.eval()
|
||||
device = model._find_layer_device()
|
||||
if model.memory_os_v2 is None:
|
||||
raise RuntimeError("the selected model does not contain hierarchical memory")
|
||||
|
||||
# Work on an empty ephemeral memory instance. The embedded package on
|
||||
# disk is never changed by this script.
|
||||
model.clear_hierarchical_memory()
|
||||
model.memory_os_v2.read_threshold = 0.0
|
||||
model.memory_os_v2.kv_budget.max_tokens = min(
|
||||
int(args.kv_budget), int(model.memory_os_v2.kv_budget.hard_max_tokens)
|
||||
)
|
||||
model.reset_runtime_memory(batch_size=1, device=device)
|
||||
|
||||
latencies: list[float] = []
|
||||
writes = 0
|
||||
read_hits = 0
|
||||
read_attempts = 0
|
||||
compactions = 0
|
||||
archived_records = 0
|
||||
errors: list[str] = []
|
||||
start = time.perf_counter()
|
||||
next_record = 0
|
||||
try:
|
||||
for round_index in range(args.rounds):
|
||||
if args.duration_minutes > 0 and (time.perf_counter() - start) >= args.duration_minutes * 60:
|
||||
break
|
||||
round_start = time.perf_counter()
|
||||
try:
|
||||
for _ in range(args.records_per_round):
|
||||
record_index = next_record
|
||||
next_record += 1
|
||||
key = torch.randn(model.memory.hidden_size)
|
||||
token_ids = torch.tensor([1000 + (record_index % 10000), 2000 + (record_index % 10000)])
|
||||
text = f"stress record {record_index} belongs to Natural Memory endurance test"
|
||||
record, _ = model.memory_os_v2.write(
|
||||
text=text,
|
||||
key=key,
|
||||
summary=key,
|
||||
token_ids=token_ids,
|
||||
token_mask=torch.ones_like(token_ids, dtype=torch.bool),
|
||||
memory_type="stress_test",
|
||||
entity=f"stress-user-{record_index % 7}",
|
||||
attribute=f"attribute-{record_index}",
|
||||
value=str(record_index),
|
||||
importance=0.9,
|
||||
confidence=0.99,
|
||||
source="stress_test",
|
||||
trusted=True,
|
||||
force=True,
|
||||
)
|
||||
writes += 1
|
||||
read_attempts += 1
|
||||
found, _ = model.memory_os_v2.read(
|
||||
query_key=key,
|
||||
query_text=text,
|
||||
query_token_ids=token_ids,
|
||||
top_k_pages=4,
|
||||
top_k_records=8,
|
||||
max_hops=3,
|
||||
)
|
||||
if any(item.record_id == record.record_id for item in found):
|
||||
read_hits += 1
|
||||
|
||||
if args.long_context_every > 0 and (round_index + 1) % args.long_context_every == 0:
|
||||
ids, mask = _long_input(tokenizer, args.long_context_tokens)
|
||||
_, _, plan = model.compact_context_for_kv(
|
||||
ids.to(device),
|
||||
mask.to(device),
|
||||
archive=True,
|
||||
chunk_tokens=args.chunk_tokens,
|
||||
)
|
||||
if plan.get("compacted"):
|
||||
compactions += 1
|
||||
archived_records += int(plan.get("archived_records", 0))
|
||||
except Exception as error: # keep the report useful after one bad round
|
||||
errors.append(f"round {round_index}: {type(error).__name__}: {error}")
|
||||
latencies.append(time.perf_counter() - round_start)
|
||||
finally:
|
||||
final_stats = model.memory_v2_stats()
|
||||
audit = model.audit_memory()
|
||||
gpu_peak = _gpu_stats()
|
||||
model.close_memory_storage()
|
||||
del model
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
latencies_sorted = sorted(latencies)
|
||||
report = {
|
||||
"format_version": 1,
|
||||
"model_path": str(model_path),
|
||||
"rounds_completed": len(latencies),
|
||||
"records_per_round": args.records_per_round,
|
||||
"writes": writes,
|
||||
"read_attempts": read_attempts,
|
||||
"read_hits": read_hits,
|
||||
"read_hit_rate": read_hits / max(1, read_attempts),
|
||||
"compactions": compactions,
|
||||
"archived_records": archived_records,
|
||||
"round_latency_seconds": {
|
||||
"mean": statistics.fmean(latencies) if latencies else 0.0,
|
||||
"p50": latencies_sorted[len(latencies_sorted) // 2] if latencies_sorted else 0.0,
|
||||
"p95": latencies_sorted[min(len(latencies_sorted) - 1, int(len(latencies_sorted) * 0.95))] if latencies_sorted else 0.0,
|
||||
"max": max(latencies, default=0.0),
|
||||
},
|
||||
"max_gpu": gpu_peak,
|
||||
"final_memory_stats": final_stats,
|
||||
"audit": audit,
|
||||
"errors": errors,
|
||||
"package_mutated": False,
|
||||
"warning": "This is an endurance/safety smoke test, not a quality benchmark.",
|
||||
}
|
||||
output = _project_path(args.output)
|
||||
output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(json.dumps(report, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,56 @@
|
||||
"""Synthetic streaming tasks for testing dynamic learning behavior."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
@dataclass
|
||||
class AssociativeBatch:
|
||||
learn_chunks: list[Tensor]
|
||||
query_input: Tensor
|
||||
query_labels: Tensor
|
||||
expected: Tensor
|
||||
|
||||
|
||||
def _rand_tokens(batch_size: int, low: int, high: int, device: torch.device) -> Tensor:
|
||||
return torch.randint(low, high, (batch_size,), device=device)
|
||||
|
||||
|
||||
def sample_associative_batch(
|
||||
*,
|
||||
batch_size: int,
|
||||
vocab_size: int,
|
||||
device: torch.device,
|
||||
overwrite: bool = False,
|
||||
) -> AssociativeBatch:
|
||||
"""Generate key-value observations followed by a separated query.
|
||||
|
||||
The query cannot see the learning chunks through attention. It can only
|
||||
answer by using the returned dynamic memory state.
|
||||
"""
|
||||
|
||||
key_low, key_high = 4, vocab_size // 2
|
||||
value_low, value_high = vocab_size // 2, vocab_size
|
||||
keys = _rand_tokens(batch_size, key_low, key_high, device)
|
||||
value = _rand_tokens(batch_size, value_low, value_high, device)
|
||||
learn_chunks = [torch.stack((keys, value), dim=1)]
|
||||
|
||||
expected = value
|
||||
if overwrite:
|
||||
replacement = _rand_tokens(batch_size, value_low, value_high, device)
|
||||
learn_chunks.append(torch.stack((keys, replacement), dim=1))
|
||||
expected = replacement
|
||||
|
||||
query_input = torch.stack((keys, expected), dim=1)
|
||||
query_labels = torch.full_like(query_input, -100)
|
||||
query_labels[:, 1] = expected
|
||||
return AssociativeBatch(
|
||||
learn_chunks=learn_chunks,
|
||||
query_input=query_input,
|
||||
query_labels=query_labels,
|
||||
expected=expected,
|
||||
)
|
||||
@@ -0,0 +1,158 @@
|
||||
"""Verify Natural Memory v2 across a real model restart.
|
||||
|
||||
The test deliberately uses a normal user turn, not ``/remember`` and not a
|
||||
replayed chat history. It writes the model-owned V2 snapshot into the
|
||||
embedded memory shard, destroys the first model, reloads the package, and
|
||||
checks both the bounded router decision and the generated answer.
|
||||
|
||||
By default the test restores an empty memory snapshot at the end. Use
|
||||
``--keep-memory`` only when the test fact should remain in the package.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
if __package__ in {None, ""}:
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
from dynamic_memory_lab.qwen_integration import load_qwen_dynamic, load_tokenizer
|
||||
from dynamic_memory_lab.stream_chat_qwen_memory import _chat_tensor, _write_turn
|
||||
|
||||
|
||||
def _raw_query(tokenizer, text: str, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
encoded = tokenizer(text, add_special_tokens=False, return_tensors="pt")
|
||||
input_ids = encoded["input_ids"].to(device)
|
||||
attention_mask = encoded.get("attention_mask")
|
||||
if attention_mask is None:
|
||||
attention_mask = torch.ones_like(input_ids)
|
||||
return input_ids, attention_mask.to(device)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _answer(model, tokenizer, query: str, max_new_tokens: int) -> dict[str, object]:
|
||||
device = model._find_layer_device()
|
||||
encoded = _chat_tensor(tokenizer, query)
|
||||
encoded = {key: value.to(device) for key, value in encoded.items()}
|
||||
query_ids, query_mask = _raw_query(tokenizer, query, device)
|
||||
prefix_ids, prefix_mask, prefix_length = model._build_text_prefix(
|
||||
query_ids,
|
||||
query_mask,
|
||||
query_text=query,
|
||||
)
|
||||
decisions = [dict(item) for item in model.runtime.v2_last_decisions]
|
||||
output = model.generate(
|
||||
**encoded,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
memory_query_input_ids=query_ids,
|
||||
memory_query_attention_mask=query_mask,
|
||||
memory_query_text=query,
|
||||
use_cache=False,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
prompt_length = int(encoded["input_ids"].shape[1])
|
||||
response = tokenizer.decode(
|
||||
output[0, prompt_length:].detach().cpu().tolist(),
|
||||
skip_special_tokens=True,
|
||||
).strip()
|
||||
return {
|
||||
"response": response,
|
||||
"prefix_length": int(prefix_length),
|
||||
"prefix_used": bool(model.runtime.text_prefix_used),
|
||||
"decision": decisions,
|
||||
"prefix_shape": list(prefix_ids.shape) if isinstance(prefix_ids, torch.Tensor) else None,
|
||||
"prefix_mask_shape": list(prefix_mask.shape) if isinstance(prefix_mask, torch.Tensor) else None,
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--model-path",
|
||||
default=r"W:\Flash\model\dynamic_memory_lab\qwen3_5_4b_natural_memory_v2",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--report",
|
||||
default=r"W:\Flash\model\dynamic_memory_lab\natural_memory_v2_restart_test.json",
|
||||
)
|
||||
parser.add_argument("--max-new-tokens", type=int, default=16)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
parser.add_argument(
|
||||
"--keep-memory",
|
||||
action="store_true",
|
||||
help="leave the controlled test fact embedded after the test",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
model_path = Path(args.model_path)
|
||||
if not (model_path / "memory_merge.json").exists():
|
||||
raise FileNotFoundError(f"Natural Memory v2 package not found: {model_path}")
|
||||
|
||||
fact = "我正在开发一个长期项目，项目内部代号是NM-V2-RESTART，使用中文。"
|
||||
query = "我正在开发的长期项目内部代号是什么？只回答代号。"
|
||||
expected = "NM-V2-RESTART"
|
||||
tokenizer = load_tokenizer(model_path)
|
||||
|
||||
first = load_qwen_dynamic(model_path, load_in_4bit=not args.no_4bit)
|
||||
first.eval()
|
||||
device = first._find_layer_device()
|
||||
# This is a destructive reset of the selected package's durable memory,
|
||||
# so the command is an explicit test tool rather than a chat startup hook.
|
||||
first.reset_memory(batch_size=1, device=device)
|
||||
changed = _write_turn(first, tokenizer, fact, device, force_write=False)
|
||||
before_save = first.memory_v2_stats()
|
||||
first.save_embedded_memory_weights(model_path)
|
||||
del first
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
restarted = load_qwen_dynamic(model_path, load_in_4bit=not args.no_4bit)
|
||||
restarted.eval()
|
||||
after_restart = _answer(restarted, tokenizer, query, args.max_new_tokens)
|
||||
selected_text = " ".join(
|
||||
str(item.get("text", ""))
|
||||
for item in restarted.runtime.v2_last_decisions
|
||||
if isinstance(item, dict)
|
||||
)
|
||||
route_recalled = expected in selected_text or (
|
||||
bool(after_restart["prefix_used"]) and int(after_restart["prefix_length"]) > 0
|
||||
)
|
||||
|
||||
after_cleanup = None
|
||||
if not args.keep_memory:
|
||||
restarted.reset_memory(batch_size=1, device=restarted._find_layer_device())
|
||||
restarted.save_embedded_memory_weights(model_path)
|
||||
after_cleanup = restarted.memory_v2_stats()
|
||||
|
||||
report = {
|
||||
"model_path": str(model_path),
|
||||
"history_passed_to_restart": False,
|
||||
"fact": fact,
|
||||
"query": query,
|
||||
"expected": expected,
|
||||
"automatic_write": True,
|
||||
"write_changed": bool(changed),
|
||||
"before_save": before_save,
|
||||
"after_restart": after_restart,
|
||||
"router_recalled_after_restart": bool(route_recalled),
|
||||
"generated_contains_expected": expected in str(after_restart["response"]),
|
||||
"cleanup_applied": not args.keep_memory,
|
||||
"after_cleanup": after_cleanup,
|
||||
}
|
||||
report_path = Path(args.report)
|
||||
report_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(json.dumps(report, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,226 @@
|
||||
"""End-to-end test: write a fact in a dialogue, restart without chat history, recall it."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from .qwen_integration import (
|
||||
DEFAULT_MEMORY_RESET_TOKEN,
|
||||
load_memory_config,
|
||||
load_qwen_dynamic,
|
||||
load_tokenizer,
|
||||
resolve_memory_reset_token,
|
||||
)
|
||||
|
||||
|
||||
def _chat_tensor(tokenizer, messages, *, add_generation_prompt: bool):
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
add_generation_prompt=add_generation_prompt,
|
||||
return_tensors="pt",
|
||||
return_dict=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
return {
|
||||
key: value
|
||||
for key, value in encoded.items()
|
||||
if isinstance(value, torch.Tensor)
|
||||
}
|
||||
|
||||
|
||||
def _memory_system_prefix(tokenizer, content: str):
|
||||
"""Encode a valid system-message prefix without adding a fake query."""
|
||||
|
||||
full = tokenizer.apply_chat_template(
|
||||
[
|
||||
{"role": "system", "content": content},
|
||||
{"role": "user", "content": "__memory_query_boundary__"},
|
||||
],
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
return_tensors="pt",
|
||||
return_dict=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
input_ids = full["input_ids"]
|
||||
im_start = tokenizer.convert_tokens_to_ids("<|im_start|>")
|
||||
positions = (input_ids[0] == int(im_start)).nonzero(as_tuple=False).flatten()
|
||||
if positions.numel() < 2:
|
||||
raise RuntimeError("could not locate the system/user memory boundary")
|
||||
end = int(positions[1].item())
|
||||
return {
|
||||
"input_ids": input_ids[:, :end],
|
||||
"attention_mask": torch.ones((1, end), dtype=torch.long),
|
||||
}
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _generate(model, tokenizer, messages, max_new_tokens: int) -> str:
|
||||
encoded = _chat_tensor(tokenizer, messages, add_generation_prompt=True)
|
||||
device = model._find_layer_device()
|
||||
encoded = {key: value.to(device) for key, value in encoded.items()}
|
||||
output = model.generate(
|
||||
**encoded,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
use_cache=False,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
response_ids = output[0, encoded["input_ids"].shape[1] :]
|
||||
return tokenizer.decode(response_ids.detach().cpu().tolist(), skip_special_tokens=True).strip()
|
||||
|
||||
|
||||
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_native_v3")
|
||||
parser.add_argument(
|
||||
"--output-adapter",
|
||||
default="dynamic_memory_lab/qwen_memory_adapter_restart_dialogue_test",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--report",
|
||||
default="dynamic_memory_lab/restart_memory_dialogue_test.json",
|
||||
)
|
||||
parser.add_argument("--max-new-tokens", type=int, default=16)
|
||||
parser.add_argument("--text-memory-threshold", type=float, default=0.0)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
fact = "请记住：我的工作地点代号是R7。"
|
||||
acknowledgement = "好的，我会记住这条个人信息。"
|
||||
query = "我的工作地点代号是什么？"
|
||||
expected = "R7"
|
||||
|
||||
tokenizer = load_tokenizer(args.model_path)
|
||||
config = load_memory_config(args.adapter)
|
||||
config.persistent_memory = True
|
||||
config.natural_language_memory = True
|
||||
config.text_memory_threshold = args.text_memory_threshold
|
||||
config.reset_token_id = resolve_memory_reset_token(tokenizer, DEFAULT_MEMORY_RESET_TOKEN)
|
||||
|
||||
model = load_qwen_dynamic(
|
||||
args.model_path,
|
||||
memory_config=config,
|
||||
load_in_4bit=not args.no_4bit,
|
||||
)
|
||||
model.load_memory_adapter(args.adapter)
|
||||
model.reset_memory()
|
||||
device = model._find_layer_device()
|
||||
|
||||
# This is the write turn: a normal user/assistant dialogue, not a prebuilt
|
||||
# memory tensor and not a query with the answer included in the prompt.
|
||||
write_messages = [
|
||||
{"role": "user", "content": fact},
|
||||
{"role": "assistant", "content": acknowledgement},
|
||||
]
|
||||
write_inputs = _chat_tensor(tokenizer, write_messages, add_generation_prompt=False)
|
||||
write_inputs = {key: value.to(device) for key, value in write_inputs.items()}
|
||||
memory_text = _memory_system_prefix(
|
||||
tokenizer,
|
||||
"以下是与当前用户相关的已保存长期记忆。仅在问题相关时使用，不要编造：\n" + fact,
|
||||
)
|
||||
memory_text_input_ids = memory_text["input_ids"].to(device)
|
||||
memory_text_attention_mask = memory_text.get("attention_mask")
|
||||
if memory_text_attention_mask is None:
|
||||
memory_text_attention_mask = torch.ones_like(memory_text_input_ids)
|
||||
memory_text_attention_mask = memory_text_attention_mask.to(device)
|
||||
memory_key = tokenizer(fact, add_special_tokens=False, return_tensors="pt")
|
||||
memory_key_input_ids = memory_key["input_ids"].to(device)
|
||||
memory_key_attention_mask = memory_key.get("attention_mask")
|
||||
if memory_key_attention_mask is None:
|
||||
memory_key_attention_mask = torch.ones_like(memory_key_input_ids)
|
||||
memory_key_attention_mask = memory_key_attention_mask.to(device)
|
||||
model(
|
||||
**write_inputs,
|
||||
read_memory=True,
|
||||
update_memory=True,
|
||||
return_memory=True,
|
||||
use_cache=False,
|
||||
memory_text_input_ids=memory_text_input_ids,
|
||||
memory_text_attention_mask=memory_text_attention_mask,
|
||||
memory_key_input_ids=memory_key_input_ids,
|
||||
memory_key_attention_mask=memory_key_attention_mask,
|
||||
)
|
||||
saved_norm = float(model.runtime.state.detach().float().norm())
|
||||
output_adapter = Path(args.output_adapter)
|
||||
model.save_persistent_memory_checkpoint(output_adapter)
|
||||
|
||||
# Destroy the first model completely. The second model receives only the
|
||||
# base model plus the persistent adapter checkpoint; no chat history or
|
||||
# runtime memory_state file is passed to it.
|
||||
del model
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
restart_config = load_memory_config(output_adapter)
|
||||
restarted = load_qwen_dynamic(
|
||||
args.model_path,
|
||||
memory_config=restart_config,
|
||||
load_in_4bit=not args.no_4bit,
|
||||
)
|
||||
restarted.load_memory_adapter(output_adapter)
|
||||
restarted.eval()
|
||||
restart_norm = float(restarted.runtime.state.detach().float().norm())
|
||||
generated_after_restart = _generate(
|
||||
restarted,
|
||||
tokenizer,
|
||||
[{"role": "user", "content": query}],
|
||||
args.max_new_tokens,
|
||||
)
|
||||
|
||||
# Also verify that the external reset token clears the model-owned state.
|
||||
reset_inputs = _chat_tensor(
|
||||
tokenizer,
|
||||
[{"role": "user", "content": DEFAULT_MEMORY_RESET_TOKEN}],
|
||||
add_generation_prompt=True,
|
||||
)
|
||||
reset_inputs = {key: value.to(restarted._find_layer_device()) for key, value in reset_inputs.items()}
|
||||
restarted.generate(
|
||||
**reset_inputs,
|
||||
max_new_tokens=1,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
use_cache=False,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
reset_norm = float(restarted.runtime.state.detach().float().norm())
|
||||
generated_after_reset = _generate(
|
||||
restarted,
|
||||
tokenizer,
|
||||
[{"role": "user", "content": query}],
|
||||
args.max_new_tokens,
|
||||
)
|
||||
|
||||
report = {
|
||||
"fact_dialogue": write_messages,
|
||||
"restart_query": [{"role": "user", "content": query}],
|
||||
"history_passed_to_restart": False,
|
||||
"expected": expected,
|
||||
"generated_after_restart": generated_after_restart,
|
||||
"recalled_after_restart": expected in generated_after_restart,
|
||||
"saved_memory_norm": saved_norm,
|
||||
"loaded_memory_norm_after_restart": restart_norm,
|
||||
"persistent_adapter": str(output_adapter),
|
||||
"reset_token": DEFAULT_MEMORY_RESET_TOKEN,
|
||||
"reset_token_id": restart_config.reset_token_id,
|
||||
"memory_norm_after_reset": reset_norm,
|
||||
"reset_cleared_memory": reset_norm < 1e-5,
|
||||
"generated_after_reset": generated_after_reset,
|
||||
}
|
||||
report_path = Path(args.report)
|
||||
report_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(json.dumps(report, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,309 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
from tempfile import TemporaryDirectory
|
||||
|
||||
import torch
|
||||
|
||||
from dynamic_memory_lab.memory_os_v2 import (
|
||||
KVBudgetManagerV2,
|
||||
MemoryOSV2,
|
||||
MemoryRouterV2,
|
||||
PagedMemoryBankV2,
|
||||
STATUS_ACTIVE,
|
||||
STATUS_QUARANTINED,
|
||||
STATUS_SUPERSEDED,
|
||||
)
|
||||
from dynamic_memory_lab.tiered_memory_store_v2 import TieredMemoryStoreV2
|
||||
|
||||
|
||||
class MemoryOSV2Test(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
torch.manual_seed(7)
|
||||
self.router = MemoryRouterV2(16, router_dim=8, num_heads=2, max_hops=3)
|
||||
self.bank = PagedMemoryBankV2(
|
||||
16,
|
||||
router=self.router,
|
||||
page_capacity=2,
|
||||
max_pages=512,
|
||||
hot_pages=2,
|
||||
top_k_pages=2,
|
||||
top_k_records=4,
|
||||
max_hops=3,
|
||||
coarse_index_bits=8,
|
||||
)
|
||||
|
||||
def test_router_scores_and_compressed_address(self) -> None:
|
||||
query = torch.randn(4, 16)
|
||||
candidates = torch.randn(4, 5, 16)
|
||||
output = self.router(query, candidates)
|
||||
self.assertEqual(tuple(output["scores"].shape), (4, 5))
|
||||
self.assertEqual(tuple(output["head_scores"].shape), (4, 5, 2))
|
||||
self.assertEqual(tuple(self.router.encode_key(query).shape), (4, 8))
|
||||
|
||||
def test_write_version_and_conflict_resolution(self) -> None:
|
||||
first, first_action = self.bank.write(
|
||||
text="我住在上海",
|
||||
key=torch.randn(16),
|
||||
entity="user",
|
||||
attribute="city",
|
||||
value="上海",
|
||||
confidence=0.9,
|
||||
)
|
||||
second, second_action = self.bank.write(
|
||||
text="我搬到了杭州",
|
||||
key=torch.randn(16),
|
||||
entity="user",
|
||||
attribute="city",
|
||||
value="杭州",
|
||||
confidence=0.95,
|
||||
)
|
||||
self.assertEqual(first_action, "inserted")
|
||||
self.assertEqual(second_action, "updated")
|
||||
self.assertEqual(first.status, STATUS_SUPERSEDED)
|
||||
self.assertEqual(second.status, STATUS_ACTIVE)
|
||||
self.assertEqual(second.version, 1)
|
||||
self.assertEqual(self.bank.active_by_conflict["user::city"], second.record_id)
|
||||
|
||||
def test_quarantine_and_approval(self) -> None:
|
||||
record, action = self.bank.write(
|
||||
text="未经确认的推断",
|
||||
key=torch.randn(16),
|
||||
trusted=False,
|
||||
confidence=0.1,
|
||||
)
|
||||
self.assertEqual(action, "quarantined")
|
||||
self.assertEqual(record.status, STATUS_QUARANTINED)
|
||||
self.assertNotIn(record.record_id, self.bank.records)
|
||||
approved = self.bank.approve(record.record_id)
|
||||
self.assertEqual(approved.status, STATUS_ACTIVE)
|
||||
self.assertIn(approved.record_id, self.bank.records)
|
||||
|
||||
def test_multi_hop_and_slot_replacement(self) -> None:
|
||||
second, _ = self.bank.write(text="项目的第二个节点", key=torch.randn(16), slot_index=2)
|
||||
third, _ = self.bank.write(text="项目的第三个节点", key=torch.randn(16), slot_index=3)
|
||||
first, _ = self.bank.write(
|
||||
text="项目的第一个节点",
|
||||
key=torch.randn(16),
|
||||
related_ids=[second.record_id, third.record_id],
|
||||
slot_index=1,
|
||||
)
|
||||
replacement, action = self.bank.write(
|
||||
text="项目的第一个节点修正版",
|
||||
key=torch.randn(16),
|
||||
related_ids=[second.record_id],
|
||||
slot_index=1,
|
||||
)
|
||||
self.assertEqual(action, "updated")
|
||||
self.assertEqual(first.status, STATUS_SUPERSEDED)
|
||||
self.assertEqual(replacement.status, STATUS_ACTIVE)
|
||||
records, decision = self.bank.query(
|
||||
query_key=replacement.key,
|
||||
top_k_pages=2,
|
||||
top_k_records=4,
|
||||
max_hops=3,
|
||||
)
|
||||
ids = {record.record_id for record in records}
|
||||
self.assertIn(replacement.record_id, ids)
|
||||
self.assertGreaterEqual(decision.hop_count, 1)
|
||||
|
||||
def test_coarse_index_bounds_candidate_pages(self) -> None:
|
||||
for index in range(300):
|
||||
key = torch.zeros(16)
|
||||
key[index % 16] = 1.0
|
||||
key[(index * 7 + 3) % 16] += 0.05
|
||||
self.bank.write(text=f"memory-{index}", key=key, importance=0.2)
|
||||
query_key = torch.zeros(16)
|
||||
query_key[3] = 1.0
|
||||
self.bank.query(query_key=query_key, top_k_pages=2, top_k_records=2)
|
||||
stats = self.bank.stats()
|
||||
self.assertGreater(stats["pages"], 128)
|
||||
self.assertLess(stats["last_coarse_candidates"], stats["pages"])
|
||||
|
||||
def test_export_and_restore(self) -> None:
|
||||
record, _ = self.bank.write(
|
||||
text="可持久化事实",
|
||||
key=torch.randn(16),
|
||||
token_ids=torch.tensor([4, 5, 6]),
|
||||
token_mask=torch.tensor([True, True, True]),
|
||||
)
|
||||
payload = self.bank.export_payload()
|
||||
restored = PagedMemoryBankV2.from_payload(payload, router=self.router)
|
||||
self.assertEqual(restored.stats()["active_records"], 1)
|
||||
self.assertTrue(torch.equal(restored.records[record.record_id].token_ids, torch.tensor([4, 5, 6])))
|
||||
self.assertEqual(restored.records[record.record_id].page_id, record.page_id)
|
||||
|
||||
def test_lazy_capacity_is_bounded(self) -> None:
|
||||
bank = PagedMemoryBankV2(
|
||||
16,
|
||||
router=self.router,
|
||||
page_capacity=1,
|
||||
max_pages=2,
|
||||
hot_pages=0,
|
||||
coarse_index_bits=8,
|
||||
)
|
||||
bank.write(text="容量一", key=torch.randn(16))
|
||||
bank.write(text="容量二", key=torch.randn(16))
|
||||
self.assertEqual(bank.stats()["pages"], 2)
|
||||
with self.assertRaises(RuntimeError):
|
||||
bank.write(text="容量三", key=torch.randn(16))
|
||||
|
||||
def test_memory_os_and_kv_budget(self) -> None:
|
||||
os_v2 = MemoryOSV2(16, router=self.router)
|
||||
record, action = os_v2.write(
|
||||
text="可靠事实",
|
||||
key=torch.randn(16),
|
||||
importance=0.9,
|
||||
confidence=0.9,
|
||||
)
|
||||
self.assertEqual(action, "inserted")
|
||||
self.assertIn(record.record_id, os_v2.bank.records)
|
||||
budget = KVBudgetManagerV2(max_tokens=128, hard_max_tokens=512, keep_recent_tokens=32)
|
||||
self.assertFalse(budget.needs_compaction(100))
|
||||
self.assertTrue(budget.needs_compaction(120))
|
||||
self.assertEqual(budget.overflow(140), 12)
|
||||
|
||||
def test_batch_context_records_keep_all_chunks_active(self) -> None:
|
||||
os_v2 = MemoryOSV2(16, router=self.router)
|
||||
output = os_v2.write_batch(
|
||||
[
|
||||
{
|
||||
"text": "context_chunk:0:0:4",
|
||||
"key": torch.randn(16),
|
||||
"memory_type": "context_chunk",
|
||||
"importance": 0.55,
|
||||
"confidence": 0.8,
|
||||
"trusted": True,
|
||||
"force": True,
|
||||
},
|
||||
{
|
||||
"text": "context_chunk:0:4:8",
|
||||
"key": torch.randn(16),
|
||||
"memory_type": "context_chunk",
|
||||
"importance": 0.55,
|
||||
"confidence": 0.8,
|
||||
"trusted": True,
|
||||
"force": True,
|
||||
},
|
||||
]
|
||||
)
|
||||
self.assertEqual(len(output), 2)
|
||||
self.assertEqual(os_v2.stats()["active_records"], 2)
|
||||
|
||||
def test_management_list_edit_retract_and_audit(self) -> None:
|
||||
os_v2 = MemoryOSV2(16, router=self.router)
|
||||
record, _ = os_v2.write(
|
||||
text="用户喜欢蓝色",
|
||||
key=torch.randn(16),
|
||||
entity="user",
|
||||
attribute="color",
|
||||
value="蓝色",
|
||||
confidence=0.95,
|
||||
importance=0.9,
|
||||
)
|
||||
listed = os_v2.list_records(query_text="蓝色", status="active", limit=10)
|
||||
self.assertEqual([item.record_id for item in listed], [record.record_id])
|
||||
edited = os_v2.edit_record(
|
||||
record.record_id,
|
||||
text="用户喜欢绿色",
|
||||
entity="user",
|
||||
attribute="color",
|
||||
value="绿色",
|
||||
evidence=["user_correction"],
|
||||
)
|
||||
self.assertEqual(edited.version, 1)
|
||||
self.assertEqual(edited.supersedes, record.record_id)
|
||||
self.assertEqual(os_v2.bank.records[record.record_id].status, STATUS_SUPERSEDED)
|
||||
self.assertEqual(os_v2.list_records(query_text="绿色")[0].record_id, edited.record_id)
|
||||
os_v2.retract_record(edited.record_id)
|
||||
self.assertEqual(os_v2.bank.records[edited.record_id].status, "retracted")
|
||||
audit = os_v2.audit()
|
||||
self.assertTrue(audit["healthy"], audit)
|
||||
all_records = os_v2.list_records(status="all", limit=10)
|
||||
self.assertEqual(len(all_records), 2)
|
||||
|
||||
def test_tiered_storage_restarts_and_evicts_cold_records(self) -> None:
|
||||
with TemporaryDirectory() as directory:
|
||||
path = f"{directory}/memory.sqlite"
|
||||
store = TieredMemoryStoreV2(path, key_dim=8, page_capacity=2)
|
||||
bank = PagedMemoryBankV2(
|
||||
16,
|
||||
router=self.router,
|
||||
page_capacity=2,
|
||||
max_pages=64,
|
||||
hot_pages=1,
|
||||
top_k_pages=2,
|
||||
top_k_records=2,
|
||||
tier_store=store,
|
||||
max_resident_pages=1,
|
||||
coarse_index_bits=8,
|
||||
)
|
||||
for index in range(8):
|
||||
key = torch.zeros(16)
|
||||
key[index % 8] = 1.0
|
||||
bank.write(
|
||||
text=f"tiered-memory-{index}",
|
||||
key=key,
|
||||
entity="user",
|
||||
attribute=f"attr-{index}",
|
||||
value=f"value-{index}",
|
||||
importance=0.1 if index < 7 else 1.0,
|
||||
confidence=0.95,
|
||||
)
|
||||
stats = bank.stats()
|
||||
self.assertEqual(stats["storage_mode"], "tiered")
|
||||
self.assertGreaterEqual(stats["pages"], 4)
|
||||
self.assertGreater(stats["cold_pages"], 0)
|
||||
self.assertLess(stats["resident_records"], stats["records"])
|
||||
store.close()
|
||||
|
||||
reopened_store = TieredMemoryStoreV2(path, key_dim=8, page_capacity=2)
|
||||
reopened = PagedMemoryBankV2(
|
||||
16,
|
||||
router=self.router,
|
||||
page_capacity=2,
|
||||
max_pages=64,
|
||||
hot_pages=1,
|
||||
top_k_pages=2,
|
||||
top_k_records=2,
|
||||
tier_store=reopened_store,
|
||||
max_resident_pages=1,
|
||||
coarse_index_bits=8,
|
||||
)
|
||||
records, decision = reopened.query(
|
||||
query_key=torch.nn.functional.one_hot(torch.tensor(3), num_classes=16).float(),
|
||||
query_text="tiered-memory-3",
|
||||
top_k_pages=2,
|
||||
top_k_records=2,
|
||||
)
|
||||
self.assertTrue(records)
|
||||
self.assertTrue(any(item.text == "tiered-memory-3" for item in records))
|
||||
self.assertGreaterEqual(decision.hop_count, 1)
|
||||
quarantined, action = reopened.write(
|
||||
text="待审批事实",
|
||||
key=torch.randn(16),
|
||||
trusted=False,
|
||||
confidence=0.1,
|
||||
)
|
||||
self.assertEqual(action, "quarantined")
|
||||
reopened_store.close()
|
||||
|
||||
final_store = TieredMemoryStoreV2(path, key_dim=8, page_capacity=2)
|
||||
final_bank = PagedMemoryBankV2(
|
||||
16,
|
||||
router=self.router,
|
||||
page_capacity=2,
|
||||
max_pages=64,
|
||||
hot_pages=1,
|
||||
tier_store=final_store,
|
||||
max_resident_pages=2,
|
||||
coarse_index_bits=8,
|
||||
)
|
||||
self.assertIn(quarantined.record_id, final_bank.quarantine)
|
||||
approved = final_bank.approve(quarantined.record_id)
|
||||
self.assertEqual(approved.status, STATUS_ACTIVE)
|
||||
final_store.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,34 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from dynamic_memory_lab.model import DynamicMemoryConfig, DynamicMemoryLM
|
||||
from dynamic_memory_lab.tasks import sample_associative_batch
|
||||
|
||||
|
||||
class DynamicMemoryModelTest(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.device = torch.device("cpu")
|
||||
self.config = DynamicMemoryConfig(vocab_size=32, max_seq_len=16, d_model=32, n_layers=1, n_heads=4, memory_slots=2)
|
||||
self.model = DynamicMemoryLM(self.config).to(self.device)
|
||||
|
||||
def test_shapes_and_loss(self) -> None:
|
||||
batch = sample_associative_batch(batch_size=3, vocab_size=self.config.vocab_size, device=self.device)
|
||||
memory = self.model(batch.learn_chunks[0]).memory
|
||||
output = self.model(batch.query_input, memory=memory, update_memory=False, labels=batch.query_labels)
|
||||
self.assertEqual(tuple(output.logits.shape), (3, 2, self.config.vocab_size))
|
||||
self.assertEqual(tuple(output.memory.shape), (3, self.config.memory_slots, self.config.d_model))
|
||||
self.assertIsNotNone(output.loss)
|
||||
output.loss.backward()
|
||||
|
||||
def test_memory_changes_after_learning_chunk(self) -> None:
|
||||
batch = sample_associative_batch(batch_size=2, vocab_size=self.config.vocab_size, device=self.device)
|
||||
initial = self.model.memory.initial_state(2, device=self.device, dtype=torch.float32)
|
||||
updated = self.model(batch.learn_chunks[0], memory=initial, update_memory=True).memory
|
||||
self.assertFalse(torch.allclose(initial, updated))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,142 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from dynamic_memory_lab.qwen_integration import (
|
||||
AutomaticMemoryPolicy,
|
||||
MemoryLayerAdapter,
|
||||
NaturalLanguageRetriever,
|
||||
NativeQwenDynamicMemory,
|
||||
QwenMemoryConfig,
|
||||
QwenDynamicMemory,
|
||||
_MemoryRuntime,
|
||||
looks_like_question,
|
||||
split_memory_candidates,
|
||||
)
|
||||
|
||||
|
||||
class _FakeAttention(nn.Module):
|
||||
def __init__(self, *, fail_if_called: bool = False) -> None:
|
||||
super().__init__()
|
||||
self.fail_if_called = fail_if_called
|
||||
self.called = False
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor, **kwargs):
|
||||
self.called = True
|
||||
if self.fail_if_called:
|
||||
raise AssertionError("original token mixer was called in replace mode")
|
||||
return hidden_states * 2.0, None
|
||||
|
||||
|
||||
class _FakeQwenLayer(nn.Module):
|
||||
layer_type = "full_attention"
|
||||
|
||||
def __init__(self, *, fail_if_called: bool = False) -> None:
|
||||
super().__init__()
|
||||
self.input_layernorm = nn.Identity()
|
||||
self.post_attention_layernorm = nn.Identity()
|
||||
self.self_attn = _FakeAttention(fail_if_called=fail_if_called)
|
||||
self.mlp = nn.Identity()
|
||||
|
||||
def forward(self, hidden_states, position_embeddings=None, attention_mask=None, position_ids=None, past_key_values=None, **kwargs):
|
||||
return hidden_states + self.self_attn(hidden_states)[0]
|
||||
|
||||
|
||||
class QwenSurgeryTest(unittest.TestCase):
|
||||
def _runtime(self):
|
||||
memory = QwenDynamicMemory(
|
||||
hidden_size=8,
|
||||
config=QwenMemoryConfig(memory_slots=2, memory_dim=4),
|
||||
)
|
||||
runtime = _MemoryRuntime(memory)
|
||||
runtime.state = memory.initial_state(1, device=torch.device("cpu"))
|
||||
return runtime
|
||||
|
||||
def test_blend_exposes_trainable_mixer_weight(self) -> None:
|
||||
runtime = self._runtime()
|
||||
adapter = MemoryLayerAdapter(
|
||||
_FakeQwenLayer(),
|
||||
runtime,
|
||||
read=True,
|
||||
write=False,
|
||||
mode="blend",
|
||||
blend_init=0.5,
|
||||
)
|
||||
output = adapter(torch.ones(1, 3, 8))
|
||||
output.sum().backward()
|
||||
self.assertEqual(tuple(output.shape), (1, 3, 8))
|
||||
self.assertIsNotNone(adapter.blend_logit.grad)
|
||||
|
||||
def test_replace_skips_original_token_mixer(self) -> None:
|
||||
runtime = self._runtime()
|
||||
layer = _FakeQwenLayer(fail_if_called=True)
|
||||
adapter = MemoryLayerAdapter(layer, runtime, read=True, write=False, mode="replace")
|
||||
output = adapter(torch.ones(1, 3, 8))
|
||||
self.assertEqual(tuple(output.shape), (1, 3, 8))
|
||||
self.assertFalse(layer.self_attn.called)
|
||||
|
||||
def test_raw_token_write_uses_output_projection_row(self) -> None:
|
||||
memory = QwenDynamicMemory(
|
||||
hidden_size=8,
|
||||
config=QwenMemoryConfig(
|
||||
memory_slots=2,
|
||||
memory_dim=4,
|
||||
write_token_offset=2,
|
||||
raw_token_write=True,
|
||||
broadcast_write=True,
|
||||
),
|
||||
)
|
||||
runtime = _MemoryRuntime(memory)
|
||||
runtime.state = memory.initial_state(1, device=torch.device("cpu"))
|
||||
runtime.read_enabled = False
|
||||
runtime.update_enabled = True
|
||||
runtime.input_ids = torch.tensor([[5, 6, 7, 8]])
|
||||
runtime.attention_mask = torch.ones_like(runtime.input_ids)
|
||||
runtime.output_embeddings = nn.Linear(8, 16, bias=False)
|
||||
adapter = MemoryLayerAdapter(_FakeQwenLayer(), runtime, read=True, write=True, mode="residual")
|
||||
|
||||
adapter(torch.ones(1, 4, 8))
|
||||
|
||||
expected = runtime.output_embeddings.weight[7]
|
||||
self.assertIsNotNone(runtime.raw_memory)
|
||||
self.assertTrue(torch.allclose(runtime.raw_memory[0], expected))
|
||||
|
||||
def test_native_controller_exposes_write_forget_and_value_state(self) -> None:
|
||||
memory = NativeQwenDynamicMemory(
|
||||
hidden_size=8,
|
||||
config=QwenMemoryConfig(memory_slots=2, memory_dim=4),
|
||||
)
|
||||
hidden = torch.randn(1, 3, 8)
|
||||
state = memory.initial_state(1, device=torch.device("cpu"))
|
||||
updated = memory.update(hidden, state, attention_mask=torch.ones(1, 5, dtype=torch.long))
|
||||
self.assertEqual(tuple(updated.shape), (1, 2, 4))
|
||||
self.assertEqual(tuple(memory.last_write_probability.shape), (1, 1))
|
||||
self.assertEqual(tuple(memory.last_forget_probability.shape), (1, 2))
|
||||
self.assertEqual(tuple(memory.last_write_summary.shape), (1, 8))
|
||||
self.assertEqual(tuple(memory.last_write_representation.shape), (1, 8))
|
||||
|
||||
def test_natural_language_retriever_scores_single_and_multiple_slots(self) -> None:
|
||||
retriever = NaturalLanguageRetriever(hidden_size=8, projection_size=4)
|
||||
query = torch.randn(2, 8)
|
||||
one_key = torch.randn(2, 8)
|
||||
many_keys = torch.randn(2, 3, 8)
|
||||
self.assertEqual(tuple(retriever(query, one_key).shape), (2,))
|
||||
self.assertEqual(tuple(retriever(query, many_keys).shape), (2, 3))
|
||||
|
||||
def test_automatic_memory_policy_and_candidate_segmentation(self) -> None:
|
||||
policy = AutomaticMemoryPolicy(hidden_size=8)
|
||||
output = policy(torch.randn(3, 8))
|
||||
self.assertEqual(tuple(output.shape), (3,))
|
||||
self.assertEqual(
|
||||
split_memory_candidates("我叫林浩，我正在开发星火项目。"),
|
||||
["我叫林浩，我正在开发星火项目。"],
|
||||
)
|
||||
self.assertTrue(looks_like_question("如果我选择 GPU，会发生什么？"))
|
||||
self.assertFalse(looks_like_question("我住在上海，正在开发星火项目。"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,561 @@
|
||||
"""Disk-backed page storage for Natural Memory v2.
|
||||
|
||||
The neural router never reads this module directly. ``PagedMemoryBankV2``
|
||||
uses it as an owned storage tier when a deployment needs more records than
|
||||
RAM can comfortably retain. Keys and token ids are stored as compact binary
|
||||
blobs; SQLite is used only for durable metadata, page locality and recovery.
|
||||
The backend is deliberately dependency-free beyond PyTorch and the Python
|
||||
standard library.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
import struct
|
||||
import threading
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable, Optional
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
def _pack_tensor(value: Optional[Tensor], *, dtype: str) -> Optional[bytes]:
|
||||
if value is None:
|
||||
return None
|
||||
tensor = value.detach().cpu().contiguous()
|
||||
if dtype == "float32":
|
||||
tensor = tensor.float()
|
||||
raw = tensor.numpy().tobytes()
|
||||
elif dtype == "int32":
|
||||
tensor = tensor.to(torch.int32)
|
||||
raw = tensor.numpy().tobytes()
|
||||
elif dtype == "bool":
|
||||
raw = tensor.bool().numpy().tobytes()
|
||||
else:
|
||||
raise ValueError(f"unsupported tensor dtype: {dtype}")
|
||||
return struct.pack("<I", int(tensor.numel())) + raw
|
||||
|
||||
|
||||
def _unpack_tensor(value: Optional[bytes], *, dtype: str) -> Optional[Tensor]:
|
||||
if value is None:
|
||||
return None
|
||||
if len(value) < 4:
|
||||
raise ValueError("corrupt tensor blob")
|
||||
count = struct.unpack("<I", value[:4])[0]
|
||||
payload = value[4:]
|
||||
if dtype == "float32":
|
||||
item_size = 4
|
||||
tensor = torch.frombuffer(bytearray(payload), dtype=torch.float32).clone()
|
||||
elif dtype == "int32":
|
||||
item_size = 4
|
||||
tensor = torch.frombuffer(bytearray(payload), dtype=torch.int32).clone().to(torch.long)
|
||||
elif dtype == "bool":
|
||||
item_size = 1
|
||||
tensor = torch.frombuffer(bytearray(payload), dtype=torch.uint8).clone().bool()
|
||||
else:
|
||||
raise ValueError(f"unsupported tensor dtype: {dtype}")
|
||||
if len(payload) != count * item_size or tensor.numel() != count:
|
||||
raise ValueError("corrupt tensor blob length")
|
||||
return tensor
|
||||
|
||||
|
||||
class TieredMemoryStoreV2:
|
||||
"""Recoverable page/record store used by the warm and cold tiers."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
path: str | Path,
|
||||
*,
|
||||
key_dim: int,
|
||||
page_capacity: int = 32,
|
||||
) -> None:
|
||||
self.path = Path(path)
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self.key_dim = int(key_dim)
|
||||
self.page_capacity = int(page_capacity)
|
||||
self._lock = threading.RLock()
|
||||
self.connection = sqlite3.connect(
|
||||
str(self.path),
|
||||
check_same_thread=False,
|
||||
isolation_level=None,
|
||||
)
|
||||
self.connection.execute("PRAGMA journal_mode=WAL")
|
||||
self.connection.execute("PRAGMA synchronous=NORMAL")
|
||||
self.connection.execute("PRAGMA temp_store=MEMORY")
|
||||
self.connection.execute("PRAGMA foreign_keys=ON")
|
||||
self._create_schema()
|
||||
|
||||
def _create_schema(self) -> None:
|
||||
with self._lock:
|
||||
self.connection.executescript(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS store_meta (
|
||||
key TEXT PRIMARY KEY,
|
||||
value TEXT NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS pages (
|
||||
page_id TEXT PRIMARY KEY,
|
||||
tier TEXT NOT NULL,
|
||||
capacity INTEGER NOT NULL,
|
||||
record_ids TEXT NOT NULL,
|
||||
key BLOB,
|
||||
summary BLOB,
|
||||
importance REAL NOT NULL,
|
||||
created_at INTEGER NOT NULL,
|
||||
last_access INTEGER NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS records (
|
||||
record_id TEXT PRIMARY KEY,
|
||||
page_id TEXT NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
key BLOB NOT NULL,
|
||||
summary BLOB NOT NULL,
|
||||
memory_type TEXT NOT NULL,
|
||||
entity TEXT NOT NULL,
|
||||
attribute TEXT NOT NULL,
|
||||
value TEXT NOT NULL,
|
||||
timestamp INTEGER NOT NULL,
|
||||
importance REAL NOT NULL,
|
||||
confidence REAL NOT NULL,
|
||||
source TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
version INTEGER NOT NULL,
|
||||
supersedes TEXT NOT NULL,
|
||||
related_ids TEXT NOT NULL,
|
||||
evidence TEXT NOT NULL,
|
||||
slot_index INTEGER NOT NULL,
|
||||
token_ids BLOB,
|
||||
token_mask BLOB,
|
||||
access_count INTEGER NOT NULL,
|
||||
last_access INTEGER NOT NULL,
|
||||
FOREIGN KEY(page_id) REFERENCES pages(page_id) ON DELETE CASCADE
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS records_page_idx ON records(page_id);
|
||||
CREATE INDEX IF NOT EXISTS records_status_idx ON records(status);
|
||||
CREATE INDEX IF NOT EXISTS records_text_idx ON records(text COLLATE NOCASE);
|
||||
CREATE INDEX IF NOT EXISTS records_conflict_idx
|
||||
ON records(entity COLLATE NOCASE, attribute COLLATE NOCASE, status);
|
||||
CREATE TABLE IF NOT EXISTS quarantine (
|
||||
record_id TEXT PRIMARY KEY,
|
||||
payload TEXT NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS coarse_buckets (
|
||||
signature INTEGER NOT NULL,
|
||||
page_id TEXT NOT NULL,
|
||||
PRIMARY KEY(signature, page_id),
|
||||
FOREIGN KEY(page_id) REFERENCES pages(page_id) ON DELETE CASCADE
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS coarse_page_idx ON coarse_buckets(page_id);
|
||||
"""
|
||||
)
|
||||
self.connection.execute(
|
||||
"INSERT OR IGNORE INTO store_meta(key, value) VALUES('format_version', '2')"
|
||||
)
|
||||
self.connection.execute(
|
||||
"INSERT OR REPLACE INTO store_meta(key, value) VALUES('key_dim', ?)",
|
||||
(str(self.key_dim),),
|
||||
)
|
||||
self.connection.execute(
|
||||
"INSERT OR REPLACE INTO store_meta(key, value) VALUES('page_capacity', ?)",
|
||||
(str(self.page_capacity),),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _record_values(record: Any) -> tuple[Any, ...]:
|
||||
return (
|
||||
record.record_id,
|
||||
record.page_id,
|
||||
record.text,
|
||||
_pack_tensor(record.key, dtype="float32"),
|
||||
_pack_tensor(record.summary, dtype="float32"),
|
||||
record.memory_type,
|
||||
record.entity,
|
||||
record.attribute,
|
||||
record.value,
|
||||
int(record.timestamp),
|
||||
float(record.importance),
|
||||
float(record.confidence),
|
||||
record.source,
|
||||
record.status,
|
||||
int(record.version),
|
||||
record.supersedes,
|
||||
json.dumps(record.related_ids, ensure_ascii=False, separators=(",", ":")),
|
||||
json.dumps(record.evidence, ensure_ascii=False, separators=(",", ":")),
|
||||
int(record.slot_index),
|
||||
_pack_tensor(record.token_ids, dtype="int32"),
|
||||
_pack_tensor(record.token_mask, dtype="bool"),
|
||||
int(record.access_count),
|
||||
int(record.last_access),
|
||||
)
|
||||
|
||||
def upsert_page(self, page: Any) -> None:
|
||||
with self._lock:
|
||||
self.connection.execute(
|
||||
"""
|
||||
INSERT INTO pages(page_id, tier, capacity, record_ids, key, summary,
|
||||
importance, created_at, last_access)
|
||||
VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(page_id) DO UPDATE SET
|
||||
tier=excluded.tier,
|
||||
capacity=excluded.capacity,
|
||||
record_ids=excluded.record_ids,
|
||||
key=excluded.key,
|
||||
summary=excluded.summary,
|
||||
importance=excluded.importance,
|
||||
last_access=excluded.last_access
|
||||
""",
|
||||
(
|
||||
page.page_id,
|
||||
page.tier,
|
||||
int(page.capacity),
|
||||
json.dumps(page.record_ids, ensure_ascii=False, separators=(",", ":")),
|
||||
_pack_tensor(page.key, dtype="float32"),
|
||||
_pack_tensor(page.summary, dtype="float32"),
|
||||
float(page.importance),
|
||||
int(page.created_at),
|
||||
int(page.last_access),
|
||||
),
|
||||
)
|
||||
|
||||
def upsert_record(self, record: Any) -> None:
|
||||
with self._lock:
|
||||
self.connection.execute(
|
||||
"""
|
||||
INSERT INTO records(
|
||||
record_id, page_id, text, key, summary, memory_type, entity,
|
||||
attribute, value, timestamp, importance, confidence, source,
|
||||
status, version, supersedes, related_ids, evidence, slot_index,
|
||||
token_ids, token_mask, access_count, last_access
|
||||
) VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(record_id) DO UPDATE SET
|
||||
page_id=excluded.page_id,
|
||||
text=excluded.text,
|
||||
key=excluded.key,
|
||||
summary=excluded.summary,
|
||||
memory_type=excluded.memory_type,
|
||||
entity=excluded.entity,
|
||||
attribute=excluded.attribute,
|
||||
value=excluded.value,
|
||||
timestamp=excluded.timestamp,
|
||||
importance=excluded.importance,
|
||||
confidence=excluded.confidence,
|
||||
source=excluded.source,
|
||||
status=excluded.status,
|
||||
version=excluded.version,
|
||||
supersedes=excluded.supersedes,
|
||||
related_ids=excluded.related_ids,
|
||||
evidence=excluded.evidence,
|
||||
slot_index=excluded.slot_index,
|
||||
token_ids=excluded.token_ids,
|
||||
token_mask=excluded.token_mask,
|
||||
access_count=excluded.access_count,
|
||||
last_access=excluded.last_access
|
||||
""",
|
||||
self._record_values(record),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _quarantine_payload(record: Any) -> str:
|
||||
payload = {
|
||||
"record_id": record.record_id,
|
||||
"text": record.text,
|
||||
"key": record.key.detach().cpu().tolist(),
|
||||
"summary": record.summary.detach().cpu().tolist(),
|
||||
"memory_type": record.memory_type,
|
||||
"entity": record.entity,
|
||||
"attribute": record.attribute,
|
||||
"value": record.value,
|
||||
"timestamp": int(record.timestamp),
|
||||
"importance": float(record.importance),
|
||||
"confidence": float(record.confidence),
|
||||
"source": record.source,
|
||||
"status": record.status,
|
||||
"version": int(record.version),
|
||||
"page_id": record.page_id,
|
||||
"supersedes": record.supersedes,
|
||||
"related_ids": list(record.related_ids),
|
||||
"evidence": list(record.evidence),
|
||||
"slot_index": int(record.slot_index),
|
||||
"token_ids": record.token_ids.detach().cpu().tolist() if record.token_ids is not None else None,
|
||||
"token_mask": record.token_mask.detach().cpu().tolist() if record.token_mask is not None else None,
|
||||
"access_count": int(record.access_count),
|
||||
"last_access": int(record.last_access),
|
||||
}
|
||||
return json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
|
||||
|
||||
def upsert_quarantine(self, record: Any) -> None:
|
||||
with self._lock:
|
||||
self.connection.execute(
|
||||
"INSERT INTO quarantine(record_id, payload) VALUES(?, ?) ON CONFLICT(record_id) DO UPDATE SET payload=excluded.payload",
|
||||
(record.record_id, self._quarantine_payload(record)),
|
||||
)
|
||||
|
||||
def load_quarantine(self) -> list[dict[str, Any]]:
|
||||
with self._lock:
|
||||
rows = self.connection.execute("SELECT payload FROM quarantine ORDER BY record_id").fetchall()
|
||||
output = []
|
||||
for (payload,) in rows:
|
||||
item = json.loads(payload)
|
||||
item["key"] = torch.tensor(item["key"], dtype=torch.float32)
|
||||
item["summary"] = torch.tensor(item["summary"], dtype=torch.float32)
|
||||
if item.get("token_ids") is not None:
|
||||
item["token_ids"] = torch.tensor(item["token_ids"], dtype=torch.long)
|
||||
if item.get("token_mask") is not None:
|
||||
item["token_mask"] = torch.tensor(item["token_mask"], dtype=torch.bool)
|
||||
output.append(item)
|
||||
return output
|
||||
|
||||
def delete_quarantine(self, record_id: str) -> None:
|
||||
with self._lock:
|
||||
self.connection.execute("DELETE FROM quarantine WHERE record_id=?", (record_id,))
|
||||
|
||||
def upsert_page_with_records(self, page: Any, records: Iterable[Any]) -> None:
|
||||
with self._lock:
|
||||
self.connection.execute("BEGIN")
|
||||
try:
|
||||
self.upsert_page(page)
|
||||
for record in records:
|
||||
self.upsert_record(record)
|
||||
self.connection.execute("COMMIT")
|
||||
except BaseException:
|
||||
self.connection.execute("ROLLBACK")
|
||||
raise
|
||||
|
||||
@contextmanager
|
||||
def transaction(self):
|
||||
with self._lock:
|
||||
self.connection.execute("BEGIN")
|
||||
try:
|
||||
yield self
|
||||
self.connection.execute("COMMIT")
|
||||
except BaseException:
|
||||
self.connection.execute("ROLLBACK")
|
||||
raise
|
||||
|
||||
def replace_page_buckets(self, page_id: str, signatures: Iterable[int]) -> None:
|
||||
with self._lock:
|
||||
self.connection.execute("DELETE FROM coarse_buckets WHERE page_id=?", (page_id,))
|
||||
self.connection.executemany(
|
||||
"INSERT OR IGNORE INTO coarse_buckets(signature, page_id) VALUES(?, ?)",
|
||||
[(int(signature), page_id) for signature in set(signatures)],
|
||||
)
|
||||
|
||||
def clear_coarse_buckets(self) -> None:
|
||||
with self._lock:
|
||||
self.connection.execute("DELETE FROM coarse_buckets")
|
||||
|
||||
def record_keys(self) -> list[tuple[str, Tensor]]:
|
||||
"""Return compact record addresses for an explicit index rebuild."""
|
||||
|
||||
with self._lock:
|
||||
rows = self.connection.execute("SELECT record_id, key FROM records").fetchall()
|
||||
return [
|
||||
(str(record_id), _unpack_tensor(blob, dtype="float32"))
|
||||
for record_id, blob in rows
|
||||
]
|
||||
|
||||
def page_headers(self) -> list[dict[str, Any]]:
|
||||
with self._lock:
|
||||
rows = self.connection.execute(
|
||||
"SELECT page_id, tier, capacity, record_ids, key, summary, importance, created_at, last_access FROM pages ORDER BY created_at, page_id"
|
||||
).fetchall()
|
||||
return [
|
||||
{
|
||||
"page_id": row[0],
|
||||
"tier": row[1],
|
||||
"capacity": int(row[2]),
|
||||
"record_ids": list(json.loads(row[3])),
|
||||
"key": _unpack_tensor(row[4], dtype="float32"),
|
||||
"summary": _unpack_tensor(row[5], dtype="float32"),
|
||||
"importance": float(row[6]),
|
||||
"created_at": int(row[7]),
|
||||
"last_access": int(row[8]),
|
||||
}
|
||||
for row in rows
|
||||
]
|
||||
|
||||
def load_records(self, record_ids: Iterable[str]) -> list[dict[str, Any]]:
|
||||
ids = list(dict.fromkeys(str(item) for item in record_ids))
|
||||
if not ids:
|
||||
return []
|
||||
output: list[dict[str, Any]] = []
|
||||
with self._lock:
|
||||
for start in range(0, len(ids), 500):
|
||||
chunk = ids[start : start + 500]
|
||||
placeholders = ",".join("?" for _ in chunk)
|
||||
rows = self.connection.execute(
|
||||
f"SELECT record_id, page_id, text, key, summary, memory_type, entity, attribute, value, timestamp, importance, confidence, source, status, version, supersedes, related_ids, evidence, slot_index, token_ids, token_mask, access_count, last_access FROM records WHERE record_id IN ({placeholders})",
|
||||
chunk,
|
||||
).fetchall()
|
||||
for row in rows:
|
||||
output.append(
|
||||
{
|
||||
"record_id": row[0],
|
||||
"page_id": row[1],
|
||||
"text": row[2],
|
||||
"key": _unpack_tensor(row[3], dtype="float32"),
|
||||
"summary": _unpack_tensor(row[4], dtype="float32"),
|
||||
"memory_type": row[5],
|
||||
"entity": row[6],
|
||||
"attribute": row[7],
|
||||
"value": row[8],
|
||||
"timestamp": int(row[9]),
|
||||
"importance": float(row[10]),
|
||||
"confidence": float(row[11]),
|
||||
"source": row[12],
|
||||
"status": row[13],
|
||||
"version": int(row[14]),
|
||||
"supersedes": row[15],
|
||||
"related_ids": list(json.loads(row[16])),
|
||||
"evidence": list(json.loads(row[17])),
|
||||
"slot_index": int(row[18]),
|
||||
"token_ids": _unpack_tensor(row[19], dtype="int32"),
|
||||
"token_mask": _unpack_tensor(row[20], dtype="bool"),
|
||||
"access_count": int(row[21]),
|
||||
"last_access": int(row[22]),
|
||||
}
|
||||
)
|
||||
return output
|
||||
|
||||
def find_by_text(self, text: str, *, active_status: str) -> Optional[dict[str, Any]]:
|
||||
with self._lock:
|
||||
row = self.connection.execute(
|
||||
"SELECT record_id FROM records WHERE text=? COLLATE NOCASE AND status=? LIMIT 1",
|
||||
(text.strip(), active_status),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
loaded = self.load_records([row[0]])
|
||||
return loaded[0] if loaded else None
|
||||
|
||||
def find_by_conflict(
|
||||
self,
|
||||
entity: str,
|
||||
attribute: str,
|
||||
*,
|
||||
active_status: str,
|
||||
) -> Optional[dict[str, Any]]:
|
||||
with self._lock:
|
||||
row = self.connection.execute(
|
||||
"SELECT record_id FROM records WHERE entity=? COLLATE NOCASE AND attribute=? COLLATE NOCASE AND status=? ORDER BY version DESC LIMIT 1",
|
||||
(entity.strip(), attribute.strip(), active_status),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
loaded = self.load_records([row[0]])
|
||||
return loaded[0] if loaded else None
|
||||
|
||||
def active_conflicts(self, *, active_status: str) -> dict[str, str]:
|
||||
with self._lock:
|
||||
rows = self.connection.execute(
|
||||
"SELECT entity, attribute, record_id FROM records WHERE status=? AND entity<>'' AND attribute<>'' ORDER BY version DESC, last_access DESC",
|
||||
(active_status,),
|
||||
).fetchall()
|
||||
output: dict[str, str] = {}
|
||||
for entity, attribute, record_id in rows:
|
||||
output.setdefault(f"{entity.strip().lower()}::{attribute.strip().lower()}", record_id)
|
||||
return output
|
||||
|
||||
def candidate_page_ids(
|
||||
self,
|
||||
signatures: Iterable[int],
|
||||
*,
|
||||
hot_page_ids: Iterable[str] = (),
|
||||
limit: int = 4096,
|
||||
) -> list[str]:
|
||||
values = list(dict.fromkeys(int(item) for item in signatures))
|
||||
selected: list[str] = []
|
||||
with self._lock:
|
||||
# Preserve the exact bucket before Hamming-neighbor probes. A
|
||||
# single combined IN query is subtly unsafe: SQLite may return
|
||||
# neighbor pages first and truncate the exact match away.
|
||||
for index, signature in enumerate(values):
|
||||
remaining = int(limit) - len(selected)
|
||||
if index == 0:
|
||||
rows = self.connection.execute(
|
||||
"SELECT page_id FROM coarse_buckets WHERE signature=? ORDER BY page_id",
|
||||
(int(signature),),
|
||||
).fetchall()
|
||||
elif remaining > 0:
|
||||
rows = self.connection.execute(
|
||||
"SELECT page_id FROM coarse_buckets WHERE signature=? ORDER BY page_id LIMIT ?",
|
||||
(int(signature), remaining),
|
||||
).fetchall()
|
||||
else:
|
||||
break
|
||||
selected.extend(row[0] for row in rows)
|
||||
selected.extend(str(item) for item in hot_page_ids)
|
||||
selected = list(dict.fromkeys(selected))[:limit]
|
||||
if not selected:
|
||||
rows = self.connection.execute(
|
||||
"SELECT page_id FROM pages ORDER BY last_access DESC, page_id LIMIT ?",
|
||||
(min(128, int(limit)),),
|
||||
).fetchall()
|
||||
selected = [row[0] for row in rows]
|
||||
return selected
|
||||
|
||||
def set_page_tier(self, page_id: str, tier: str) -> None:
|
||||
with self._lock:
|
||||
self.connection.execute("UPDATE pages SET tier=? WHERE page_id=?", (tier, page_id))
|
||||
|
||||
def count(self) -> dict[str, int]:
|
||||
with self._lock:
|
||||
pages = int(self.connection.execute("SELECT COUNT(*) FROM pages").fetchone()[0])
|
||||
records = int(self.connection.execute("SELECT COUNT(*) FROM records").fetchone()[0])
|
||||
status_rows = self.connection.execute(
|
||||
"SELECT status, COUNT(*) FROM records GROUP BY status"
|
||||
).fetchall()
|
||||
cold = int(self.connection.execute("SELECT COUNT(*) FROM pages WHERE tier='cold'").fetchone()[0])
|
||||
warm = int(self.connection.execute("SELECT COUNT(*) FROM pages WHERE tier='warm'").fetchone()[0])
|
||||
hot = int(self.connection.execute("SELECT COUNT(*) FROM pages WHERE tier='hot'").fetchone()[0])
|
||||
quarantined = int(self.connection.execute("SELECT COUNT(*) FROM quarantine").fetchone()[0])
|
||||
output = {
|
||||
"pages": pages,
|
||||
"records": records,
|
||||
"hot_pages": hot,
|
||||
"warm_pages": warm,
|
||||
"cold_pages": cold,
|
||||
"quarantined": quarantined,
|
||||
}
|
||||
output.update({f"status_{status}": int(count) for status, count in status_rows})
|
||||
return output
|
||||
|
||||
def coarse_bucket_count(self) -> int:
|
||||
with self._lock:
|
||||
return int(self.connection.execute("SELECT COUNT(DISTINCT signature) FROM coarse_buckets").fetchone()[0])
|
||||
|
||||
def flush(self) -> None:
|
||||
with self._lock:
|
||||
self.connection.execute("PRAGMA wal_checkpoint(PASSIVE)")
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear durable memory records while preserving the store schema."""
|
||||
|
||||
with self._lock:
|
||||
self.connection.execute("BEGIN")
|
||||
try:
|
||||
self.connection.execute("DELETE FROM coarse_buckets")
|
||||
self.connection.execute("DELETE FROM records")
|
||||
self.connection.execute("DELETE FROM pages")
|
||||
self.connection.execute("DELETE FROM quarantine")
|
||||
self.connection.execute("COMMIT")
|
||||
except BaseException:
|
||||
self.connection.execute("ROLLBACK")
|
||||
raise
|
||||
|
||||
def close(self) -> None:
|
||||
with self._lock:
|
||||
self.flush()
|
||||
self.connection.close()
|
||||
|
||||
def __enter__(self) -> "TieredMemoryStoreV2":
|
||||
return self
|
||||
|
||||
def __exit__(self, *_: Any) -> None:
|
||||
self.close()
|
||||
|
||||
|
||||
__all__ = ["TieredMemoryStoreV2"]
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Train the dynamic memory mechanism on a separated key-value task."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import random
|
||||
from contextlib import nullcontext
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from torch.nn.utils import clip_grad_norm_
|
||||
|
||||
from .model import DynamicMemoryConfig, DynamicMemoryLM, count_parameters
|
||||
from .tasks import sample_associative_batch
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--steps", type=int, default=1000)
|
||||
parser.add_argument("--batch-size", type=int, default=64)
|
||||
parser.add_argument("--lr", type=float, default=3e-4)
|
||||
parser.add_argument("--eval-every", type=int, default=100)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--device", default="auto", choices=("auto", "cpu", "cuda"))
|
||||
parser.add_argument("--overwrite", action="store_true")
|
||||
parser.add_argument("--checkpoint-dir", default="dynamic_memory_lab/checkpoints")
|
||||
parser.add_argument("--resume", default=None, help="path to a checkpoint produced by this script")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def choose_device(requested: str) -> torch.device:
|
||||
if requested == "cuda" and not torch.cuda.is_available():
|
||||
raise RuntimeError("CUDA was requested but is not available")
|
||||
if requested == "auto":
|
||||
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
return torch.device(requested)
|
||||
|
||||
|
||||
def evaluate(model: DynamicMemoryLM, *, device: torch.device, overwrite: bool, batches: int = 20) -> tuple[float, float]:
|
||||
model.eval()
|
||||
total_loss = 0.0
|
||||
total_correct = 0
|
||||
total_items = 0
|
||||
with torch.no_grad():
|
||||
for _ in range(batches):
|
||||
batch = sample_associative_batch(
|
||||
batch_size=256,
|
||||
vocab_size=model.config.vocab_size,
|
||||
device=device,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
memory = None
|
||||
for chunk in batch.learn_chunks:
|
||||
memory = model(chunk, memory=memory, update_memory=True).memory
|
||||
output = model(batch.query_input, memory=memory, update_memory=False, labels=batch.query_labels)
|
||||
total_loss += float(output.loss)
|
||||
prediction = output.logits[:, 0].argmax(dim=-1)
|
||||
total_correct += int((prediction == batch.expected).sum())
|
||||
total_items += batch.expected.numel()
|
||||
model.train()
|
||||
return total_loss / batches, total_correct / total_items
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
random.seed(args.seed)
|
||||
torch.manual_seed(args.seed)
|
||||
device = choose_device(args.device)
|
||||
if device.type == "cuda":
|
||||
torch.cuda.manual_seed_all(args.seed)
|
||||
|
||||
config = DynamicMemoryConfig()
|
||||
model = DynamicMemoryLM(config).to(device)
|
||||
optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.01)
|
||||
use_amp = device.type == "cuda"
|
||||
checkpoint_dir = Path(args.checkpoint_dir)
|
||||
checkpoint_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
start_step = 0
|
||||
if args.resume:
|
||||
checkpoint = torch.load(args.resume, map_location=device, weights_only=False)
|
||||
model.load_state_dict(checkpoint["model"])
|
||||
optimizer.load_state_dict(checkpoint["optimizer"])
|
||||
start_step = int(checkpoint.get("step", 0))
|
||||
print(f"resumed_from={args.resume} step={start_step}")
|
||||
|
||||
print(f"device={device} parameters={count_parameters(model):,} overwrite={args.overwrite}")
|
||||
for step in range(start_step + 1, start_step + args.steps + 1):
|
||||
model.train()
|
||||
batch = sample_associative_batch(
|
||||
batch_size=args.batch_size,
|
||||
vocab_size=config.vocab_size,
|
||||
device=device,
|
||||
overwrite=args.overwrite,
|
||||
)
|
||||
memory = None
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
amp_context = torch.autocast(device_type="cuda", dtype=torch.bfloat16) if use_amp else nullcontext()
|
||||
with amp_context:
|
||||
for chunk in batch.learn_chunks:
|
||||
memory = model(chunk, memory=memory, update_memory=True).memory
|
||||
output = model(batch.query_input, memory=memory, update_memory=False, labels=batch.query_labels)
|
||||
if output.loss is None:
|
||||
raise RuntimeError("training loss was not produced")
|
||||
loss = output.loss
|
||||
loss.backward()
|
||||
clip_grad_norm_(model.parameters(), 1.0)
|
||||
optimizer.step()
|
||||
|
||||
if step == 1 or step % args.eval_every == 0 or step == args.steps:
|
||||
eval_loss, accuracy = evaluate(model, device=device, overwrite=args.overwrite)
|
||||
print(f"step={step:5d} train_loss={loss.detach().item():.4f} eval_loss={eval_loss:.4f} accuracy={accuracy:.3f}")
|
||||
state = {
|
||||
"config": vars(config),
|
||||
"model": model.state_dict(),
|
||||
"optimizer": optimizer.state_dict(),
|
||||
"step": step,
|
||||
"seed": args.seed,
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
torch.save(state, checkpoint_dir / "latest.pt")
|
||||
(checkpoint_dir / "run.json").write_text(json.dumps({"args": vars(args), "device": str(device)}, indent=2), encoding="utf-8")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,284 @@
|
||||
"""Train the internal high-recall automatic memory write policy.
|
||||
|
||||
The frozen Qwen backbone and the existing native memory controller provide
|
||||
the representation. Only a small policy head is trained. Positive examples
|
||||
cover durable personal facts, preferences, plans, project constraints and
|
||||
corrections; negative examples cover questions, requests, hypotheticals and
|
||||
casual conversation. The runtime still stores the exact user token sequence,
|
||||
so this head decides *whether* to remember rather than compressing the fact.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import random
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .qwen_integration import load_memory_config, load_qwen_dynamic, load_tokenizer
|
||||
|
||||
|
||||
POSITIVE = (
|
||||
"我叫{value}。",
|
||||
"我的常住城市是{value}。",
|
||||
"我最喜欢的水果是{value}。",
|
||||
"我正在开发{value}项目。",
|
||||
"以后请把代码默认写成{value}。",
|
||||
"我计划在{value}完成这个任务。",
|
||||
"我的工作地点是{value}。",
|
||||
"我的常用时区是{value}。",
|
||||
"请记住这条信息：{value}。",
|
||||
"更正一下，刚才的内容应该是{value}。",
|
||||
"这是我的长期偏好：{value}。",
|
||||
"这个项目的重要约束是{value}。",
|
||||
)
|
||||
|
||||
NEGATIVE = (
|
||||
"帮我写一段关于{value}的代码。",
|
||||
"请解释{value}是什么意思。",
|
||||
"{value}是什么？",
|
||||
"你觉得{value}怎么样？",
|
||||
"如果以后遇到{value}，应该怎么办？",
|
||||
"今天天气不错，随便聊聊{value}。",
|
||||
"请把{value}翻译成英文。",
|
||||
"计算一下{value}。",
|
||||
"给我介绍一下{value}。",
|
||||
"哈哈，{value}真有意思。",
|
||||
"我想知道之前有没有提到{value}。",
|
||||
"假设我选择{value}，会发生什么？",
|
||||
"我叫什么？",
|
||||
"我的名字是什么？",
|
||||
"你还记得我叫什么吗？",
|
||||
"我正在开发什么项目？",
|
||||
"我的项目叫什么？",
|
||||
"请问我的项目叫什么？",
|
||||
"我的工作地点是什么？",
|
||||
"我的工作地点代号是什么？",
|
||||
"工作地点代号是多少？",
|
||||
"你记得我的工作地点吗？",
|
||||
"请告诉我之前有没有说过{value}。",
|
||||
"我之前有没有告诉过你{value}？",
|
||||
"能不能帮我完成{value}？",
|
||||
"如何处理{value}？",
|
||||
"请给我一个{value}的方案。",
|
||||
)
|
||||
|
||||
VALUES = (
|
||||
"小明",
|
||||
"上海",
|
||||
"红富士苹果",
|
||||
"个人记忆系统",
|
||||
"Python",
|
||||
"下周五",
|
||||
"R7",
|
||||
"Asia/Shanghai",
|
||||
"不要删除用户数据",
|
||||
"使用简洁中文",
|
||||
"每天晚上八点",
|
||||
"蓝鲸-47",
|
||||
)
|
||||
|
||||
|
||||
def make_examples(seed: int, count: int) -> list[tuple[str, float]]:
|
||||
rng = random.Random(seed)
|
||||
examples: list[tuple[str, float]] = []
|
||||
for _ in range(count):
|
||||
examples.append((rng.choice(POSITIVE).format(value=rng.choice(VALUES)), 1.0))
|
||||
examples.append((rng.choice(NEGATIVE).format(value=rng.choice(VALUES)), 0.0))
|
||||
rng.shuffle(examples)
|
||||
return examples
|
||||
|
||||
|
||||
def encode_batch(tokenizer, texts: list[str], device: torch.device):
|
||||
rows = []
|
||||
masks = []
|
||||
for text in texts:
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
[{"role": "user", "content": text}],
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
return_tensors="pt",
|
||||
return_dict=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
rows.append(encoded["input_ids"][0])
|
||||
masks.append(encoded.get("attention_mask", torch.ones_like(encoded["input_ids"]))[0])
|
||||
max_length = max(row.numel() for row in rows)
|
||||
input_ids = torch.zeros(len(rows), max_length, dtype=torch.long, device=device)
|
||||
attention_mask = torch.zeros(len(rows), max_length, dtype=torch.long, device=device)
|
||||
for index, (row, mask) in enumerate(zip(rows, masks)):
|
||||
input_ids[index, : row.numel()] = row.to(device)
|
||||
attention_mask[index, : mask.numel()] = mask.to(device)
|
||||
return input_ids, attention_mask
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def collect_representations(model, tokenizer, examples, *, batch_size: int, device):
|
||||
vectors = []
|
||||
labels = []
|
||||
for start in range(0, len(examples), batch_size):
|
||||
batch = examples[start : start + batch_size]
|
||||
input_ids, attention_mask = encode_batch(
|
||||
tokenizer,
|
||||
[item[0] for item in batch],
|
||||
device,
|
||||
)
|
||||
model.reset_memory(batch_size=len(batch), device=device)
|
||||
model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
read_memory=False,
|
||||
update_memory=True,
|
||||
return_memory=True,
|
||||
use_cache=False,
|
||||
)
|
||||
representation = getattr(model.memory, "last_write_representation", None)
|
||||
if representation is None:
|
||||
raise RuntimeError("native controller did not expose a write representation")
|
||||
vectors.append(representation.detach().float().cpu())
|
||||
labels.extend(item[1] for item in batch)
|
||||
return torch.cat(vectors), torch.tensor(labels, dtype=torch.float32)
|
||||
|
||||
|
||||
def evaluate(policy, vectors, labels, threshold: float) -> dict[str, float]:
|
||||
with torch.inference_mode():
|
||||
probabilities = torch.sigmoid(policy(vectors)).cpu()
|
||||
labels = labels.cpu()
|
||||
predictions = probabilities >= threshold
|
||||
positive = labels >= 0.5
|
||||
negative = ~positive
|
||||
true_positive = (predictions & positive).sum().item()
|
||||
false_negative = ((~predictions) & positive).sum().item()
|
||||
false_positive = (predictions & negative).sum().item()
|
||||
true_negative = ((~predictions) & negative).sum().item()
|
||||
return {
|
||||
"threshold": threshold,
|
||||
"accuracy": float((predictions == positive).float().mean()),
|
||||
"positive_recall": true_positive / max(1, true_positive + false_negative),
|
||||
"negative_specificity": true_negative / max(1, true_negative + false_positive),
|
||||
"false_positive_rate": false_positive / max(1, false_positive + true_negative),
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-path", default=".")
|
||||
parser.add_argument(
|
||||
"--base-adapter",
|
||||
default="dynamic_memory_lab/qwen_memory_adapter_natural_controller_v3",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-adapter",
|
||||
default="dynamic_memory_lab/qwen_memory_adapter_natural_auto_v3",
|
||||
)
|
||||
parser.add_argument("--steps", type=int, default=2400)
|
||||
parser.add_argument("--example-count", type=int, default=1280)
|
||||
parser.add_argument("--batch-size", type=int, default=32)
|
||||
parser.add_argument("--lr", type=float, default=2e-4)
|
||||
parser.add_argument("--threshold", type=float, default=0.35)
|
||||
parser.add_argument(
|
||||
"--text-memory-threshold",
|
||||
type=float,
|
||||
default=0.30,
|
||||
help="retrieval threshold used by the automatic-memory adapter",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=20260904)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
random.seed(args.seed)
|
||||
torch.manual_seed(args.seed)
|
||||
tokenizer = load_tokenizer(args.model_path)
|
||||
config = load_memory_config(args.base_adapter)
|
||||
config.natural_language_memory = True
|
||||
config.automatic_memory = True
|
||||
config.auto_memory_threshold = args.threshold
|
||||
config.text_memory_threshold = args.text_memory_threshold
|
||||
config.persistent_memory = False
|
||||
model = load_qwen_dynamic(
|
||||
args.model_path,
|
||||
memory_config=config,
|
||||
load_in_4bit=not args.no_4bit,
|
||||
)
|
||||
model.load_memory_adapter(args.base_adapter, strict=True)
|
||||
if model.memory_policy is None:
|
||||
raise RuntimeError("automatic memory policy was not created")
|
||||
model.eval()
|
||||
model.memory_policy.train()
|
||||
device = model._find_layer_device()
|
||||
|
||||
examples = make_examples(args.seed, args.example_count)
|
||||
split = int(len(examples) * 0.8)
|
||||
train_examples = examples[:split]
|
||||
eval_examples = examples[split:]
|
||||
train_vectors, train_labels = collect_representations(
|
||||
model,
|
||||
tokenizer,
|
||||
train_examples,
|
||||
batch_size=args.batch_size,
|
||||
device=device,
|
||||
)
|
||||
eval_vectors, eval_labels = collect_representations(
|
||||
model,
|
||||
tokenizer,
|
||||
eval_examples,
|
||||
batch_size=args.batch_size,
|
||||
device=device,
|
||||
)
|
||||
train_vectors = train_vectors.to(device)
|
||||
train_labels = train_labels.to(device)
|
||||
eval_vectors = eval_vectors.to(device)
|
||||
eval_labels = eval_labels.to(device)
|
||||
|
||||
optimizer = torch.optim.AdamW(model.memory_policy.parameters(), lr=args.lr, weight_decay=0.01)
|
||||
positive_weight = torch.tensor([1.5], device=device)
|
||||
rng = random.Random(args.seed + 1)
|
||||
for step in range(1, args.steps + 1):
|
||||
indices = torch.tensor(
|
||||
[rng.randrange(train_vectors.shape[0]) for _ in range(args.batch_size)],
|
||||
dtype=torch.long,
|
||||
device=device,
|
||||
)
|
||||
logits = model.memory_policy(train_vectors[indices])
|
||||
weights = torch.where(train_labels[indices] >= 0.5, positive_weight, torch.ones_like(logits))
|
||||
loss = F.binary_cross_entropy_with_logits(logits, train_labels[indices], weight=weights)
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(model.memory_policy.parameters(), 1.0)
|
||||
optimizer.step()
|
||||
if step == 1 or step % 100 == 0 or step == args.steps:
|
||||
stats = evaluate(model.memory_policy, train_vectors, train_labels, args.threshold)
|
||||
print(
|
||||
f"step={step} loss={float(loss.detach()):.5f} "
|
||||
f"train_accuracy={stats['accuracy']:.3f} "
|
||||
f"positive_recall={stats['positive_recall']:.3f} "
|
||||
f"false_positive_rate={stats['false_positive_rate']:.3f}"
|
||||
)
|
||||
|
||||
model.memory_policy.eval()
|
||||
model._memory_policy_ready = True
|
||||
output_dir = Path(args.output_adapter)
|
||||
model.save_memory_adapter(output_dir)
|
||||
report = {
|
||||
"steps": args.steps,
|
||||
"example_count_per_class": args.example_count,
|
||||
"train_examples": len(train_examples),
|
||||
"eval_examples": len(eval_examples),
|
||||
"source_adapter": str(args.base_adapter),
|
||||
"policy": "high_recall_automatic_memory_importance",
|
||||
"threshold": args.threshold,
|
||||
"train": evaluate(model.memory_policy, train_vectors, train_labels, args.threshold),
|
||||
"eval": evaluate(model.memory_policy, eval_vectors, eval_labels, args.threshold),
|
||||
}
|
||||
(output_dir / "auto_policy_training.json").write_text(
|
||||
json.dumps(report, ensure_ascii=False, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
print(json.dumps(report, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,245 @@
|
||||
"""Train the Natural Memory v2 sparse router on hard-negative episodes.
|
||||
|
||||
The default dataset is generated from a shared latent factor space. Train
|
||||
and validation topics are sampled independently from that same space, so the
|
||||
validation numbers measure generalization rather than memorization of topic
|
||||
ids. A JSONL collector can be added later without changing the router loss;
|
||||
the important unit is still a query, candidate addresses, a positive index,
|
||||
and a write/read policy label.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import random
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
if __package__ in {None, ""}:
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
from dynamic_memory_lab.memory_os_v2 import MemoryRouterV2
|
||||
|
||||
|
||||
def _device(name: str) -> torch.device:
|
||||
if name == "auto":
|
||||
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
return torch.device(name)
|
||||
|
||||
|
||||
def _make_basis(hidden_size: int, latent_size: int, device: torch.device) -> torch.Tensor:
|
||||
basis = torch.randn(hidden_size, latent_size, device=device)
|
||||
return F.normalize(basis, dim=0)
|
||||
|
||||
|
||||
def _latent_to_hidden(latent: torch.Tensor, basis: torch.Tensor, noise: float) -> torch.Tensor:
|
||||
hidden = latent @ basis.T
|
||||
if noise > 0:
|
||||
hidden = hidden + noise * torch.randn_like(hidden)
|
||||
return hidden
|
||||
|
||||
|
||||
def sample_episode(
|
||||
*,
|
||||
batch_size: int,
|
||||
candidate_count: int,
|
||||
basis: torch.Tensor,
|
||||
device: torch.device,
|
||||
no_memory_rate: float = 0.20,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Create semantic positives, hard negatives and no-memory episodes."""
|
||||
|
||||
latent_size = basis.shape[1]
|
||||
query_latent = torch.randn(batch_size, latent_size, device=device)
|
||||
no_memory = torch.rand(batch_size, device=device) < no_memory_rate
|
||||
# A no-memory query is deliberately drawn from a low-information region
|
||||
# instead of being another random topic. This gives the policy head a
|
||||
# learnable abstention signal while still leaving unrelated hard negatives
|
||||
# in the candidate set.
|
||||
query_latent[no_memory] = 0.0
|
||||
query = _latent_to_hidden(query_latent, basis, 0.04)
|
||||
candidates_latent = torch.randn(batch_size, candidate_count, latent_size, device=device)
|
||||
candidates_latent[:, 0] = query_latent + 0.04 * torch.randn_like(query_latent)
|
||||
if candidate_count > 1:
|
||||
# A hard negative shares most factors but changes two dimensions.
|
||||
hard = query_latent + 0.16 * torch.randn_like(query_latent)
|
||||
hard[:, :2] = -hard[:, :2]
|
||||
candidates_latent[:, 1] = hard
|
||||
candidates_latent[no_memory] = torch.randn(
|
||||
int(no_memory.sum().item()), candidate_count, latent_size, device=device
|
||||
)
|
||||
candidates = _latent_to_hidden(candidates_latent, basis, 0.04)
|
||||
need_memory = (~no_memory).float()
|
||||
hop_label = torch.where(
|
||||
query_latent.norm(dim=-1) > math.sqrt(latent_size),
|
||||
torch.full((batch_size,), 2, device=device, dtype=torch.long),
|
||||
torch.where(
|
||||
query_latent[:, 0] > 0,
|
||||
torch.ones(batch_size, device=device, dtype=torch.long),
|
||||
torch.zeros(batch_size, device=device, dtype=torch.long),
|
||||
),
|
||||
)
|
||||
positive_index = torch.zeros(batch_size, device=device, dtype=torch.long)
|
||||
return query, candidates, positive_index, need_memory, hop_label
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate(
|
||||
router: MemoryRouterV2,
|
||||
*,
|
||||
basis: torch.Tensor,
|
||||
device: torch.device,
|
||||
batches: int,
|
||||
batch_size: int,
|
||||
candidate_count: int,
|
||||
) -> dict[str, float]:
|
||||
router.eval()
|
||||
route_correct = 0
|
||||
route_total = 0
|
||||
need_tp = need_tn = need_fp = need_fn = 0
|
||||
hop_correct = 0
|
||||
hop_total = 0
|
||||
for _ in range(batches):
|
||||
query, candidates, positive, need, hop = sample_episode(
|
||||
batch_size=batch_size,
|
||||
candidate_count=candidate_count,
|
||||
basis=basis,
|
||||
device=device,
|
||||
no_memory_rate=0.25,
|
||||
)
|
||||
output = router(query, candidates)
|
||||
predicted = output["scores"].argmax(dim=-1)
|
||||
route_correct += int(((predicted == positive) & need.bool()).sum().item())
|
||||
route_total += int(need.sum().item())
|
||||
need_pred = (torch.sigmoid(output["need_memory_logits"]) >= 0.5).float()
|
||||
need_tp += int(((need_pred == 1) & (need == 1)).sum().item())
|
||||
need_tn += int(((need_pred == 0) & (need == 0)).sum().item())
|
||||
need_fp += int(((need_pred == 1) & (need == 0)).sum().item())
|
||||
need_fn += int(((need_pred == 0) & (need == 1)).sum().item())
|
||||
hop_correct += int((output["hop_logits"].argmax(dim=-1) == hop).sum().item())
|
||||
hop_total += batch_size
|
||||
precision = need_tp / max(1, need_tp + need_fp)
|
||||
recall = need_tp / max(1, need_tp + need_fn)
|
||||
return {
|
||||
"route_accuracy": route_correct / max(1, route_total),
|
||||
"need_memory_precision": precision,
|
||||
"need_memory_recall": recall,
|
||||
"need_memory_specificity": need_tn / max(1, need_tn + need_fp),
|
||||
"hop_accuracy": hop_correct / max(1, hop_total),
|
||||
"need_tp": float(need_tp),
|
||||
"need_tn": float(need_tn),
|
||||
"need_fp": float(need_fp),
|
||||
"need_fn": float(need_fn),
|
||||
}
|
||||
|
||||
|
||||
def train(args: argparse.Namespace) -> dict[str, Any]:
|
||||
random.seed(args.seed)
|
||||
torch.manual_seed(args.seed)
|
||||
device = _device(args.device)
|
||||
if device.type == "cuda":
|
||||
torch.cuda.manual_seed_all(args.seed)
|
||||
torch.set_float32_matmul_precision("high")
|
||||
basis = _make_basis(args.hidden_size, args.latent_size, device)
|
||||
router = MemoryRouterV2(
|
||||
args.hidden_size,
|
||||
router_dim=args.router_dim,
|
||||
num_heads=args.num_heads,
|
||||
max_hops=args.max_hops,
|
||||
).to(device)
|
||||
optimizer = torch.optim.AdamW(router.parameters(), lr=args.learning_rate, weight_decay=1e-4)
|
||||
history: list[dict[str, float]] = []
|
||||
router.train()
|
||||
for step in range(1, args.steps + 1):
|
||||
query, candidates, positive, need, hop = sample_episode(
|
||||
batch_size=args.batch_size,
|
||||
candidate_count=args.candidate_count,
|
||||
basis=basis,
|
||||
device=device,
|
||||
)
|
||||
output = router(query, candidates)
|
||||
route_mask = need.bool()
|
||||
candidate_loss = (
|
||||
F.cross_entropy(output["scores"][route_mask], positive[route_mask])
|
||||
if bool(route_mask.any())
|
||||
else output["scores"].sum() * 0.0
|
||||
)
|
||||
need_loss = F.binary_cross_entropy_with_logits(output["need_memory_logits"], need)
|
||||
hop_loss = F.cross_entropy(output["hop_logits"], hop)
|
||||
loss = candidate_loss + args.need_loss_weight * need_loss + args.hop_loss_weight * hop_loss
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(router.parameters(), 1.0)
|
||||
optimizer.step()
|
||||
if step == 1 or step % args.log_every == 0 or step == args.steps:
|
||||
history.append(
|
||||
{
|
||||
"step": float(step),
|
||||
"loss": float(loss.detach().cpu()),
|
||||
"candidate_loss": float(candidate_loss.detach().cpu()),
|
||||
"need_loss": float(need_loss.detach().cpu()),
|
||||
"hop_loss": float(hop_loss.detach().cpu()),
|
||||
}
|
||||
)
|
||||
validation = evaluate(
|
||||
router,
|
||||
basis=basis,
|
||||
device=device,
|
||||
batches=args.eval_batches,
|
||||
batch_size=args.batch_size,
|
||||
candidate_count=args.candidate_count,
|
||||
)
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
torch.save(router.state_dict(), output_dir / "memory_router_v2.pt")
|
||||
torch.save(basis.detach().cpu(), output_dir / "memory_router_v2_basis.pt")
|
||||
summary = {
|
||||
"format_version": 2,
|
||||
"seed": args.seed,
|
||||
"device": str(device),
|
||||
"hidden_size": args.hidden_size,
|
||||
"router_dim": args.router_dim,
|
||||
"num_heads": args.num_heads,
|
||||
"max_hops": args.max_hops,
|
||||
"steps": args.steps,
|
||||
"batch_size": args.batch_size,
|
||||
"candidate_count": args.candidate_count,
|
||||
"training_history": history,
|
||||
"validation": validation,
|
||||
}
|
||||
(output_dir / "memory_router_v2_training.json").write_text(
|
||||
json.dumps(summary, indent=2, ensure_ascii=False), encoding="utf-8"
|
||||
)
|
||||
return summary
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--output-dir", default="W:/Flash/model/dynamic_memory_lab/checkpoints/natural_memory_v2_router")
|
||||
parser.add_argument("--device", default="auto")
|
||||
parser.add_argument("--hidden-size", type=int, default=2560)
|
||||
parser.add_argument("--router-dim", type=int, default=128)
|
||||
parser.add_argument("--num-heads", type=int, default=8)
|
||||
parser.add_argument("--max-hops", type=int, default=3)
|
||||
parser.add_argument("--latent-size", type=int, default=32)
|
||||
parser.add_argument("--steps", type=int, default=1200)
|
||||
parser.add_argument("--batch-size", type=int, default=64)
|
||||
parser.add_argument("--candidate-count", type=int, default=32)
|
||||
parser.add_argument("--eval-batches", type=int, default=80)
|
||||
parser.add_argument("--learning-rate", type=float, default=2e-3)
|
||||
parser.add_argument("--need-loss-weight", type=float, default=0.75)
|
||||
parser.add_argument("--hop-loss-weight", type=float, default=0.35)
|
||||
parser.add_argument("--log-every", type=int, default=100)
|
||||
parser.add_argument("--seed", type=int, default=20260904)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
result = train(parse_args())
|
||||
print(json.dumps(result["validation"], ensure_ascii=False, indent=2))
|
||||
@@ -0,0 +1,222 @@
|
||||
"""Train the learned persistent-memory controller with policy auxiliary losses."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import random
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.utils import clip_grad_norm_
|
||||
|
||||
from .qwen_integration import QwenMemoryConfig, load_qwen_dynamic, load_tokenizer
|
||||
from .train_qwen_memory import encode_messages, pad_batch
|
||||
|
||||
|
||||
def load_records(path: str | Path) -> list[dict[str, Any]]:
|
||||
records = []
|
||||
for line in Path(path).read_text(encoding="utf-8").splitlines():
|
||||
if line.strip():
|
||||
records.append(json.loads(line))
|
||||
if not records:
|
||||
raise ValueError(f"no records found in {path}")
|
||||
return records
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-path", default=".")
|
||||
parser.add_argument("--data", default="dynamic_memory_lab/data/native_memory/train.jsonl")
|
||||
parser.add_argument("--output-dir", default="dynamic_memory_lab/qwen_memory_adapter_native")
|
||||
parser.add_argument("--steps", type=int, default=1000)
|
||||
parser.add_argument("--lr", type=float, default=1e-4)
|
||||
parser.add_argument("--max-length", type=int, default=192)
|
||||
parser.add_argument("--seed", type=int, default=20260904)
|
||||
parser.add_argument("--save-every", type=int, default=50)
|
||||
parser.add_argument("--direct-logit-scale", type=float, default=4.0)
|
||||
parser.add_argument("--write-loss-weight", type=float, default=0.25)
|
||||
parser.add_argument("--forget-loss-weight", type=float, default=0.25)
|
||||
parser.add_argument(
|
||||
"--value-loss-weight",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="weight for aligning each labeled fact's write representation to its value token",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--value-cosine-weight",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="additional cosine alignment weight against the frozen output embedding row",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--forget-positive-weight",
|
||||
type=float,
|
||||
default=4.0,
|
||||
help="extra BCE weight for positive replacement/forget examples",
|
||||
)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
random.seed(args.seed)
|
||||
torch.manual_seed(args.seed)
|
||||
records = load_records(args.data)
|
||||
tokenizer = load_tokenizer(args.model_path)
|
||||
config = QwenMemoryConfig(
|
||||
mode="blend",
|
||||
blend_init=0.1,
|
||||
direct_logit_scale=args.direct_logit_scale,
|
||||
native_mode=True,
|
||||
persistent_memory=False,
|
||||
summary_pooling=True,
|
||||
)
|
||||
model = load_qwen_dynamic(
|
||||
args.model_path,
|
||||
memory_config=config,
|
||||
load_in_4bit=not args.no_4bit,
|
||||
)
|
||||
model.train()
|
||||
parameters = list(model.trainable_parameters)
|
||||
optimizer = torch.optim.AdamW(parameters, lr=args.lr, weight_decay=0.01)
|
||||
device = model._find_layer_device()
|
||||
pad_id = int(tokenizer.pad_token_id)
|
||||
output_dir = Path(args.output_dir)
|
||||
|
||||
print(
|
||||
f"device={device} records={len(records)} layers={model.layer_indices} "
|
||||
f"direct_logit_scale={args.direct_logit_scale}"
|
||||
)
|
||||
for step in range(1, args.steps + 1):
|
||||
record = records[(step - 1) % len(records)]
|
||||
model.reset_memory()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
write_losses: list[torch.Tensor] = []
|
||||
forget_losses: list[torch.Tensor] = []
|
||||
value_losses: list[torch.Tensor] = []
|
||||
value_cosine_losses: list[torch.Tensor] = []
|
||||
chunks = record.get("memory_chunks")
|
||||
if not isinstance(chunks, list) or not chunks:
|
||||
raise ValueError("each native-memory record needs a non-empty memory_chunks list")
|
||||
|
||||
for chunk in chunks:
|
||||
memory_item = encode_messages(tokenizer, chunk["messages"], args.max_length)
|
||||
memory_input, memory_mask, _ = pad_batch([memory_item], pad_id)
|
||||
memory_output = 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,
|
||||
)
|
||||
del memory_output
|
||||
write_probability = model.memory.last_write_probability
|
||||
forget_probability = model.memory.last_forget_probability
|
||||
if write_probability is None or forget_probability is None:
|
||||
raise RuntimeError("native memory controller did not expose write statistics")
|
||||
write_target = torch.full_like(write_probability, float(chunk.get("write_label", 1.0)))
|
||||
forget_target = torch.full_like(
|
||||
forget_probability,
|
||||
float(chunk.get("forget_label", 0.0)),
|
||||
)
|
||||
write_losses.append(F.binary_cross_entropy(write_probability, write_target))
|
||||
forget_weight = 1.0 + (args.forget_positive_weight - 1.0) * forget_target
|
||||
forget_losses.append(
|
||||
F.binary_cross_entropy(forget_probability, forget_target, weight=forget_weight)
|
||||
)
|
||||
value = chunk.get("value")
|
||||
write_representation = model.memory.last_write_representation
|
||||
if value and write_representation is not None:
|
||||
value_tokens = tokenizer(
|
||||
str(value),
|
||||
add_special_tokens=False,
|
||||
)["input_ids"]
|
||||
if value_tokens and isinstance(value_tokens[0], list):
|
||||
value_tokens = value_tokens[0]
|
||||
if value_tokens:
|
||||
target_id = torch.tensor(
|
||||
[int(value_tokens[0])],
|
||||
dtype=torch.long,
|
||||
device=model.base_model.get_output_embeddings().weight.device,
|
||||
)
|
||||
output_embeddings = model.base_model.get_output_embeddings()
|
||||
write_logits = output_embeddings(
|
||||
write_representation.to(
|
||||
device=output_embeddings.weight.device,
|
||||
dtype=output_embeddings.weight.dtype,
|
||||
)
|
||||
).float()
|
||||
value_losses.append(F.cross_entropy(write_logits, target_id))
|
||||
target_embedding = output_embeddings.weight[target_id].detach().float()
|
||||
predicted_embedding = write_representation.float()
|
||||
value_cosine_losses.append(
|
||||
1.0
|
||||
- F.cosine_similarity(predicted_embedding, target_embedding, dim=-1).mean()
|
||||
)
|
||||
|
||||
query_item = encode_messages(tokenizer, record["query"], args.max_length)
|
||||
query_input, query_mask, query_labels = pad_batch([query_item], pad_id)
|
||||
query_output = model(
|
||||
input_ids=query_input.to(device),
|
||||
attention_mask=query_mask.to(device),
|
||||
labels=query_labels.to(device),
|
||||
read_memory=True,
|
||||
update_memory=False,
|
||||
return_memory=True,
|
||||
use_cache=False,
|
||||
)
|
||||
if query_output.loss is None:
|
||||
raise RuntimeError("native Qwen query returned no loss")
|
||||
write_loss = torch.stack(write_losses).mean()
|
||||
forget_loss = torch.stack(forget_losses).mean()
|
||||
value_loss = torch.stack(value_losses).mean() if value_losses else query_output.loss.new_zeros(())
|
||||
value_cosine_loss = (
|
||||
torch.stack(value_cosine_losses).mean()
|
||||
if value_cosine_losses
|
||||
else query_output.loss.new_zeros(())
|
||||
)
|
||||
loss = (
|
||||
query_output.loss
|
||||
+ args.write_loss_weight * write_loss
|
||||
+ args.forget_loss_weight * forget_loss
|
||||
+ args.value_loss_weight * value_loss
|
||||
+ args.value_cosine_weight * value_cosine_loss
|
||||
)
|
||||
loss.backward()
|
||||
clip_grad_norm_(parameters, 1.0)
|
||||
optimizer.step()
|
||||
|
||||
if step == 1 or step % 10 == 0 or step == args.steps:
|
||||
write_mean = float(torch.stack(write_losses).detach().mean())
|
||||
forget_mean = float(torch.stack(forget_losses).detach().mean())
|
||||
print(
|
||||
f"step={step:4d} loss={loss.detach().item():.4f} "
|
||||
f"query={query_output.loss.detach().item():.4f} "
|
||||
f"write_bce={write_mean:.4f} forget_bce={forget_mean:.4f} "
|
||||
f"value={value_loss.detach().item():.4f} "
|
||||
f"value_cos={value_cosine_loss.detach().item():.4f}"
|
||||
)
|
||||
if step % args.save_every == 0 or step == args.steps:
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
model.save_memory_adapter(output_dir)
|
||||
(output_dir / "training_state.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"step": step,
|
||||
"data": str(args.data),
|
||||
"model_path": str(args.model_path),
|
||||
"controller": "native_learned_write_forget_summary",
|
||||
},
|
||||
indent=2,
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,747 @@
|
||||
"""Train the small internal query-to-text-memory retrieval head."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import random
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .qwen_integration import load_memory_config, load_qwen_dynamic, load_tokenizer
|
||||
|
||||
|
||||
ATTRIBUTES = {
|
||||
"name": {
|
||||
"facts": (
|
||||
"我的名字是{value}。",
|
||||
"请记住，我叫{value}。",
|
||||
"用户姓名记录为{value}。",
|
||||
"我叫{value}",
|
||||
"我是{value}",
|
||||
"叫我{value}就行",
|
||||
"姓名：{value}",
|
||||
"用户叫{value}",
|
||||
),
|
||||
"queries": (
|
||||
"我叫什么名字？",
|
||||
"请告诉我已经记录的姓名。",
|
||||
"你还记得我的名字吗？",
|
||||
"我叫什么？",
|
||||
"我的姓名是什么？",
|
||||
"怎么称呼我？",
|
||||
"我的名字呢？",
|
||||
"你记得我叫什么吗？",
|
||||
"我在这里登记的名字是什么？",
|
||||
"我在这里叫什么",
|
||||
"你这里记录的我叫什么",
|
||||
"我在你这边叫什么",
|
||||
"你这边怎么称呼我",
|
||||
"系统里记录的我的名字是什么",
|
||||
),
|
||||
"holdout_queries": (
|
||||
"我是谁？",
|
||||
"你知道我是谁吗？",
|
||||
"你还记得我是谁吗？",
|
||||
"我在你这里叫什么？",
|
||||
),
|
||||
},
|
||||
"project": {
|
||||
"facts": (
|
||||
"我正在开发{value}项目。",
|
||||
"请记住，我当前负责的项目是{value}。",
|
||||
"我的当前项目名称是{value}。",
|
||||
"我做的是{value}",
|
||||
"项目：{value}",
|
||||
"最近在搞{value}",
|
||||
"我手头做{value}",
|
||||
),
|
||||
"queries": (
|
||||
"我正在开发什么项目？",
|
||||
"请查询我当前的项目。",
|
||||
"我之前说过正在做什么吗？",
|
||||
"我最近在做什么？",
|
||||
"我手头的项目叫什么？",
|
||||
"我在忙哪个项目？",
|
||||
"我最近搞的是什么？",
|
||||
"手上的活是什么项目？",
|
||||
),
|
||||
"holdout_queries": (
|
||||
"我现在主要做什么？",
|
||||
"我最近在开发哪一个东西？",
|
||||
"之前提到的项目是什么？",
|
||||
),
|
||||
},
|
||||
"plan": {
|
||||
"facts": (
|
||||
"我计划在{value}完成这件事。",
|
||||
"请记住我的计划：{value}。",
|
||||
"我的下一步安排是{value}。",
|
||||
"计划：{value}",
|
||||
"我打算{value}",
|
||||
"准备在{value}完成",
|
||||
"安排：{value}",
|
||||
"待办：{value}",
|
||||
),
|
||||
"queries": (
|
||||
"我接下来的计划是什么？",
|
||||
"请查询我记录过的安排。",
|
||||
"我之前说过下一步要做什么？",
|
||||
"我下一步准备做什么？",
|
||||
"我安排在什么时候完成？",
|
||||
"我接下来怎么安排？",
|
||||
"我的安排是什么？",
|
||||
"我有什么计划？",
|
||||
),
|
||||
"holdout_queries": (
|
||||
"我接下来打算怎么安排？",
|
||||
"我已经计划好的事情是什么？",
|
||||
),
|
||||
},
|
||||
"constraint": {
|
||||
"facts": (
|
||||
"这个项目的重要约束是{value}。",
|
||||
"请记住这个开发约束：{value}。",
|
||||
"以后处理这个项目时必须遵守：{value}。",
|
||||
"要求：{value}",
|
||||
"必须{value}",
|
||||
"别忘了：{value}",
|
||||
),
|
||||
"queries": (
|
||||
"这个项目的重要约束是什么？",
|
||||
"请查询我记录的开发约束。",
|
||||
"之前说过这个项目需要遵守什么吗？",
|
||||
"这个项目有什么限制？",
|
||||
"开发时需要注意哪条规则？",
|
||||
"有哪些要求不能忘？",
|
||||
),
|
||||
"holdout_queries": (
|
||||
"这个项目有哪些不能违反的要求？",
|
||||
"我之前定下的开发规则是什么？",
|
||||
),
|
||||
},
|
||||
"work_code": {
|
||||
"facts": (
|
||||
"我的工作地点代号是{value}。",
|
||||
"请记住，我的工作地点代号为{value}。",
|
||||
"以后如果问到工作地点，请记住代号{value}。",
|
||||
"工作地点：{value}",
|
||||
"地点编号{value}",
|
||||
"我在{value}办公",
|
||||
"办公地点：{value}",
|
||||
"工作地点编号：{value}",
|
||||
),
|
||||
"queries": (
|
||||
"我的工作地点代号是什么？",
|
||||
"请告诉我已经记录的工作地点代号。",
|
||||
"我之前说过的工作地点代号是多少？",
|
||||
"工作地点对应哪个代号？",
|
||||
"我工作的地方编号是什么？",
|
||||
"我办公地点是哪儿？",
|
||||
"我在哪办公？",
|
||||
"我在哪工作？",
|
||||
"工作地点是哪儿？",
|
||||
),
|
||||
"holdout_queries": (
|
||||
"我的办公地点编号是多少？",
|
||||
"我在哪个工作地点？",
|
||||
),
|
||||
},
|
||||
"fruit": {
|
||||
"facts": (
|
||||
"我最喜欢的水果是{value}。",
|
||||
"请记住我的水果偏好：我喜欢{value}。",
|
||||
"我的个人偏好是最喜欢吃{value}。",
|
||||
"我爱吃{value}",
|
||||
"水果偏好：{value}",
|
||||
"我喜欢{value}",
|
||||
),
|
||||
"queries": (
|
||||
"我最喜欢吃什么水果？",
|
||||
"请查询我记录过的水果偏好。",
|
||||
"我之前告诉你的水果喜好是什么？",
|
||||
"我平时爱吃哪种水果？",
|
||||
"我的水果口味偏好是什么？",
|
||||
"我爱吃哪种？",
|
||||
),
|
||||
"holdout_queries": (
|
||||
"我喜欢吃哪一类水果？",
|
||||
"哪种水果是我的首选？",
|
||||
),
|
||||
},
|
||||
"pet": {
|
||||
"facts": (
|
||||
"我养的宠物名字叫{value}。",
|
||||
"请记住，我的宠物是{value}。",
|
||||
"我的宠物信息：名字是{value}。",
|
||||
"宠物是{value}",
|
||||
"我养了{value}",
|
||||
"宠物：{value}",
|
||||
),
|
||||
"queries": (
|
||||
"我养的宠物叫什么名字？",
|
||||
"请查询我的宠物姓名。",
|
||||
"你记得我的宠物是谁吗？",
|
||||
"我家的宠物叫什么？",
|
||||
"我的宠物是哪一只？",
|
||||
"我养的是什么？",
|
||||
),
|
||||
"holdout_queries": (
|
||||
"我养了什么动物？",
|
||||
"我的宠物信息是什么？",
|
||||
),
|
||||
},
|
||||
"editor": {
|
||||
"facts": (
|
||||
"我平时最常用的编辑器是{value}。",
|
||||
"记住我的开发工具偏好：编辑器使用{value}。",
|
||||
"我的编程编辑器偏好为{value}。",
|
||||
"我用{value}写代码",
|
||||
"编辑器：{value}",
|
||||
"开发工具是{value}",
|
||||
"代码用{value}",
|
||||
"工具：{value}",
|
||||
"编程工具：{value}",
|
||||
"写代码用{value}",
|
||||
),
|
||||
"queries": (
|
||||
"我最常用哪个编辑器？",
|
||||
"请查询我的开发工具偏好。",
|
||||
"我的编程编辑器是什么？",
|
||||
"我平时用什么编辑器？",
|
||||
"我习惯用哪款开发工具？",
|
||||
"我写代码用什么？",
|
||||
"我用什么写代码？",
|
||||
"我平时写程序用什么工具？",
|
||||
"我编程时使用什么工具？",
|
||||
"写代码用的是哪款工具？",
|
||||
"我用什么工具编程？",
|
||||
),
|
||||
"holdout_queries": (
|
||||
"我写代码通常使用什么工具？",
|
||||
"我常用的 IDE 是哪个？",
|
||||
),
|
||||
},
|
||||
"city": {
|
||||
"facts": (
|
||||
"我现在长期居住在{value}。",
|
||||
"请记住，我的常住城市是{value}。",
|
||||
"我的个人资料显示常住地为{value}。",
|
||||
"我住在{value}",
|
||||
"常住地：{value}",
|
||||
"我在{value}生活",
|
||||
),
|
||||
"queries": (
|
||||
"我的常住城市是哪里？",
|
||||
"请查询我的居住地。",
|
||||
"我平时住在哪座城市？",
|
||||
"我长期住在哪里？",
|
||||
"我的居住城市是什么？",
|
||||
"我现在住哪儿？",
|
||||
),
|
||||
"holdout_queries": (
|
||||
"我现在定居在哪儿？",
|
||||
"我的长期住址城市是哪座？",
|
||||
),
|
||||
},
|
||||
"timezone": {
|
||||
"facts": (
|
||||
"我的常用时区是{value}。",
|
||||
"请把我的时区偏好记为{value}。",
|
||||
"个人资料：我的时区设置为{value}。",
|
||||
"时区：{value}",
|
||||
"我在{value}时区",
|
||||
"本地时区是{value}",
|
||||
),
|
||||
"queries": (
|
||||
"我的常用时区是什么？",
|
||||
"请查询我的时区设置。",
|
||||
"我使用哪个时区？",
|
||||
"我平时按哪个时区生活？",
|
||||
"我的时间设置是哪一个时区？",
|
||||
"我所在的时区是什么？",
|
||||
),
|
||||
"holdout_queries": (
|
||||
"我的本地时间属于哪个时区？",
|
||||
"我应该使用什么时区？",
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
VALUES = {
|
||||
"name": ("林浩", "小明", "周宁", "陈雪", "Alice", "Zoe", "X7"),
|
||||
"project": ("星火记忆", "自然语言记忆", "Qwen 架构实验", "个人助手"),
|
||||
"plan": ("下周五", "今晚八点", "本周末", "明天上午"),
|
||||
"constraint": ("不要删除用户数据", "使用简洁中文", "保持原版能力", "优先保证可恢复"),
|
||||
"work_code": ("R7", "K9", "蓝鲸-47", "M2"),
|
||||
"fruit": ("红富士苹果", "阳光玫瑰葡萄", "海南芒果", "脆甜梨"),
|
||||
"pet": ("豆包", "团子", "可可", "雪球"),
|
||||
"editor": ("VS Code", "Neovim", "PyCharm", "Emacs"),
|
||||
"city": ("上海", "成都", "深圳", "杭州"),
|
||||
"timezone": ("Asia/Shanghai", "UTC+8", "Europe/London", "America/Los_Angeles"),
|
||||
}
|
||||
|
||||
HARD_NEGATIVE_QUERIES = (
|
||||
"你是谁",
|
||||
"你叫什么",
|
||||
"请介绍你自己",
|
||||
"你能做什么",
|
||||
"你的名字是什么",
|
||||
"今天天气怎么样",
|
||||
"帮我写一段代码",
|
||||
"Python是什么",
|
||||
"解释一下这个概念",
|
||||
)
|
||||
|
||||
|
||||
def _chat_ids(tokenizer, text: str) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
[{"role": "user", "content": text}],
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
return_tensors="pt",
|
||||
return_dict=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
ids = encoded["input_ids"]
|
||||
mask = encoded.get("attention_mask", torch.ones_like(ids))
|
||||
return ids, mask
|
||||
|
||||
|
||||
def _plain_ids(tokenizer, text: str) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
encoded = tokenizer(text, add_special_tokens=False, return_tensors="pt")
|
||||
ids = encoded["input_ids"]
|
||||
mask = encoded.get("attention_mask", torch.ones_like(ids))
|
||||
return ids, mask
|
||||
|
||||
|
||||
def _training_query_candidates(attribute: str) -> list[str]:
|
||||
queries = list(ATTRIBUTES[attribute]["queries"])
|
||||
queries.extend(
|
||||
query.rstrip("？?。！!，,")
|
||||
for query in ATTRIBUTES[attribute]["queries"]
|
||||
if query.rstrip("？?。！!，,")
|
||||
)
|
||||
return list(dict.fromkeys(queries))
|
||||
|
||||
|
||||
def make_pairs(seed: int, count: int) -> list[dict[str, object]]:
|
||||
rng = random.Random(seed)
|
||||
keys = list(ATTRIBUTES)
|
||||
positive_records: list[dict[str, object]] = []
|
||||
|
||||
# Cover every short/long key form and every train-time query form before
|
||||
# falling back to random samples. This prevents the optimizer from
|
||||
# seeing a mostly easy subset of the cross-product.
|
||||
for attribute in keys:
|
||||
for value in VALUES[attribute][:2]:
|
||||
for fact_template in ATTRIBUTES[attribute]["facts"]:
|
||||
for query in _training_query_candidates(attribute):
|
||||
positive_records.append(
|
||||
{
|
||||
"fact": fact_template.format(value=value),
|
||||
"query": query,
|
||||
"label": 1.0,
|
||||
"attribute": attribute,
|
||||
}
|
||||
)
|
||||
|
||||
while len(positive_records) < count:
|
||||
attribute = rng.choice(keys)
|
||||
value = rng.choice(VALUES[attribute])
|
||||
positive_records.append(
|
||||
{
|
||||
"fact": rng.choice(ATTRIBUTES[attribute]["facts"]).format(value=value),
|
||||
"query": rng.choice(_training_query_candidates(attribute)),
|
||||
"label": 1.0,
|
||||
"attribute": attribute,
|
||||
}
|
||||
)
|
||||
rng.shuffle(positive_records)
|
||||
positive_records = positive_records[:count]
|
||||
|
||||
pairs: list[dict[str, object]] = []
|
||||
for positive in positive_records:
|
||||
pairs.append(positive)
|
||||
attribute = str(positive["attribute"])
|
||||
negative_attribute = rng.choice([item for item in keys if item != attribute])
|
||||
negative_query_pool = _training_query_candidates(negative_attribute)
|
||||
if rng.random() < 0.5:
|
||||
negative_query_pool = list(negative_query_pool) + list(HARD_NEGATIVE_QUERIES)
|
||||
negative_query = rng.choice(negative_query_pool)
|
||||
pairs.append(
|
||||
{
|
||||
"fact": positive["fact"],
|
||||
"query": negative_query,
|
||||
"label": 0.0,
|
||||
"attribute": f"{attribute}->{negative_attribute}",
|
||||
}
|
||||
)
|
||||
rng.shuffle(pairs)
|
||||
return pairs
|
||||
|
||||
|
||||
def make_holdout_queries(seed: int) -> list[dict[str, object]]:
|
||||
"""Create query paraphrases that are never used during optimization."""
|
||||
|
||||
rng = random.Random(seed + 17)
|
||||
records: list[dict[str, object]] = []
|
||||
for attribute, definition in ATTRIBUTES.items():
|
||||
value = rng.choice(VALUES[attribute])
|
||||
fact = str(definition["facts"][0]).format(value=value)
|
||||
for query in definition.get("holdout_queries", ()):
|
||||
records.append({"fact": fact, "query": query, "attribute": attribute})
|
||||
rng.shuffle(records)
|
||||
return records
|
||||
|
||||
|
||||
def make_short_fact_holdout() -> list[dict[str, object]]:
|
||||
"""Stress-test colloquial, short and previously unseen fact strings.
|
||||
|
||||
The values are deliberately different from ``VALUES``. This checks that
|
||||
the retriever recognizes the attribute-bearing language around a value,
|
||||
instead of memorizing names, cities or project identifiers.
|
||||
"""
|
||||
|
||||
return [
|
||||
{"fact": "我叫Wpy", "query": "你知道我是谁吗", "attribute": "name"},
|
||||
{"fact": "我叫Wpy", "query": "我是谁", "attribute": "name"},
|
||||
{"fact": "我做的是量子账本", "query": "我最近在忙什么项目", "attribute": "project"},
|
||||
{"fact": "准备在周三完成", "query": "我的待办安排是什么", "attribute": "plan"},
|
||||
{"fact": "必须保留原始数据", "query": "有什么要求必须遵守", "attribute": "constraint"},
|
||||
{"fact": "地点编号Z3", "query": "我在哪儿办公", "attribute": "work_code"},
|
||||
{"fact": "我喜欢白桃", "query": "我爱吃什么", "attribute": "fruit"},
|
||||
{"fact": "我养了阿福", "query": "家里养的是什么动物", "attribute": "pet"},
|
||||
{"fact": "我用Cursor写代码", "query": "我编程时用哪个工具", "attribute": "editor"},
|
||||
{"fact": "我住在苏州", "query": "我人住哪座城", "attribute": "city"},
|
||||
{"fact": "时区：Asia/Tokyo", "query": "本地采用哪个时区", "attribute": "timezone"},
|
||||
]
|
||||
|
||||
|
||||
def make_hard_negative_records() -> list[dict[str, object]]:
|
||||
"""Build explicit non-memory queries for every representative key form."""
|
||||
|
||||
records: list[dict[str, object]] = []
|
||||
for attribute, definition in ATTRIBUTES.items():
|
||||
key_templates = (definition["facts"][0], definition["facts"][3])
|
||||
for value in VALUES[attribute][:2]:
|
||||
for fact_template in key_templates:
|
||||
fact = fact_template.format(value=value)
|
||||
for query in HARD_NEGATIVE_QUERIES:
|
||||
records.append(
|
||||
{
|
||||
"fact": fact,
|
||||
"query": query,
|
||||
"label": 0.0,
|
||||
"attribute": attribute,
|
||||
}
|
||||
)
|
||||
return records
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def encode_records(model, tokenizer, records, *, batch_size: int, device):
|
||||
"""Encode pairs without padding-induced representation drift.
|
||||
|
||||
The Qwen linear-attention path is not perfectly invariant to right-padded
|
||||
batches. Runtime retrieval encodes one query/fact at a time, so training
|
||||
must use the same effective sequence shape. Grouping by exact token
|
||||
length keeps batching efficient while guaranteeing that every row is
|
||||
unpadded and therefore matches single-example inference.
|
||||
"""
|
||||
|
||||
query_rows: list[torch.Tensor] = []
|
||||
key_rows: list[torch.Tensor] = []
|
||||
for record in records:
|
||||
# Runtime retrieval receives the raw user query, not a chat-template
|
||||
# wrapped prompt. Keeping this protocol identical is essential:
|
||||
# otherwise a retriever can score its offline test set well while
|
||||
# failing on the actual restart path.
|
||||
q_ids, _ = _plain_ids(tokenizer, str(record["query"]))
|
||||
k_ids, _ = _plain_ids(tokenizer, str(record["fact"]))
|
||||
query_rows.append(q_ids[0])
|
||||
key_rows.append(k_ids[0])
|
||||
|
||||
def encode_without_padding(rows: list[torch.Tensor]) -> torch.Tensor:
|
||||
vectors: list[Optional[torch.Tensor]] = [None] * len(rows)
|
||||
groups: dict[int, list[int]] = {}
|
||||
for index, row in enumerate(rows):
|
||||
groups.setdefault(int(row.numel()), []).append(index)
|
||||
for indices in groups.values():
|
||||
for start in range(0, len(indices), batch_size):
|
||||
selected = indices[start : start + batch_size]
|
||||
length = rows[selected[0]].numel()
|
||||
batch = torch.stack([rows[index] for index in selected]).to(device)
|
||||
mask = torch.ones((len(selected), length), dtype=torch.long, device=device)
|
||||
encoded = model._encode_model_key(batch, mask).cpu()
|
||||
for row_index, vector in zip(selected, encoded):
|
||||
vectors[row_index] = vector
|
||||
if any(vector is None for vector in vectors):
|
||||
raise RuntimeError("failed to encode every memory record")
|
||||
return torch.stack([vector for vector in vectors if vector is not None])
|
||||
|
||||
return encode_without_padding(query_rows), encode_without_padding(key_rows)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def evaluate_retriever(model, query_vectors, key_vectors, query_attributes, key_attributes) -> dict[str, float]:
|
||||
"""Measure attribute retrieval on held-out paraphrases."""
|
||||
|
||||
scores = model.text_retriever(
|
||||
query_vectors,
|
||||
key_vectors.unsqueeze(0).expand(query_vectors.shape[0], -1, -1),
|
||||
)
|
||||
key_attributes = list(key_attributes)
|
||||
positive_mask = torch.tensor(
|
||||
[[query_attribute == key_attribute for key_attribute in key_attributes]
|
||||
for query_attribute in query_attributes],
|
||||
dtype=torch.bool,
|
||||
device=scores.device,
|
||||
)
|
||||
positive_scores = scores.masked_fill(~positive_mask, torch.finfo(scores.dtype).min).max(dim=1).values
|
||||
negative_scores = scores.masked_fill(positive_mask, torch.finfo(scores.dtype).min).max(dim=1).values
|
||||
positive_probabilities = torch.sigmoid(positive_scores)
|
||||
negative_probabilities = torch.sigmoid(negative_scores)
|
||||
predicted = scores.argmax(dim=1).detach().cpu().tolist()
|
||||
predicted_attributes = [key_attributes[index] for index in predicted]
|
||||
accuracy = sum(
|
||||
predicted_attribute == query_attribute
|
||||
for predicted_attribute, query_attribute in zip(predicted_attributes, query_attributes)
|
||||
) / max(1, len(query_attributes))
|
||||
return {
|
||||
"accuracy": float(accuracy),
|
||||
"positive_score_mean": float(positive_scores.mean().detach().cpu()),
|
||||
"negative_score_mean": float(negative_scores.mean().detach().cpu()),
|
||||
"margin_mean": float((positive_scores - negative_scores).mean().detach().cpu()),
|
||||
"positive_probability_min": float(positive_probabilities.min().detach().cpu()),
|
||||
"negative_probability_max": float(negative_probabilities.max().detach().cpu()),
|
||||
"positive_threshold_recall": float((positive_probabilities >= 0.5).float().mean().detach().cpu()),
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-path", default=".")
|
||||
parser.add_argument("--base-adapter", default="dynamic_memory_lab/qwen_memory_adapter_native_v3")
|
||||
parser.add_argument(
|
||||
"--output-adapter",
|
||||
default="dynamic_memory_lab/qwen_memory_adapter_natural_controller_v1",
|
||||
)
|
||||
parser.add_argument("--steps", type=int, default=2200)
|
||||
parser.add_argument("--pair-count", type=int, default=1600)
|
||||
parser.add_argument("--batch-size", type=int, default=32)
|
||||
parser.add_argument("--lr", type=float, default=2e-4)
|
||||
parser.add_argument("--seed", type=int, default=20260904)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
random.seed(args.seed)
|
||||
torch.manual_seed(args.seed)
|
||||
tokenizer = load_tokenizer(args.model_path)
|
||||
config = load_memory_config(args.base_adapter)
|
||||
config.natural_language_memory = True
|
||||
config.persistent_memory = False
|
||||
config.direct_logit_scale = 0.0
|
||||
model = load_qwen_dynamic(
|
||||
args.model_path,
|
||||
memory_config=config,
|
||||
load_in_4bit=not args.no_4bit,
|
||||
)
|
||||
model.load_memory_adapter(args.base_adapter, strict=True)
|
||||
model.eval()
|
||||
if model.text_retriever is None:
|
||||
raise RuntimeError("natural-language retriever was not created")
|
||||
model.text_retriever.train()
|
||||
device = model._find_layer_device()
|
||||
|
||||
pairs = make_pairs(args.seed, args.pair_count)
|
||||
training_records = [pair for pair in pairs if float(pair["label"]) == 1.0]
|
||||
query_tensor, key_tensor = encode_records(
|
||||
model,
|
||||
tokenizer,
|
||||
training_records,
|
||||
batch_size=args.batch_size,
|
||||
device=device,
|
||||
)
|
||||
query_tensor = query_tensor.to(device)
|
||||
key_tensor = key_tensor.to(device)
|
||||
training_attributes = [str(record["attribute"]) for record in training_records]
|
||||
attribute_names = list(ATTRIBUTES)
|
||||
attribute_indices = {
|
||||
attribute: [index for index, item in enumerate(training_attributes) if item == attribute]
|
||||
for attribute in attribute_names
|
||||
}
|
||||
if any(not indices for indices in attribute_indices.values()):
|
||||
raise RuntimeError("training data did not cover every memory attribute")
|
||||
|
||||
holdout_records = make_holdout_queries(args.seed)
|
||||
holdout_query_tensor, holdout_key_tensor = encode_records(
|
||||
model,
|
||||
tokenizer,
|
||||
holdout_records,
|
||||
batch_size=args.batch_size,
|
||||
device=device,
|
||||
)
|
||||
holdout_query_tensor = holdout_query_tensor.to(device)
|
||||
holdout_key_tensor = holdout_key_tensor.to(device)
|
||||
holdout_attributes = [str(record["attribute"]) for record in holdout_records]
|
||||
short_holdout_records = make_short_fact_holdout()
|
||||
short_holdout_query_tensor, short_holdout_key_tensor = encode_records(
|
||||
model,
|
||||
tokenizer,
|
||||
short_holdout_records,
|
||||
batch_size=args.batch_size,
|
||||
device=device,
|
||||
)
|
||||
short_holdout_query_tensor = short_holdout_query_tensor.to(device)
|
||||
short_holdout_key_tensor = short_holdout_key_tensor.to(device)
|
||||
short_holdout_attributes = [str(record["attribute"]) for record in short_holdout_records]
|
||||
hard_negative_records = make_hard_negative_records()
|
||||
hard_negative_query_tensor, hard_negative_key_tensor = encode_records(
|
||||
model,
|
||||
tokenizer,
|
||||
hard_negative_records,
|
||||
batch_size=args.batch_size,
|
||||
device=device,
|
||||
)
|
||||
hard_negative_query_tensor = hard_negative_query_tensor.to(device)
|
||||
hard_negative_key_tensor = hard_negative_key_tensor.to(device)
|
||||
optimizer = torch.optim.AdamW(model.text_retriever.parameters(), lr=args.lr, weight_decay=0.01)
|
||||
rng = random.Random(args.seed + 1)
|
||||
for step in range(1, args.steps + 1):
|
||||
selected_indices = [
|
||||
rng.choice(attribute_indices[attribute]) for attribute in attribute_names
|
||||
]
|
||||
selected_indices.extend(
|
||||
rng.randrange(query_tensor.shape[0])
|
||||
for _ in range(max(0, args.batch_size - len(selected_indices)))
|
||||
)
|
||||
indices = torch.tensor(
|
||||
selected_indices[: args.batch_size],
|
||||
dtype=torch.long,
|
||||
device=device,
|
||||
)
|
||||
batch_queries = query_tensor[indices]
|
||||
batch_keys = key_tensor[indices]
|
||||
batch_attributes = [training_attributes[index] for index in indices.detach().cpu().tolist()]
|
||||
pair_logits = model.text_retriever(
|
||||
batch_queries,
|
||||
batch_keys.unsqueeze(0).expand(batch_queries.shape[0], -1, -1),
|
||||
)
|
||||
positive_mask = torch.tensor(
|
||||
[[left == right for right in batch_attributes] for left in batch_attributes],
|
||||
dtype=torch.bool,
|
||||
device=device,
|
||||
)
|
||||
positive_logsum = torch.logsumexp(
|
||||
pair_logits.masked_fill(~positive_mask, torch.finfo(pair_logits.dtype).min),
|
||||
dim=1,
|
||||
)
|
||||
contrastive_loss = -(positive_logsum - torch.logsumexp(pair_logits, dim=1)).mean()
|
||||
|
||||
negative_indices = [
|
||||
rng.choice(
|
||||
attribute_indices[
|
||||
rng.choice([name for name in attribute_names if name != attribute])
|
||||
]
|
||||
)
|
||||
for attribute in batch_attributes
|
||||
]
|
||||
negative_keys = key_tensor[torch.tensor(negative_indices, dtype=torch.long, device=device)]
|
||||
positive_logits = pair_logits.diagonal()
|
||||
negative_logits = model.text_retriever(batch_queries, negative_keys)
|
||||
hard_negative_logits = model.text_retriever(
|
||||
hard_negative_query_tensor,
|
||||
hard_negative_key_tensor,
|
||||
)
|
||||
bce_logits = torch.cat((positive_logits, negative_logits), dim=0)
|
||||
bce_labels = torch.cat(
|
||||
(
|
||||
torch.ones_like(positive_logits),
|
||||
torch.zeros_like(negative_logits),
|
||||
),
|
||||
dim=0,
|
||||
)
|
||||
classification_loss = F.binary_cross_entropy_with_logits(bce_logits, bce_labels)
|
||||
hard_negative_loss = F.binary_cross_entropy_with_logits(
|
||||
hard_negative_logits,
|
||||
torch.zeros_like(hard_negative_logits),
|
||||
)
|
||||
loss = contrastive_loss + 0.5 * classification_loss + 0.75 * hard_negative_loss
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(model.text_retriever.parameters(), 1.0)
|
||||
optimizer.step()
|
||||
if step == 1 or step % 100 == 0 or step == args.steps:
|
||||
with torch.inference_mode():
|
||||
predictions = pair_logits.argmax(dim=1)
|
||||
batch_accuracy = sum(
|
||||
batch_attributes[index] == attribute
|
||||
for index, attribute in zip(
|
||||
predictions.detach().cpu().tolist(), batch_attributes
|
||||
)
|
||||
) / max(1, len(batch_attributes))
|
||||
holdout_stats = evaluate_retriever(
|
||||
model,
|
||||
holdout_query_tensor,
|
||||
holdout_key_tensor,
|
||||
holdout_attributes,
|
||||
holdout_attributes,
|
||||
)
|
||||
short_holdout_stats = evaluate_retriever(
|
||||
model,
|
||||
short_holdout_query_tensor,
|
||||
short_holdout_key_tensor,
|
||||
short_holdout_attributes,
|
||||
short_holdout_attributes,
|
||||
)
|
||||
print(
|
||||
f"step={step} loss={float(loss.detach()):.5f} "
|
||||
f"batch_attribute_accuracy={batch_accuracy:.3f} "
|
||||
f"holdout_accuracy={holdout_stats['accuracy']:.3f} "
|
||||
f"holdout_margin={holdout_stats['margin_mean']:.3f} "
|
||||
f"short_fact_accuracy={short_holdout_stats['accuracy']:.3f} "
|
||||
f"short_fact_margin={short_holdout_stats['margin_mean']:.3f} "
|
||||
f"short_fact_threshold_recall={short_holdout_stats['positive_threshold_recall']:.3f}"
|
||||
)
|
||||
|
||||
model.text_retriever.eval()
|
||||
model._text_retriever_ready = True
|
||||
model.memory_config.persistent_memory = False
|
||||
output_dir = Path(args.output_adapter)
|
||||
model.save_memory_adapter(output_dir)
|
||||
stats = {
|
||||
"steps": args.steps,
|
||||
"pair_count": len(pairs),
|
||||
"positive_training_records": len(training_records),
|
||||
"hard_negative_records": len(hard_negative_records),
|
||||
"holdout_records": len(holdout_records),
|
||||
"source_adapter": str(args.base_adapter),
|
||||
"retriever": "qwen_hidden_pair_mlp_multisample_contrastive_with_hard_negatives",
|
||||
"holdout": evaluate_retriever(
|
||||
model,
|
||||
holdout_query_tensor,
|
||||
holdout_key_tensor,
|
||||
holdout_attributes,
|
||||
holdout_attributes,
|
||||
),
|
||||
"short_fact_holdout": evaluate_retriever(
|
||||
model,
|
||||
short_holdout_query_tensor,
|
||||
short_holdout_key_tensor,
|
||||
short_holdout_attributes,
|
||||
short_holdout_attributes,
|
||||
),
|
||||
}
|
||||
(output_dir / "retriever_training.json").write_text(
|
||||
json.dumps(stats, ensure_ascii=False, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
print(json.dumps(stats, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,289 @@
|
||||
"""Train the automatic write/forget policy from normalized conversation JSONL."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import random
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .qwen_integration import load_memory_config, load_qwen_dynamic, load_tokenizer
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
def _project_path(value: str | Path) -> Path:
|
||||
path = Path(value)
|
||||
if path.is_absolute() or path.exists():
|
||||
return path
|
||||
return PROJECT_ROOT / path
|
||||
|
||||
|
||||
def _read_jsonl(path: Path, max_examples: int | None = None) -> list[dict[str, Any]]:
|
||||
rows: list[dict[str, Any]] = []
|
||||
with path.open("r", encoding="utf-8") as handle:
|
||||
for raw in handle:
|
||||
if not raw.strip():
|
||||
continue
|
||||
rows.append(json.loads(raw))
|
||||
if max_examples is not None and len(rows) >= max_examples:
|
||||
break
|
||||
if not rows:
|
||||
raise ValueError(f"no examples found in {path}")
|
||||
return rows
|
||||
|
||||
|
||||
def _encode_batch(tokenizer, texts: list[str], device: torch.device):
|
||||
encoded_rows = []
|
||||
for text in texts:
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
[{"role": "user", "content": text}],
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
return_tensors="pt",
|
||||
return_dict=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
encoded_rows.append(
|
||||
(
|
||||
encoded["input_ids"][0],
|
||||
encoded.get("attention_mask", torch.ones_like(encoded["input_ids"]))[0],
|
||||
)
|
||||
)
|
||||
max_length = max(row.numel() for row, _ in encoded_rows)
|
||||
input_ids = torch.zeros(len(encoded_rows), max_length, dtype=torch.long, device=device)
|
||||
attention_mask = torch.zeros_like(input_ids)
|
||||
for index, (row, mask) in enumerate(encoded_rows):
|
||||
input_ids[index, : row.numel()] = row.to(device)
|
||||
attention_mask[index, : mask.numel()] = mask.to(device)
|
||||
return input_ids, attention_mask
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _collect(model, tokenizer, rows: list[dict[str, Any]], *, batch_size: int, device: torch.device):
|
||||
vectors: list[torch.Tensor] = []
|
||||
write_labels: list[float] = []
|
||||
forget_labels: list[float] = []
|
||||
for start in range(0, len(rows), batch_size):
|
||||
batch = rows[start : start + batch_size]
|
||||
input_ids, attention_mask = _encode_batch(
|
||||
tokenizer,
|
||||
[str(row["text"]) for row in batch],
|
||||
device,
|
||||
)
|
||||
model.reset_memory(batch_size=len(batch), device=device)
|
||||
model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
read_memory=False,
|
||||
update_memory=True,
|
||||
return_memory=True,
|
||||
use_cache=False,
|
||||
)
|
||||
representation = getattr(
|
||||
model.memory,
|
||||
"last_write_summary",
|
||||
getattr(model.memory, "last_write_representation", None),
|
||||
)
|
||||
if representation is None:
|
||||
raise RuntimeError("the loaded memory controller did not expose last_write_representation")
|
||||
vectors.append(representation.detach().float().cpu())
|
||||
write_labels.extend(float(row.get("write_label", 0.0)) for row in batch)
|
||||
forget_labels.extend(float(row.get("forget_label", 0.0)) for row in batch)
|
||||
return (
|
||||
torch.cat(vectors),
|
||||
torch.tensor(write_labels, dtype=torch.float32),
|
||||
torch.tensor(forget_labels, dtype=torch.float32),
|
||||
)
|
||||
|
||||
|
||||
def _metrics(logits: torch.Tensor, labels: torch.Tensor, threshold: float) -> dict[str, float]:
|
||||
probabilities = torch.sigmoid(logits.detach()).reshape(-1).cpu()
|
||||
labels = labels.reshape(-1).cpu() >= 0.5
|
||||
predictions = probabilities >= threshold
|
||||
positive = labels
|
||||
negative = ~labels
|
||||
tp = int((predictions & positive).sum())
|
||||
fn = int((~predictions & positive).sum())
|
||||
fp = int((predictions & negative).sum())
|
||||
tn = int((~predictions & negative).sum())
|
||||
return {
|
||||
"threshold": float(threshold),
|
||||
"accuracy": float((predictions == labels).float().mean()),
|
||||
"precision": tp / max(1, tp + fp),
|
||||
"recall": tp / max(1, tp + fn),
|
||||
"specificity": tn / max(1, tn + fp),
|
||||
"false_positive_rate": fp / max(1, fp + tn),
|
||||
"f1": (2.0 * tp) / max(1, 2 * tp + fp + fn),
|
||||
"positive_count": int(positive.sum()),
|
||||
"negative_count": int(negative.sum()),
|
||||
}
|
||||
|
||||
|
||||
def _choose_threshold(logits: torch.Tensor, labels: torch.Tensor, max_fpr: float) -> dict[str, float]:
|
||||
candidates = [index / 100.0 for index in range(10, 91, 2)]
|
||||
reports = [_metrics(logits, labels, threshold) for threshold in candidates]
|
||||
acceptable = [item for item in reports if item["false_positive_rate"] <= max_fpr]
|
||||
return max(acceptable or reports, key=lambda item: (item["recall"], item["specificity"]))
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-path", default="qwen3_5_4b_natural_memory_v2")
|
||||
parser.add_argument("--base-adapter", default="qwen_memory_adapter_natural_auto_v13")
|
||||
parser.add_argument("--dataset-dir", default="data/production_memory")
|
||||
parser.add_argument("--output-adapter", default="qwen_memory_adapter_natural_production_candidate")
|
||||
parser.add_argument("--steps", type=int, default=240)
|
||||
parser.add_argument("--batch-size", type=int, default=4)
|
||||
parser.add_argument("--lr", type=float, default=1e-4)
|
||||
parser.add_argument("--threshold", type=float, default=None)
|
||||
parser.add_argument("--max-fpr", type=float, default=0.02)
|
||||
parser.add_argument("--max-examples", type=int, default=None)
|
||||
parser.add_argument("--seed", type=int, default=20260905)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
args = parser.parse_args()
|
||||
random.seed(args.seed)
|
||||
torch.manual_seed(args.seed)
|
||||
|
||||
dataset_dir = _project_path(args.dataset_dir)
|
||||
train_rows = _read_jsonl(dataset_dir / "train.jsonl", args.max_examples)
|
||||
eval_rows = _read_jsonl(dataset_dir / "eval.jsonl", args.max_examples)
|
||||
model_path = _project_path(args.model_path)
|
||||
base_adapter = _project_path(args.base_adapter)
|
||||
tokenizer = load_tokenizer(model_path)
|
||||
# A merged v2 package carries the router architecture metadata. Loading
|
||||
# the older v1 adapter config here would disable hierarchical memory before
|
||||
# the embedded shard is read, so the package config is authoritative.
|
||||
config_source = model_path if (model_path / "memory_merge.json").exists() else base_adapter
|
||||
config = load_memory_config(config_source)
|
||||
config.natural_language_memory = True
|
||||
config.automatic_memory = True
|
||||
config.automatic_memory_policy_version = 2
|
||||
config.auto_forget_threshold = 0.50
|
||||
config.persistent_memory = False
|
||||
model = load_qwen_dynamic(model_path, memory_config=config, load_in_4bit=not args.no_4bit)
|
||||
# The base adapter has the old one-logit policy. Its write head is a
|
||||
# useful initialization; the new forget head starts trainable and is
|
||||
# intentionally loaded with strict=False.
|
||||
model.load_memory_adapter(base_adapter, strict=False)
|
||||
if model.memory_policy is None:
|
||||
raise RuntimeError("automatic memory policy is disabled by the selected configuration")
|
||||
# The policy is now trained on the frozen Qwen semantic summary rather
|
||||
# than the value-path projection used by the bootstrap adapter. Reset
|
||||
# only this small controller so stale input-space weights cannot poison
|
||||
# the new feature space; the main model, memory bank, and retriever stay
|
||||
# untouched.
|
||||
model.memory_policy.apply(model.memory_policy._init_weights)
|
||||
model.eval()
|
||||
model.memory_policy.train()
|
||||
device = model._find_layer_device()
|
||||
|
||||
train_vectors, train_labels, train_forget_labels = _collect(
|
||||
model, tokenizer, train_rows, batch_size=args.batch_size, device=device
|
||||
)
|
||||
eval_vectors, eval_labels, eval_forget_labels = _collect(
|
||||
model, tokenizer, eval_rows, batch_size=args.batch_size, device=device
|
||||
)
|
||||
train_vectors = train_vectors.to(device)
|
||||
train_labels = train_labels.to(device)
|
||||
train_forget_labels = train_forget_labels.to(device)
|
||||
eval_vectors = eval_vectors.to(device)
|
||||
eval_labels = eval_labels.to(device)
|
||||
eval_forget_labels = eval_forget_labels.to(device)
|
||||
optimizer = torch.optim.AdamW(model.memory_policy.parameters(), lr=args.lr, weight_decay=0.01)
|
||||
positive_weight = torch.tensor([1.5], device=device)
|
||||
forget_positive_count = max(1, int(train_forget_labels.sum().item()))
|
||||
forget_negative_count = max(1, int(train_forget_labels.numel() - forget_positive_count))
|
||||
forget_positive_weight = max(4.0, 0.75 * forget_negative_count / forget_positive_count)
|
||||
rng = random.Random(args.seed + 1)
|
||||
steps = max(1, int(args.steps))
|
||||
for step in range(1, steps + 1):
|
||||
indices = torch.tensor(
|
||||
[rng.randrange(train_vectors.shape[0]) for _ in range(min(args.batch_size, train_vectors.shape[0]))],
|
||||
dtype=torch.long,
|
||||
device=device,
|
||||
)
|
||||
logits = model.memory_policy(train_vectors[indices]).reshape(-1)
|
||||
labels = train_labels[indices]
|
||||
weights = torch.where(labels >= 0.5, positive_weight.expand_as(labels), torch.ones_like(labels))
|
||||
write_loss = F.binary_cross_entropy_with_logits(logits, labels, weight=weights)
|
||||
forget_logits = model.memory_policy.forget_logits(train_vectors[indices]).reshape(-1)
|
||||
forget_labels = train_forget_labels[indices]
|
||||
forget_weights = torch.where(
|
||||
forget_labels >= 0.5,
|
||||
torch.full_like(forget_labels, forget_positive_weight),
|
||||
torch.ones_like(forget_labels),
|
||||
)
|
||||
forget_loss = F.binary_cross_entropy_with_logits(
|
||||
forget_logits, forget_labels, weight=forget_weights
|
||||
)
|
||||
loss = write_loss + forget_loss
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(model.memory_policy.parameters(), 1.0)
|
||||
optimizer.step()
|
||||
|
||||
model.memory_policy.eval()
|
||||
model._memory_policy_ready = True
|
||||
train_logits = model.memory_policy(train_vectors).reshape(-1)
|
||||
eval_logits = model.memory_policy(eval_vectors).reshape(-1)
|
||||
train_forget_logits = model.memory_policy.forget_logits(train_vectors).reshape(-1)
|
||||
eval_forget_logits = model.memory_policy.forget_logits(eval_vectors).reshape(-1)
|
||||
selected = _choose_threshold(eval_logits, eval_labels, args.max_fpr)
|
||||
selected_forget = _choose_threshold(eval_forget_logits, eval_forget_labels, args.max_fpr)
|
||||
threshold = float(args.threshold if args.threshold is not None else selected["threshold"])
|
||||
forget_threshold = float(selected_forget["threshold"])
|
||||
# The selected thresholds are part of the adapter contract. Keeping the
|
||||
# write threshold local to the report would silently revert to the config
|
||||
# default when the adapter is loaded by the benchmark or service.
|
||||
model.memory_config.auto_memory_threshold = threshold
|
||||
model.memory_config.auto_forget_threshold = forget_threshold
|
||||
output_dir = _project_path(args.output_adapter)
|
||||
# The embedded package loader temporarily restores its user snapshot and
|
||||
# marks the config persistent. A policy adapter must be stateless: never
|
||||
# ship the source user's memory with a training candidate.
|
||||
model.memory_config.persistent_memory = False
|
||||
stale_user_state = output_dir / "persistent_memory.pt"
|
||||
if stale_user_state.exists():
|
||||
stale_user_state.rename(output_dir / "persistent_memory.pt.disabled")
|
||||
model.save_memory_adapter(output_dir)
|
||||
report = {
|
||||
"format_version": 1,
|
||||
"dataset_dir": str(dataset_dir),
|
||||
"base_adapter": str(base_adapter),
|
||||
"steps": steps,
|
||||
"train_examples": len(train_rows),
|
||||
"eval_examples": len(eval_rows),
|
||||
"selected_threshold": selected,
|
||||
"threshold": threshold,
|
||||
"selected_forget_threshold": selected_forget,
|
||||
"forget_threshold": forget_threshold,
|
||||
"forget_positive_weight": forget_positive_weight,
|
||||
"train": _metrics(train_logits, train_labels, threshold),
|
||||
"eval": _metrics(eval_logits, eval_labels, threshold),
|
||||
"train_forget": _metrics(
|
||||
train_forget_logits,
|
||||
train_forget_labels,
|
||||
forget_threshold,
|
||||
),
|
||||
"eval_forget": _metrics(
|
||||
eval_forget_logits,
|
||||
eval_forget_labels,
|
||||
forget_threshold,
|
||||
),
|
||||
"forget_label_count": int(sum(float(row.get("forget_label", 0.0)) >= 0.5 for row in train_rows + eval_rows)),
|
||||
"warning": "This candidate must pass the full benchmark before replacing the production adapter.",
|
||||
}
|
||||
(output_dir / "production_policy_training.json").write_text(
|
||||
json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8"
|
||||
)
|
||||
print(json.dumps(report, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,249 @@
|
||||
"""Train only the dynamic memory adapter on streaming SFT records.
|
||||
|
||||
Each JSONL record must contain ``memory`` and ``query`` message lists. The
|
||||
memory turn is observed first; the query turn is evaluated afterwards using
|
||||
the updated state. The loss is therefore downstream of a differentiable write.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import random
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch.nn.utils import clip_grad_norm_
|
||||
|
||||
from .qwen_integration import QwenMemoryConfig, load_qwen_dynamic, load_tokenizer
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-path", default=".")
|
||||
parser.add_argument("--data", default="dynamic_memory_lab/data/demo_stream.jsonl")
|
||||
parser.add_argument("--output-dir", default="dynamic_memory_lab/qwen_memory_adapter")
|
||||
parser.add_argument("--steps", type=int, default=20)
|
||||
parser.add_argument("--batch-size", type=int, default=1)
|
||||
parser.add_argument("--lr", type=float, default=1e-4)
|
||||
parser.add_argument("--max-length", type=int, default=512)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--resume-adapter", default=None)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
parser.add_argument(
|
||||
"--surgery-mode",
|
||||
choices=("residual", "blend", "replace"),
|
||||
default="residual",
|
||||
help="residual adds memory, blend learns a gradual token-mixer replacement, replace removes the original token mixer",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--blend-init",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="initial fraction of the token mixer supplied by memory in blend mode",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--layer-indices",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=None,
|
||||
help="zero-based Qwen layer indices to adapt; defaults to four full-attention layers",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--direct-logit-scale",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="add the final memory readout directly to vocabulary logits; useful for exact recall experiments",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--write-token-offset",
|
||||
type=int,
|
||||
default=None,
|
||||
help="write a fixed token counted from the end of the memory sequence instead of the final token",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--broadcast-write",
|
||||
action="store_true",
|
||||
help="write the proposal to every memory slot; useful for one-fact recall ablations",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--raw-token-write",
|
||||
action="store_true",
|
||||
help="write the selected token's output-projection row into runtime memory",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--raw-logit-scale",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="scale the raw token memory logits during training and generation",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--native-mode",
|
||||
action="store_true",
|
||||
help="use the learned write/forget controller instead of the legacy memory rule",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--persistent-memory",
|
||||
action="store_true",
|
||||
help="keep the learned runtime memory as part of the model instance",
|
||||
)
|
||||
parser.add_argument("--reset-token-id", type=int, default=None)
|
||||
parser.add_argument(
|
||||
"--no-summary-pooling",
|
||||
action="store_true",
|
||||
help="use the final hidden state instead of learned summary pooling",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def load_records(path: str | Path) -> list[dict[str, Any]]:
|
||||
records = []
|
||||
for line_number, line in enumerate(Path(path).read_text(encoding="utf-8").splitlines(), 1):
|
||||
if not line.strip():
|
||||
continue
|
||||
record = json.loads(line)
|
||||
if not isinstance(record.get("memory"), list) or not isinstance(record.get("query"), list):
|
||||
raise ValueError(f"line {line_number}: expected memory/query message lists")
|
||||
records.append(record)
|
||||
if not records:
|
||||
raise ValueError(f"no records found in {path}")
|
||||
return records
|
||||
|
||||
|
||||
def encode_messages(tokenizer: Any, messages: list[dict[str, Any]], max_length: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
full_encoding = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
add_generation_prompt=False,
|
||||
enable_thinking=False,
|
||||
)
|
||||
prompt_encoding = tokenizer.apply_chat_template(
|
||||
messages[:-1],
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
full = full_encoding["input_ids"] if hasattr(full_encoding, "__getitem__") and "input_ids" in full_encoding else full_encoding
|
||||
prompt = prompt_encoding["input_ids"] if hasattr(prompt_encoding, "__getitem__") and "input_ids" in prompt_encoding else prompt_encoding
|
||||
if full and isinstance(full[0], list):
|
||||
full = full[0]
|
||||
if prompt and isinstance(prompt[0], list):
|
||||
prompt = prompt[0]
|
||||
original_full_length = len(full)
|
||||
truncated_prefix = max(0, original_full_length - max_length)
|
||||
if len(full) > max_length:
|
||||
full = full[-max_length:]
|
||||
input_ids = torch.tensor(full, dtype=torch.long)
|
||||
attention_mask = torch.ones_like(input_ids)
|
||||
# The prompt may have been truncated from the left. Keep only the target
|
||||
# part visible to the loss.
|
||||
prompt_tokens = max(0, min(len(prompt) - truncated_prefix, len(full)))
|
||||
labels = input_ids.clone()
|
||||
labels[:prompt_tokens] = -100
|
||||
return input_ids, attention_mask, labels
|
||||
|
||||
|
||||
def pad_batch(items: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]], pad_id: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
length = max(item[0].numel() for item in items)
|
||||
input_ids = torch.full((len(items), length), pad_id, dtype=torch.long)
|
||||
attention_mask = torch.zeros_like(input_ids)
|
||||
labels = torch.full_like(input_ids, -100)
|
||||
for row, (ids, mask, row_labels) in enumerate(items):
|
||||
input_ids[row, : ids.numel()] = ids
|
||||
attention_mask[row, : mask.numel()] = mask
|
||||
labels[row, : row_labels.numel()] = row_labels
|
||||
return input_ids, attention_mask, labels
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
random.seed(args.seed)
|
||||
torch.manual_seed(args.seed)
|
||||
records = load_records(args.data)
|
||||
tokenizer = load_tokenizer(args.model_path)
|
||||
memory_config = QwenMemoryConfig(
|
||||
mode=args.surgery_mode,
|
||||
blend_init=args.blend_init,
|
||||
layer_indices=tuple(args.layer_indices) if args.layer_indices else None,
|
||||
direct_logit_scale=args.direct_logit_scale,
|
||||
write_token_offset=args.write_token_offset,
|
||||
broadcast_write=args.broadcast_write,
|
||||
raw_token_write=args.raw_token_write,
|
||||
raw_logit_scale=args.raw_logit_scale,
|
||||
native_mode=args.native_mode,
|
||||
persistent_memory=args.persistent_memory,
|
||||
reset_token_id=args.reset_token_id,
|
||||
summary_pooling=not args.no_summary_pooling,
|
||||
)
|
||||
model = load_qwen_dynamic(
|
||||
args.model_path,
|
||||
memory_config=memory_config,
|
||||
load_in_4bit=not args.no_4bit,
|
||||
)
|
||||
if args.resume_adapter:
|
||||
model.load_memory_adapter(args.resume_adapter)
|
||||
model.train()
|
||||
|
||||
parameters = list(model.trainable_parameters)
|
||||
if not parameters:
|
||||
raise RuntimeError("no trainable memory parameters")
|
||||
optimizer = torch.optim.AdamW(parameters, lr=args.lr, weight_decay=0.01)
|
||||
device = model._find_layer_device()
|
||||
pad_id = int(tokenizer.pad_token_id)
|
||||
output_dir = Path(args.output_dir)
|
||||
|
||||
print(
|
||||
f"device={device} records={len(records)} memory_layers={model.layer_indices} "
|
||||
f"surgery_mode={memory_config.mode} blend_init={memory_config.blend_init}"
|
||||
)
|
||||
for step in range(1, args.steps + 1):
|
||||
if memory_config.persistent_memory:
|
||||
model.reset_memory(batch_size=args.batch_size, device=model._find_layer_device())
|
||||
chosen = [records[(step * args.batch_size + i) % len(records)] for i in range(args.batch_size)]
|
||||
memory_items = [encode_messages(tokenizer, item["memory"], args.max_length) for item in chosen]
|
||||
query_items = [encode_messages(tokenizer, item["query"], args.max_length) for item in chosen]
|
||||
memory_input, memory_mask, _ = pad_batch(memory_items, pad_id)
|
||||
query_input, query_mask, query_labels = pad_batch(query_items, pad_id)
|
||||
memory_input = memory_input.to(device)
|
||||
memory_mask = memory_mask.to(device)
|
||||
query_input = query_input.to(device)
|
||||
query_mask = query_mask.to(device)
|
||||
query_labels = query_labels.to(device)
|
||||
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
memory_output = model(
|
||||
input_ids=memory_input,
|
||||
attention_mask=memory_mask,
|
||||
update_memory=True,
|
||||
read_memory=False,
|
||||
detach_memory=False,
|
||||
return_memory=True,
|
||||
)
|
||||
query_output = model(
|
||||
input_ids=query_input,
|
||||
attention_mask=query_mask,
|
||||
labels=query_labels,
|
||||
memory_state=memory_output.memory,
|
||||
update_memory=False,
|
||||
read_memory=True,
|
||||
return_memory=True,
|
||||
)
|
||||
if query_output.loss is None:
|
||||
raise RuntimeError("Qwen did not return an SFT loss")
|
||||
query_output.loss.backward()
|
||||
clip_grad_norm_(parameters, 1.0)
|
||||
optimizer.step()
|
||||
|
||||
if step == 1 or step % 5 == 0 or step == args.steps:
|
||||
print(f"step={step:4d} loss={query_output.loss.detach().item():.4f}")
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
model.save_memory_adapter(output_dir)
|
||||
(output_dir / "training_state.json").write_text(
|
||||
json.dumps({"step": step, "data": str(args.data), "model_path": str(args.model_path)}, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,367 @@
|
||||
"""Train the V2 router on real Qwen hidden representations.
|
||||
|
||||
This is intentionally separate from the fast synthetic router pre-training.
|
||||
The production checkpoint must see the same representation distribution that
|
||||
the memory adapter will use at runtime.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import random
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
if __package__ in {None, ""}:
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
from dynamic_memory_lab.memory_os_v2 import MemoryRouterV2
|
||||
from dynamic_memory_lab.qwen_integration import QwenMemoryConfig, load_qwen_dynamic, load_tokenizer
|
||||
|
||||
|
||||
ATTRIBUTES = [
|
||||
"姓名",
|
||||
"常住城市",
|
||||
"工作地点",
|
||||
"项目代号",
|
||||
"喜欢的水果",
|
||||
"宠物名字",
|
||||
"生日月份",
|
||||
"最常用的编辑器",
|
||||
"长期目标",
|
||||
"周末习惯",
|
||||
"学习方向",
|
||||
"重要联系人",
|
||||
]
|
||||
ENTITIES = [f"用户档案{index:02d}" for index in range(64)]
|
||||
VALUES = [
|
||||
"林浩",
|
||||
"上海",
|
||||
"杭州",
|
||||
"NM-V2",
|
||||
"青提",
|
||||
"小灰",
|
||||
"十月",
|
||||
"Neovim",
|
||||
"做出新的记忆架构",
|
||||
"阅读论文",
|
||||
"稀疏路由",
|
||||
"陈老师",
|
||||
"苏州",
|
||||
"Natural Memory",
|
||||
"星河项目",
|
||||
"午夜跑步",
|
||||
]
|
||||
|
||||
|
||||
def _device(name: str, fallback: torch.device | None = None) -> torch.device:
|
||||
if name == "auto":
|
||||
return fallback or torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
return torch.device(name)
|
||||
|
||||
|
||||
def _encode_texts(model, tokenizer, texts: list[str], device: torch.device, batch_size: int) -> torch.Tensor:
|
||||
chunks: list[torch.Tensor] = []
|
||||
for start in range(0, len(texts), batch_size):
|
||||
batch = texts[start : start + batch_size]
|
||||
encoded = tokenizer(
|
||||
batch,
|
||||
add_special_tokens=False,
|
||||
padding=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
input_ids = encoded["input_ids"].to(device)
|
||||
attention_mask = encoded.get("attention_mask")
|
||||
if attention_mask is None:
|
||||
attention_mask = torch.ones_like(input_ids)
|
||||
key = model._encode_model_key(input_ids, attention_mask.to(device))
|
||||
chunks.append(key.detach().cpu())
|
||||
return torch.cat(chunks, dim=0)
|
||||
|
||||
|
||||
def _make_fact_set(count: int, seed: int) -> tuple[list[str], list[str], list[int]]:
|
||||
random.seed(seed)
|
||||
facts: list[str] = []
|
||||
queries: list[str] = []
|
||||
hops: list[int] = []
|
||||
for index in range(count):
|
||||
attribute = ATTRIBUTES[index % len(ATTRIBUTES)]
|
||||
value = VALUES[(index * 7 + 3) % len(VALUES)]
|
||||
entity = ENTITIES[index // len(ATTRIBUTES)]
|
||||
fact = f"{entity}的{attribute}是{value}。"
|
||||
query_templates = [
|
||||
(f"请问{entity}的{attribute}是什么？", 1),
|
||||
(f"我之前告诉过你的{entity}{attribute}，答案是什么？", 1),
|
||||
(f"先找出{entity}的{attribute}，再结合关联记忆回答。", 2),
|
||||
(f"回忆一下，{entity}在{attribute}这一项的信息。", 1),
|
||||
]
|
||||
facts.append(fact)
|
||||
query, hop = query_templates[index % len(query_templates)]
|
||||
queries.append(query)
|
||||
hops.append(hop)
|
||||
return facts, queries, hops
|
||||
|
||||
|
||||
def _make_episodes(
|
||||
fact_keys: torch.Tensor,
|
||||
query_keys: torch.Tensor,
|
||||
fact_hops: list[int],
|
||||
*,
|
||||
candidate_count: int,
|
||||
seed: int,
|
||||
no_memory_keys: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
random.seed(seed)
|
||||
sample_count = fact_keys.shape[0]
|
||||
queries: list[torch.Tensor] = []
|
||||
candidates: list[torch.Tensor] = []
|
||||
positives: list[int] = []
|
||||
need: list[float] = []
|
||||
hops: list[int] = []
|
||||
for index in range(sample_count):
|
||||
candidate_indices = [index]
|
||||
# Prefer same-entity/nearby attribute negatives before random ones.
|
||||
for offset in range(1, sample_count):
|
||||
candidate_indices.append((index + offset) % sample_count)
|
||||
if len(candidate_indices) >= candidate_count:
|
||||
break
|
||||
candidate_tensor = fact_keys[candidate_indices]
|
||||
queries.append(query_keys[index])
|
||||
candidates.append(candidate_tensor)
|
||||
positives.append(0)
|
||||
need.append(1.0)
|
||||
hops.append(int(fact_hops[index]))
|
||||
for index in range(min(sample_count // 3, no_memory_keys.shape[0])):
|
||||
queries.append(no_memory_keys[index])
|
||||
candidates.append(fact_keys[torch.randperm(sample_count)[:candidate_count]])
|
||||
positives.append(0)
|
||||
need.append(0.0)
|
||||
hops.append(0)
|
||||
return (
|
||||
torch.stack(queries),
|
||||
torch.stack(candidates),
|
||||
torch.tensor(positives, dtype=torch.long),
|
||||
torch.tensor(need, dtype=torch.float32),
|
||||
torch.tensor(hops, dtype=torch.long),
|
||||
)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def _evaluate(
|
||||
router: MemoryRouterV2,
|
||||
query: torch.Tensor,
|
||||
candidates: torch.Tensor,
|
||||
positive: torch.Tensor,
|
||||
need: torch.Tensor,
|
||||
hops: torch.Tensor,
|
||||
device: torch.device,
|
||||
) -> dict[str, float]:
|
||||
router.eval()
|
||||
output = router(query.to(device), candidates.to(device))
|
||||
need_mask = need.to(device).bool()
|
||||
route_pred = output["scores"].argmax(dim=-1)
|
||||
route_correct = ((route_pred == positive.to(device)) & need_mask).sum().item()
|
||||
route_total = need_mask.sum().item()
|
||||
need_pred = (torch.sigmoid(output["need_memory_logits"]) >= 0.5).float()
|
||||
need_device = need.to(device)
|
||||
tp = ((need_pred == 1) & (need_device == 1)).sum().item()
|
||||
tn = ((need_pred == 0) & (need_device == 0)).sum().item()
|
||||
fp = ((need_pred == 1) & (need_device == 0)).sum().item()
|
||||
fn = ((need_pred == 0) & (need_device == 1)).sum().item()
|
||||
hop_accuracy = (output["hop_logits"].argmax(dim=-1) == hops.to(device)).float().mean().item()
|
||||
return {
|
||||
"route_accuracy": route_correct / max(1, route_total),
|
||||
"need_memory_precision": tp / max(1, tp + fp),
|
||||
"need_memory_recall": tp / max(1, tp + fn),
|
||||
"need_memory_specificity": tn / max(1, tn + fp),
|
||||
"hop_accuracy": hop_accuracy,
|
||||
"need_tp": float(tp),
|
||||
"need_tn": float(tn),
|
||||
"need_fp": float(fp),
|
||||
"need_fn": float(fn),
|
||||
}
|
||||
|
||||
|
||||
def train(args: argparse.Namespace) -> dict[str, Any]:
|
||||
random.seed(args.seed)
|
||||
torch.manual_seed(args.seed)
|
||||
config = QwenMemoryConfig(
|
||||
memory_slots=16,
|
||||
memory_dim=512,
|
||||
layer_indices=(7, 15, 23, 31),
|
||||
mode="blend",
|
||||
blend_init=0.1,
|
||||
native_mode=True,
|
||||
persistent_memory=True,
|
||||
natural_language_memory=True,
|
||||
automatic_memory=True,
|
||||
memory_version=2,
|
||||
hierarchical_memory=True,
|
||||
memory_router_dim=args.router_dim,
|
||||
memory_router_heads=args.num_heads,
|
||||
memory_max_hops=args.max_hops,
|
||||
)
|
||||
model = load_qwen_dynamic(
|
||||
args.model_path,
|
||||
memory_config=config,
|
||||
load_in_4bit=not args.no_4bit,
|
||||
device_map="auto",
|
||||
)
|
||||
model.eval()
|
||||
model_device = model._find_layer_device()
|
||||
device = _device(args.device, model_device)
|
||||
tokenizer = load_tokenizer(args.model_path)
|
||||
fact_texts, query_texts, fact_hops = _make_fact_set(args.fact_count, args.seed)
|
||||
no_memory_texts = [
|
||||
"请写一首关于春天的短诗。",
|
||||
"解释一下二分查找的时间复杂度。",
|
||||
"帮我规划一个周末旅行。",
|
||||
"什么是矩阵乘法？",
|
||||
"把这句话翻译成英文。",
|
||||
"今天适合做什么运动？",
|
||||
]
|
||||
fact_keys = _encode_texts(model, tokenizer, fact_texts, model_device, args.encode_batch_size)
|
||||
query_keys = _encode_texts(model, tokenizer, query_texts, model_device, args.encode_batch_size)
|
||||
no_memory_keys = _encode_texts(model, tokenizer, no_memory_texts, model_device, args.encode_batch_size)
|
||||
train_count = max(1, int(fact_keys.shape[0] * 0.8))
|
||||
train_query, train_candidates, train_positive, train_need, train_hops = _make_episodes(
|
||||
fact_keys[:train_count],
|
||||
query_keys[:train_count],
|
||||
fact_hops[:train_count],
|
||||
candidate_count=args.candidate_count,
|
||||
seed=args.seed,
|
||||
no_memory_keys=no_memory_keys,
|
||||
)
|
||||
heldout_query, heldout_candidates, heldout_positive, heldout_need, heldout_hops = _make_episodes(
|
||||
fact_keys[train_count:],
|
||||
query_keys[train_count:],
|
||||
fact_hops[train_count:],
|
||||
candidate_count=args.candidate_count,
|
||||
seed=args.seed + 1,
|
||||
no_memory_keys=no_memory_keys,
|
||||
)
|
||||
router = MemoryRouterV2(
|
||||
args.hidden_size,
|
||||
router_dim=args.router_dim,
|
||||
num_heads=args.num_heads,
|
||||
max_hops=args.max_hops,
|
||||
).to(device)
|
||||
if args.init_checkpoint:
|
||||
initial_state = torch.load(args.init_checkpoint, map_location=device, weights_only=True)
|
||||
router.load_state_dict(initial_state, strict=True)
|
||||
if args.freeze_retrieval:
|
||||
for name, parameter in router.named_parameters():
|
||||
parameter.requires_grad = name.startswith("hop_controller.")
|
||||
trainable_parameters = [parameter for parameter in router.parameters() if parameter.requires_grad]
|
||||
optimizer = torch.optim.AdamW(trainable_parameters, lr=args.learning_rate, weight_decay=1e-4)
|
||||
history: list[dict[str, float]] = []
|
||||
router.train()
|
||||
for step in range(1, args.steps + 1):
|
||||
indices = torch.randint(0, train_query.shape[0], (args.batch_size,))
|
||||
query = train_query[indices].to(device)
|
||||
candidates = train_candidates[indices].to(device)
|
||||
positive = train_positive[indices].to(device)
|
||||
need = train_need[indices].to(device)
|
||||
hops = train_hops[indices].clamp(0, args.max_hops).to(device)
|
||||
output = router(query, candidates)
|
||||
need_mask = need.bool()
|
||||
candidate_loss = (
|
||||
F.cross_entropy(output["scores"][need_mask], positive[need_mask])
|
||||
if bool(need_mask.any())
|
||||
else output["scores"].sum() * 0.0
|
||||
)
|
||||
need_loss = F.binary_cross_entropy_with_logits(output["need_memory_logits"], need)
|
||||
hop_loss = F.cross_entropy(output["hop_logits"], hops)
|
||||
loss = candidate_loss + args.need_loss_weight * need_loss + args.hop_loss_weight * hop_loss
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(router.parameters(), 1.0)
|
||||
optimizer.step()
|
||||
if step == 1 or step % args.log_every == 0 or step == args.steps:
|
||||
history.append(
|
||||
{
|
||||
"step": float(step),
|
||||
"loss": float(loss.detach().cpu()),
|
||||
"candidate_loss": float(candidate_loss.detach().cpu()),
|
||||
"need_loss": float(need_loss.detach().cpu()),
|
||||
"hop_loss": float(hop_loss.detach().cpu()),
|
||||
}
|
||||
)
|
||||
validation = _evaluate(
|
||||
router,
|
||||
heldout_query,
|
||||
heldout_candidates,
|
||||
heldout_positive,
|
||||
heldout_need,
|
||||
heldout_hops,
|
||||
device,
|
||||
)
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
torch.save(router.state_dict(), output_dir / "memory_router_v2.pt")
|
||||
torch.save(
|
||||
{
|
||||
"fact_keys": fact_keys,
|
||||
"query_keys": query_keys,
|
||||
"no_memory_keys": no_memory_keys,
|
||||
"fact_texts": fact_texts,
|
||||
"query_texts": query_texts,
|
||||
},
|
||||
output_dir / "qwen_router_v2_encoded_dataset.pt",
|
||||
)
|
||||
summary = {
|
||||
"format_version": 2,
|
||||
"representation": "qwen3.5_hidden_state",
|
||||
"model_path": args.model_path,
|
||||
"device": str(device),
|
||||
"model_device": str(model_device),
|
||||
"hidden_size": args.hidden_size,
|
||||
"router_dim": args.router_dim,
|
||||
"num_heads": args.num_heads,
|
||||
"max_hops": args.max_hops,
|
||||
"fact_count": args.fact_count,
|
||||
"train_count": train_count,
|
||||
"heldout_count": int(fact_keys.shape[0] - train_count),
|
||||
"steps": args.steps,
|
||||
"training_history": history,
|
||||
"validation": validation,
|
||||
}
|
||||
(output_dir / "qwen_router_v2_training.json").write_text(
|
||||
json.dumps(summary, indent=2, ensure_ascii=False), encoding="utf-8"
|
||||
)
|
||||
del model
|
||||
return summary
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-path", default="W:/Flash/model/dynamic_memory_lab/qwen3_5_4b_memory_merged_v13")
|
||||
parser.add_argument("--output-dir", default="W:/Flash/model/dynamic_memory_lab/checkpoints/natural_memory_v2_qwen_router")
|
||||
parser.add_argument("--device", default="auto")
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
parser.add_argument("--hidden-size", type=int, default=2560)
|
||||
parser.add_argument("--router-dim", type=int, default=128)
|
||||
parser.add_argument("--num-heads", type=int, default=8)
|
||||
parser.add_argument("--max-hops", type=int, default=3)
|
||||
parser.add_argument("--fact-count", type=int, default=160)
|
||||
parser.add_argument("--candidate-count", type=int, default=16)
|
||||
parser.add_argument("--encode-batch-size", type=int, default=8)
|
||||
parser.add_argument("--steps", type=int, default=1200)
|
||||
parser.add_argument("--batch-size", type=int, default=32)
|
||||
parser.add_argument("--learning-rate", type=float, default=2e-3)
|
||||
parser.add_argument("--need-loss-weight", type=float, default=0.75)
|
||||
parser.add_argument("--hop-loss-weight", type=float, default=0.35)
|
||||
parser.add_argument("--init-checkpoint", default=None)
|
||||
parser.add_argument("--freeze-retrieval", action="store_true")
|
||||
parser.add_argument("--log-every", type=int, default=100)
|
||||
parser.add_argument("--seed", type=int, default=20260904)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(json.dumps(train(parse_args()), ensure_ascii=False, indent=2))
|
||||
@@ -0,0 +1,109 @@
|
||||
"""Package learned memory into an adapter checkpoint and verify restart/reset."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from .evaluate_native_memory import _controller_step, _encode_prompt
|
||||
from .qwen_integration import (
|
||||
DEFAULT_MEMORY_RESET_TOKEN,
|
||||
load_memory_config,
|
||||
load_qwen_dynamic,
|
||||
load_tokenizer,
|
||||
resolve_memory_reset_token,
|
||||
)
|
||||
from .train_native_memory import load_records
|
||||
|
||||
|
||||
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_native_v3")
|
||||
parser.add_argument("--data", default="dynamic_memory_lab/data/native_memory/eval.jsonl")
|
||||
parser.add_argument(
|
||||
"--output-adapter",
|
||||
default="dynamic_memory_lab/qwen_memory_adapter_native_v3_persistent",
|
||||
)
|
||||
parser.add_argument("--max-length", type=int, default=192)
|
||||
parser.add_argument("--max-new-tokens", type=int, default=8)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
records = load_records(args.data)
|
||||
record = next(record for record in records if record.get("answerable"))
|
||||
tokenizer = load_tokenizer(args.model_path)
|
||||
config = load_memory_config(args.adapter)
|
||||
config.persistent_memory = True
|
||||
config.reset_token_id = resolve_memory_reset_token(tokenizer, DEFAULT_MEMORY_RESET_TOKEN)
|
||||
model = load_qwen_dynamic(args.model_path, memory_config=config, load_in_4bit=not args.no_4bit)
|
||||
model.load_memory_adapter(args.adapter)
|
||||
model.reset_memory()
|
||||
for chunk in record["memory_chunks"]:
|
||||
_controller_step(model, tokenizer, chunk, args.max_length)
|
||||
output_adapter = Path(args.output_adapter)
|
||||
model.save_persistent_memory_checkpoint(output_adapter)
|
||||
saved_norm = float(model.runtime.state.detach().float().norm())
|
||||
del model
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
restart_config = load_memory_config(output_adapter)
|
||||
restarted = load_qwen_dynamic(
|
||||
args.model_path,
|
||||
memory_config=restart_config,
|
||||
load_in_4bit=not args.no_4bit,
|
||||
)
|
||||
restarted.load_memory_adapter(output_adapter)
|
||||
answer = str(record["answer"])
|
||||
prompt = _encode_prompt(tokenizer, record["query"][:-1])
|
||||
device = restarted._find_layer_device()
|
||||
generated = restarted.generate(
|
||||
input_ids=prompt.to(device),
|
||||
attention_mask=torch.ones_like(prompt, device=device),
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
do_sample=False,
|
||||
use_cache=False,
|
||||
update_memory=False,
|
||||
)
|
||||
response = tokenizer.decode(
|
||||
generated[:, prompt.shape[1] :][0].detach().cpu().tolist(),
|
||||
skip_special_tokens=True,
|
||||
).strip()
|
||||
reset_prompt = _encode_prompt(tokenizer, [{"role": "user", "content": DEFAULT_MEMORY_RESET_TOKEN}])
|
||||
restarted.generate(
|
||||
input_ids=reset_prompt.to(device),
|
||||
attention_mask=torch.ones_like(reset_prompt, device=device),
|
||||
max_new_tokens=1,
|
||||
do_sample=False,
|
||||
use_cache=False,
|
||||
update_memory=False,
|
||||
)
|
||||
reset_norm = float(restarted.runtime.state.detach().float().norm())
|
||||
report = {
|
||||
"source_adapter": str(args.adapter),
|
||||
"output_adapter": str(output_adapter),
|
||||
"record_id": record.get("id"),
|
||||
"answer": answer,
|
||||
"generated_after_restart": response,
|
||||
"restart_contains_answer": answer in response,
|
||||
"saved_memory_norm": saved_norm,
|
||||
"reset_token": DEFAULT_MEMORY_RESET_TOKEN,
|
||||
"reset_token_id": restart_config.reset_token_id,
|
||||
"memory_norm_after_reset_token": reset_norm,
|
||||
"reset_zeroed_memory": reset_norm < 1e-5,
|
||||
}
|
||||
(output_adapter / "native_checkpoint_verification.json").write_text(
|
||||
json.dumps(report, ensure_ascii=False, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
print(json.dumps(report, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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()
|
||||
@@ -0,0 +1,99 @@
|
||||
"""Load the local Qwen checkpoint and run one real memory-aware forward pass."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import contextlib
|
||||
import io
|
||||
|
||||
import torch
|
||||
|
||||
from .qwen_integration import QwenMemoryConfig, load_qwen_dynamic, load_tokenizer
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model-path", default=".")
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
parser.add_argument("--surgery-mode", choices=("residual", "blend", "replace"), default="residual")
|
||||
parser.add_argument("--blend-init", type=float, default=0.0)
|
||||
parser.add_argument("--layer-indices", type=int, nargs="+", default=None)
|
||||
parser.add_argument("--direct-logit-scale", type=float, default=0.0)
|
||||
parser.add_argument("--write-token-offset", type=int, default=None)
|
||||
parser.add_argument("--broadcast-write", action="store_true")
|
||||
parser.add_argument("--raw-token-write", action="store_true")
|
||||
parser.add_argument("--raw-logit-scale", type=float, default=0.0)
|
||||
parser.add_argument("--native-mode", action="store_true")
|
||||
parser.add_argument("--persistent-memory", action="store_true")
|
||||
parser.add_argument("--reset-token-id", type=int, default=None)
|
||||
parser.add_argument("--no-summary-pooling", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
captured = io.StringIO()
|
||||
stdout = contextlib.redirect_stdout(captured)
|
||||
stderr = contextlib.redirect_stderr(captured)
|
||||
stdout.__enter__()
|
||||
stderr.__enter__()
|
||||
error = None
|
||||
model = None
|
||||
tokenizer = None
|
||||
try:
|
||||
memory_config = QwenMemoryConfig(
|
||||
mode=args.surgery_mode,
|
||||
blend_init=args.blend_init,
|
||||
layer_indices=tuple(args.layer_indices) if args.layer_indices else None,
|
||||
direct_logit_scale=args.direct_logit_scale,
|
||||
write_token_offset=args.write_token_offset,
|
||||
broadcast_write=args.broadcast_write,
|
||||
raw_token_write=args.raw_token_write,
|
||||
raw_logit_scale=args.raw_logit_scale,
|
||||
native_mode=args.native_mode,
|
||||
persistent_memory=args.persistent_memory,
|
||||
reset_token_id=args.reset_token_id,
|
||||
summary_pooling=not args.no_summary_pooling,
|
||||
)
|
||||
model = load_qwen_dynamic(
|
||||
args.model_path,
|
||||
memory_config=memory_config,
|
||||
load_in_4bit=not args.no_4bit,
|
||||
)
|
||||
tokenizer = load_tokenizer(args.model_path)
|
||||
except Exception as exc: # pragma: no cover - diagnostic entry point
|
||||
error = (type(exc).__name__, str(exc))
|
||||
finally:
|
||||
stderr.__exit__(None, None, None)
|
||||
stdout.__exit__(None, None, None)
|
||||
|
||||
print(f"load_error={error if error else 'none'}")
|
||||
if model is None or tokenizer is None:
|
||||
return
|
||||
|
||||
print(f"model={type(model.base_model).__name__}")
|
||||
print(f"device={model._find_layer_device()}")
|
||||
print(f"memory_layers={model.layer_indices}")
|
||||
print(
|
||||
"layer_types="
|
||||
+ str(
|
||||
tuple(
|
||||
getattr(getattr(model.base_model.model.language_model.layers[index], "inner", None), "layer_type", "unknown")
|
||||
for index in model.layer_indices
|
||||
)
|
||||
)
|
||||
)
|
||||
print(f"surgery_mode={model.memory_config.mode}")
|
||||
print(f"blend_init={model.memory_config.blend_init}")
|
||||
print(f"trainable_memory_parameters={sum(p.numel() for p in model.trainable_parameters):,}")
|
||||
|
||||
text = "你好，请用一句话介绍你自己。"
|
||||
encoded = tokenizer(text, return_tensors="pt")
|
||||
device = model._find_layer_device()
|
||||
encoded = {key: value.to(device) for key, value in encoded.items()}
|
||||
with torch.no_grad():
|
||||
output = model(**encoded, update_memory=True, return_memory=True)
|
||||
print(f"logits_shape={tuple(output.logits.shape)}")
|
||||
print(f"memory_shape={tuple(output.memory.shape)}")
|
||||
print("forward_ok=true")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user