320 lines
12 KiB
Python
320 lines
12 KiB
Python
"""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()
|