302 lines
11 KiB
Python
302 lines
11 KiB
Python
"""Compare original Qwen and Native Memory on general regression tasks."""
|
|
|
|
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
|
|
|
|
import torch
|
|
|
|
from .qwen_integration import load_memory_config, load_qwen_base, load_qwen_dynamic, load_tokenizer
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
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("--data", default="dynamic_memory_lab/data/comprehensive_general.jsonl")
|
|
parser.add_argument("--output", default="dynamic_memory_lab/comprehensive_benchmark_native_v3.json")
|
|
parser.add_argument("--max-new-tokens", type=int, default=32)
|
|
parser.add_argument("--perf-repeats", type=int, default=3)
|
|
parser.add_argument("--no-4bit", action="store_true")
|
|
parser.add_argument("--category-drop-limit", type=float, default=0.10)
|
|
parser.add_argument("--overall-drop-limit", type=float, default=0.05)
|
|
return parser.parse_args()
|
|
|
|
|
|
def load_cases(path: str | Path) -> list[dict[str, Any]]:
|
|
cases = []
|
|
for line in Path(path).read_text(encoding="utf-8").splitlines():
|
|
if line.strip():
|
|
cases.append(json.loads(line))
|
|
if not cases:
|
|
raise ValueError(f"no benchmark cases found in {path}")
|
|
return cases
|
|
|
|
|
|
def normalize(text: str) -> str:
|
|
return re.sub(r"[\s`*_#，。！？、；：,.!?;:'\"（）()\[\]{}]", "", text).lower()
|
|
|
|
|
|
def contains_answer(text: str, acceptable: list[str]) -> bool:
|
|
normalized = normalize(text)
|
|
for answer in acceptable:
|
|
expected = normalize(str(answer))
|
|
if not expected:
|
|
continue
|
|
if expected.isdigit() and len(expected) == 1:
|
|
if re.search(rf"(?<!\d){re.escape(expected)}(?!\d)", normalized):
|
|
return True
|
|
elif expected in normalized:
|
|
return True
|
|
return False
|
|
|
|
|
|
def prompt_inputs(tokenizer: Any, prompt: str, device: torch.device) -> dict[str, torch.Tensor]:
|
|
encoded = tokenizer.apply_chat_template(
|
|
[{"role": "user", "content": prompt}],
|
|
tokenize=True,
|
|
add_generation_prompt=True,
|
|
enable_thinking=False,
|
|
)
|
|
input_ids = encoded["input_ids"] if isinstance(encoded, dict) or hasattr(encoded, "__getitem__") else encoded
|
|
if isinstance(input_ids, torch.Tensor):
|
|
if input_ids.ndim == 1:
|
|
input_ids = input_ids.unsqueeze(0)
|
|
input_ids = input_ids.to(device)
|
|
else:
|
|
if input_ids and isinstance(input_ids[0], list):
|
|
input_ids = input_ids[0]
|
|
input_ids = torch.tensor([input_ids], dtype=torch.long, device=device)
|
|
return {"input_ids": input_ids, "attention_mask": torch.ones_like(input_ids)}
|
|
|
|
|
|
def _decode_generation(tokenizer: Any, output: torch.Tensor, prompt: dict[str, torch.Tensor]) -> str:
|
|
start = prompt["input_ids"].shape[1]
|
|
return tokenizer.decode(output[0, start:].detach().cpu().tolist(), skip_special_tokens=True).strip()
|
|
|
|
|
|
def evaluate_model(
|
|
model: Any,
|
|
tokenizer: Any,
|
|
cases: list[dict[str, Any]],
|
|
*,
|
|
adapted: bool,
|
|
max_new_tokens: int,
|
|
perf_repeats: int,
|
|
) -> dict[str, Any]:
|
|
device = model._find_layer_device() if adapted else model.get_input_embeddings().weight.device
|
|
rows: list[dict[str, Any]] = []
|
|
category_values: dict[str, list[float]] = defaultdict(list)
|
|
category_examples: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
|
started = time.perf_counter()
|
|
for case in cases:
|
|
if adapted:
|
|
model.reset_memory()
|
|
prompt = prompt_inputs(tokenizer, str(case["prompt"]), device)
|
|
with torch.inference_mode():
|
|
if adapted:
|
|
output = model.generate(
|
|
**prompt,
|
|
max_new_tokens=max_new_tokens,
|
|
do_sample=False,
|
|
update_memory=False,
|
|
use_cache=True,
|
|
pad_token_id=tokenizer.pad_token_id,
|
|
)
|
|
else:
|
|
output = model.generate(
|
|
**prompt,
|
|
max_new_tokens=max_new_tokens,
|
|
do_sample=False,
|
|
use_cache=True,
|
|
pad_token_id=tokenizer.pad_token_id,
|
|
)
|
|
generated = _decode_generation(tokenizer, output, prompt)
|
|
passed = contains_answer(generated, list(case["acceptable"]))
|
|
category = str(case["category"])
|
|
category_values[category].append(float(passed))
|
|
if len(category_examples[category]) < 3:
|
|
category_examples[category].append(
|
|
{
|
|
"id": case["id"],
|
|
"prompt": case["prompt"],
|
|
"acceptable": case["acceptable"],
|
|
"generated": generated,
|
|
"passed": passed,
|
|
}
|
|
)
|
|
rows.append(
|
|
{
|
|
"id": case["id"],
|
|
"category": category,
|
|
"acceptable": case["acceptable"],
|
|
"generated": generated,
|
|
"passed": passed,
|
|
}
|
|
)
|
|
elapsed = time.perf_counter() - started
|
|
|
|
perf_case = cases[0]
|
|
if adapted:
|
|
model.reset_memory()
|
|
perf_prompt = prompt_inputs(tokenizer, str(perf_case["prompt"]), device)
|
|
if device.type == "cuda":
|
|
torch.cuda.synchronize(device)
|
|
perf_start = time.perf_counter()
|
|
generated_tokens = 0
|
|
for _ in range(max(1, perf_repeats)):
|
|
if adapted:
|
|
model.reset_memory()
|
|
with torch.inference_mode():
|
|
if adapted:
|
|
perf_output = model.generate(
|
|
**perf_prompt,
|
|
max_new_tokens=max_new_tokens,
|
|
do_sample=False,
|
|
update_memory=False,
|
|
use_cache=True,
|
|
pad_token_id=tokenizer.pad_token_id,
|
|
)
|
|
else:
|
|
perf_output = model.generate(
|
|
**perf_prompt,
|
|
max_new_tokens=max_new_tokens,
|
|
do_sample=False,
|
|
use_cache=True,
|
|
pad_token_id=tokenizer.pad_token_id,
|
|
)
|
|
generated_tokens += int(perf_output.shape[1] - perf_prompt["input_ids"].shape[1])
|
|
if device.type == "cuda":
|
|
torch.cuda.synchronize(device)
|
|
perf_elapsed = time.perf_counter() - perf_start
|
|
score = sum(sum(values) for values in category_values.values()) / len(cases)
|
|
return {
|
|
"device": str(device),
|
|
"cases": len(cases),
|
|
"elapsed_seconds": elapsed,
|
|
"overall_score": score,
|
|
"categories": {
|
|
category: {
|
|
"count": len(values),
|
|
"score": sum(values) / len(values),
|
|
"examples": category_examples[category],
|
|
}
|
|
for category, values in sorted(category_values.items())
|
|
},
|
|
"performance": {
|
|
"repeats": max(1, perf_repeats),
|
|
"max_new_tokens": max_new_tokens,
|
|
"tokens_per_second": generated_tokens / max(perf_elapsed, 1e-9),
|
|
"seconds_per_run": perf_elapsed / max(1, perf_repeats),
|
|
},
|
|
"rows": rows,
|
|
}
|
|
|
|
|
|
def release(model: Any) -> None:
|
|
del model
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def main() -> None:
|
|
if hasattr(sys.stdout, "reconfigure"):
|
|
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
|
|
args = parse_args()
|
|
cases = load_cases(args.data)
|
|
tokenizer = load_tokenizer(args.model_path)
|
|
use_4bit = not args.no_4bit
|
|
report: dict[str, Any] = {
|
|
"model_path": str(Path(args.model_path).resolve()),
|
|
"adapter": str(Path(args.adapter).resolve()),
|
|
"data": str(Path(args.data).resolve()),
|
|
"quantization": "4bit_nf4" if use_4bit else "none",
|
|
"cases": len(cases),
|
|
}
|
|
|
|
print("loading baseline")
|
|
baseline = load_qwen_base(args.model_path, load_in_4bit=use_4bit)
|
|
baseline.eval()
|
|
baseline_device = baseline.get_input_embeddings().weight.device
|
|
if baseline_device.type == "cuda":
|
|
torch.cuda.reset_peak_memory_stats(baseline_device)
|
|
report["baseline"] = evaluate_model(
|
|
baseline,
|
|
tokenizer,
|
|
cases,
|
|
adapted=False,
|
|
max_new_tokens=args.max_new_tokens,
|
|
perf_repeats=args.perf_repeats,
|
|
)
|
|
if baseline_device.type == "cuda":
|
|
report["baseline"]["peak_memory_gb"] = torch.cuda.max_memory_allocated(baseline_device) / 1024**3
|
|
release(baseline)
|
|
|
|
config = load_memory_config(args.adapter)
|
|
config.persistent_memory = False
|
|
print(f"loading native adapter mode={config.mode} layers={config.layer_indices}")
|
|
adapted = load_qwen_dynamic(
|
|
args.model_path,
|
|
memory_config=config,
|
|
load_in_4bit=use_4bit,
|
|
)
|
|
adapted.load_memory_adapter(args.adapter)
|
|
adapted.eval()
|
|
adapted_device = adapted._find_layer_device()
|
|
if adapted_device.type == "cuda":
|
|
torch.cuda.reset_peak_memory_stats(adapted_device)
|
|
report["native_memory"] = evaluate_model(
|
|
adapted,
|
|
tokenizer,
|
|
cases,
|
|
adapted=True,
|
|
max_new_tokens=args.max_new_tokens,
|
|
perf_repeats=args.perf_repeats,
|
|
)
|
|
if adapted_device.type == "cuda":
|
|
report["native_memory"]["peak_memory_gb"] = torch.cuda.max_memory_allocated(adapted_device) / 1024**3
|
|
|
|
baseline_score = report["baseline"]["overall_score"]
|
|
native_score = report["native_memory"]["overall_score"]
|
|
baseline_categories = report["baseline"]["categories"]
|
|
native_categories = report["native_memory"]["categories"]
|
|
category_deltas = {
|
|
category: native_categories[category]["score"] - baseline_categories[category]["score"]
|
|
for category in baseline_categories.keys() & native_categories.keys()
|
|
}
|
|
report["regression"] = {
|
|
"overall_delta": native_score - baseline_score,
|
|
"category_deltas": category_deltas,
|
|
"overall_drop_limit": args.overall_drop_limit,
|
|
"category_drop_limit": args.category_drop_limit,
|
|
"overall_regression_alert": native_score < baseline_score - args.overall_drop_limit,
|
|
"category_regression_alerts": {
|
|
category: delta < -args.category_drop_limit for category, delta in category_deltas.items()
|
|
},
|
|
"pass": native_score >= baseline_score - args.overall_drop_limit
|
|
and all(delta >= -args.category_drop_limit for delta in category_deltas.values()),
|
|
}
|
|
legacy_path = Path(args.adapter).parent / "benchmark_qwen_native_v3_eval.json"
|
|
if legacy_path.exists():
|
|
report["same_memory_benchmark"] = json.loads(legacy_path.read_text(encoding="utf-8"))
|
|
native_eval_path = Path(args.adapter) / "native_eval_report.json"
|
|
if native_eval_path.exists():
|
|
report["native_memory_holdout"] = json.loads(native_eval_path.read_text(encoding="utf-8"))
|
|
output = Path(args.output)
|
|
output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
print(json.dumps({key: report[key] for key in ("baseline", "native_memory", "regression")}, ensure_ascii=False, indent=2))
|
|
print(f"saved={output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|