WpyQwq
643e22ecb9
- 引入 MemoryRouterXL 与 v5/v6 流式多线程训练/编码管线 - 修复 prepare_memory_router_dataset 候选池重建缺陷（mega 家族 3568x 加速，输出逐字节相同） - 修复 v5 被破坏的拒答与多跳标签（train 未知样本 319 -> 16319，multi_hop 平均正例 1.00 -> 2.00） - 同存储预算下 V2-128 v6 逐轴 22/22 通过：Top-1 41.12% -> 94.62%，未知拒答 0.00% -> 100.00% - 记录三条被实测推翻的显然优化（logits_to_keep=1 反而慢 55%、XL 容量未带来收益） - 记忆手术跨架构可移植性 14/14，读写关闭时与原生模型逐位相同
1008 lines
40 KiB
Python
1008 lines
40 KiB
Python
"""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是什么",
|
||
"解释一下这个概念",
|
||
)
|
||
|
||
REDTEAM_NEGATIVE_QUERIES = (
|
||
"如果我的{attribute}改成另一个值，会发生什么？",
|
||
"别人说我的{attribute}是这个值，但那不是我的资料。",
|
||
"这次临时提到{attribute}，不用保存。",
|
||
"我想了解{attribute}这个概念，不是查询我的个人资料。",
|
||
"我的不存在的{attribute}是什么？如果没有登记就说不知道。",
|
||
"从来没有登记过的{attribute}是什么？不要从别的字段推断。",
|
||
)
|
||
|
||
|
||
def make_redteam_records(seed: int, count: int = 960) -> list[dict[str, object]]:
|
||
"""Add adversarial pairs that look semantically close but are not reads.
|
||
|
||
The original corpus mostly contrasted one attribute with another. These
|
||
examples specifically target the failure mode called out in the review:
|
||
a paraphrase, quotation, or hypothetical sentence can receive a confident
|
||
retrieval score even though it should not select a personal memory.
|
||
"""
|
||
|
||
rng = random.Random(seed + 101)
|
||
attributes = list(ATTRIBUTES)
|
||
rows: list[dict[str, object]] = []
|
||
for index in range(count // 2):
|
||
attribute = rng.choice(attributes)
|
||
value = f"RT-{index:05d}"
|
||
fact = rng.choice(ATTRIBUTES[attribute]["facts"]).format(value=value)
|
||
positive_queries = (
|
||
f"我之前登记的{attribute}是哪一个？",
|
||
f"只读取长期资料，告诉我{attribute}。",
|
||
f"跨会话后，我保存的{attribute}是什么？",
|
||
f"不要猜测，回忆我的{attribute}设置。",
|
||
)
|
||
rows.append(
|
||
{
|
||
"fact": fact,
|
||
"query": rng.choice(positive_queries),
|
||
"label": 1.0,
|
||
"attribute": attribute,
|
||
}
|
||
)
|
||
rows.append(
|
||
{
|
||
"fact": fact,
|
||
"query": rng.choice(REDTEAM_NEGATIVE_QUERIES).format(attribute=attribute),
|
||
"label": 0.0,
|
||
"attribute": f"{attribute}->redteam_negative",
|
||
}
|
||
)
|
||
rng.shuffle(rows)
|
||
return rows
|
||
|
||
|
||
def make_redteam_holdout(seed: int, count: int = 240) -> list[dict[str, object]]:
|
||
"""Generate held-out adversarial paraphrases with unseen values."""
|
||
|
||
rng = random.Random(seed + 109)
|
||
attributes = list(ATTRIBUTES)
|
||
rows: list[dict[str, object]] = []
|
||
for index in range(count):
|
||
attribute = attributes[index % len(attributes)]
|
||
value = f"HELDOUT-{rng.randrange(10**8):08d}"
|
||
fact = f"我的{attribute}是{value}。"
|
||
if index % 4 == 0:
|
||
query = f"跨对话后，之前存下来的{attribute}是哪一个？"
|
||
label = 1.0
|
||
elif index % 4 == 1:
|
||
query = f"如果把我的{attribute}改掉，应该如何规划？"
|
||
label = 0.0
|
||
elif index % 4 == 2:
|
||
query = f"别人提到我的{attribute}，但请不要把这句话当作我的资料。"
|
||
label = 0.0
|
||
else:
|
||
query = f"我的不存在的{attribute}是什么？如果没有登记就说不知道。"
|
||
label = 0.0
|
||
rows.append({"fact": fact, "query": query, "label": label, "attribute": attribute})
|
||
return rows
|
||
|
||
|
||
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
|
||
|
||
|
||
def make_update_positive_records(seed: int, count: int = 640) -> list[dict[str, object]]:
|
||
"""Create same-attribute fact-to-fact update pairs."""
|
||
|
||
rng = random.Random(seed + 31)
|
||
attributes = list(ATTRIBUTES)
|
||
records: list[dict[str, object]] = []
|
||
for index in range(count):
|
||
attribute = rng.choice(attributes)
|
||
first_value = f"UPD-OLD-{index:05d}"
|
||
second_value = f"UPD-NEW-{index:05d}"
|
||
definition = ATTRIBUTES[attribute]
|
||
first = rng.choice(definition["facts"]).format(value=first_value)
|
||
second = rng.choice(definition["facts"]).format(value=second_value)
|
||
records.append(
|
||
{
|
||
"fact": first,
|
||
"query": second,
|
||
"label": 1.0,
|
||
"attribute": attribute,
|
||
}
|
||
)
|
||
return records
|
||
|
||
|
||
def make_update_negative_records(seed: int, count: int = 640) -> list[dict[str, object]]:
|
||
"""Create same-entity, different-attribute hard negatives.
|
||
|
||
Both sides intentionally use first-person language. The only durable
|
||
distinction is the attribute, which is exactly what the hot-slot update
|
||
decision must learn instead of collapsing every personal fragment.
|
||
"""
|
||
|
||
rng = random.Random(seed + 37)
|
||
attributes = list(ATTRIBUTES)
|
||
records: list[dict[str, object]] = []
|
||
for index in range(count):
|
||
left, right = rng.sample(attributes, 2)
|
||
left_definition = ATTRIBUTES[left]
|
||
right_definition = ATTRIBUTES[right]
|
||
fact = rng.choice(left_definition["facts"]).format(value=f"NEG-L-{index:05d}")
|
||
query = rng.choice(right_definition["facts"]).format(value=f"NEG-R-{index:05d}")
|
||
records.append(
|
||
{
|
||
"fact": fact,
|
||
"query": query,
|
||
"label": 0.0,
|
||
"attribute": f"{right}!={left}",
|
||
}
|
||
)
|
||
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="V2_dpskw/qwen_memory_adapter_native_v3")
|
||
parser.add_argument(
|
||
"--output-adapter",
|
||
default="V2_dpskw/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)
|
||
update_positive_records = make_update_positive_records(args.seed)
|
||
update_positive_query_tensor, update_positive_key_tensor = encode_records(
|
||
model,
|
||
tokenizer,
|
||
update_positive_records,
|
||
batch_size=args.batch_size,
|
||
device=device,
|
||
)
|
||
update_positive_query_tensor = update_positive_query_tensor.to(device)
|
||
update_positive_key_tensor = update_positive_key_tensor.to(device)
|
||
update_negative_records = make_update_negative_records(args.seed)
|
||
update_negative_query_tensor, update_negative_key_tensor = encode_records(
|
||
model,
|
||
tokenizer,
|
||
update_negative_records,
|
||
batch_size=args.batch_size,
|
||
device=device,
|
||
)
|
||
update_negative_query_tensor = update_negative_query_tensor.to(device)
|
||
update_negative_key_tensor = update_negative_key_tensor.to(device)
|
||
redteam_records = make_redteam_records(args.seed)
|
||
redteam_query_tensor, redteam_key_tensor = encode_records(
|
||
model,
|
||
tokenizer,
|
||
redteam_records,
|
||
batch_size=args.batch_size,
|
||
device=device,
|
||
)
|
||
redteam_query_tensor = redteam_query_tensor.to(device)
|
||
redteam_key_tensor = redteam_key_tensor.to(device)
|
||
redteam_labels = torch.tensor(
|
||
[float(record["label"]) for record in redteam_records],
|
||
dtype=torch.float32,
|
||
device=device,
|
||
)
|
||
redteam_holdout_records = make_redteam_holdout(args.seed)
|
||
redteam_holdout_query_tensor, redteam_holdout_key_tensor = encode_records(
|
||
model,
|
||
tokenizer,
|
||
redteam_holdout_records,
|
||
batch_size=args.batch_size,
|
||
device=device,
|
||
)
|
||
redteam_holdout_query_tensor = redteam_holdout_query_tensor.to(device)
|
||
redteam_holdout_key_tensor = redteam_holdout_key_tensor.to(device)
|
||
redteam_holdout_labels = torch.tensor(
|
||
[float(record["label"]) for record in redteam_holdout_records],
|
||
dtype=torch.float32,
|
||
device=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),
|
||
)
|
||
update_indices = torch.tensor(
|
||
[rng.randrange(update_positive_query_tensor.shape[0]) for _ in range(args.batch_size)],
|
||
dtype=torch.long,
|
||
device=device,
|
||
)
|
||
update_positive_logits = model.text_retriever(
|
||
update_positive_query_tensor[update_indices],
|
||
update_positive_key_tensor[update_indices],
|
||
)
|
||
update_negative_logits = model.text_retriever(
|
||
update_negative_query_tensor[update_indices],
|
||
update_negative_key_tensor[update_indices],
|
||
)
|
||
update_pair_loss = F.binary_cross_entropy_with_logits(
|
||
torch.cat((update_positive_logits, update_negative_logits), dim=0),
|
||
torch.cat(
|
||
(
|
||
torch.ones_like(update_positive_logits),
|
||
torch.zeros_like(update_negative_logits),
|
||
),
|
||
dim=0,
|
||
),
|
||
)
|
||
redteam_indices = torch.tensor(
|
||
[rng.randrange(redteam_query_tensor.shape[0]) for _ in range(args.batch_size)],
|
||
dtype=torch.long,
|
||
device=device,
|
||
)
|
||
redteam_logits = model.text_retriever(
|
||
redteam_query_tensor[redteam_indices],
|
||
redteam_key_tensor[redteam_indices],
|
||
)
|
||
redteam_loss = F.binary_cross_entropy_with_logits(
|
||
redteam_logits,
|
||
redteam_labels[redteam_indices],
|
||
)
|
||
loss = (
|
||
contrastive_loss
|
||
+ 0.5 * classification_loss
|
||
+ 0.75 * hard_negative_loss
|
||
+ 1.25 * update_pair_loss
|
||
+ 1.00 * redteam_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,
|
||
)
|
||
update_positive_probability = torch.sigmoid(update_positive_logits).mean()
|
||
update_negative_probability = torch.sigmoid(update_negative_logits).mean()
|
||
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"update_pos={float(update_positive_probability):.3f} "
|
||
f"update_neg={float(update_negative_probability):.3f} "
|
||
f"redteam_loss={float(redteam_loss.detach()):.5f} "
|
||
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),
|
||
"update_positive_records": len(update_positive_records),
|
||
"update_negative_records": len(update_negative_records),
|
||
"redteam_records": len(redteam_records),
|
||
"redteam_holdout_records": len(redteam_holdout_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,
|
||
),
|
||
"update_pair_holdout": {
|
||
"positive_probability_mean": float(
|
||
torch.sigmoid(
|
||
model.text_retriever(update_positive_query_tensor, update_positive_key_tensor)
|
||
).mean().detach().cpu()
|
||
),
|
||
"negative_probability_mean": float(
|
||
torch.sigmoid(
|
||
model.text_retriever(update_negative_query_tensor, update_negative_key_tensor)
|
||
).mean().detach().cpu()
|
||
),
|
||
},
|
||
"redteam_holdout": {
|
||
"positive_probability_mean": float(
|
||
torch.sigmoid(
|
||
model.text_retriever(
|
||
redteam_holdout_query_tensor[redteam_holdout_labels >= 0.5],
|
||
redteam_holdout_key_tensor[redteam_holdout_labels >= 0.5],
|
||
)
|
||
).mean().detach().cpu()
|
||
),
|
||
"negative_probability_mean": float(
|
||
torch.sigmoid(
|
||
model.text_retriever(
|
||
redteam_holdout_query_tensor[redteam_holdout_labels < 0.5],
|
||
redteam_holdout_key_tensor[redteam_holdout_labels < 0.5],
|
||
)
|
||
).mean().detach().cpu()
|
||
),
|
||
},
|
||
}
|
||
(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()
|