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