"""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()