Add Natural Memory architecture and tooling
This commit is contained in:
WpyQwq
@@ -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()
|
||||
Reference in New Issue
Block a user