✅ test: add comprehensive RULER benchmark test suite
- Add test_ruler.py supporting all 13 RULER tasks (NIAH, QA, CWE, FWE, VT) - Implement RULER official evaluation metrics (string_match_all/part) - Fix max_model_len to 32896 to prevent decode OOM on long inputs - Add ruler_benchmark_report.md with full test results (92.1% accuracy) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Zijie Tian
392
tests/test_ruler.py
Normal file
392
tests/test_ruler.py
Normal file
@@ -0,0 +1,392 @@
|
||||
"""
|
||||
RULER benchmark comprehensive test for LLM.
|
||||
|
||||
Tests multiple RULER tasks:
|
||||
- NIAH (Needle-In-A-Haystack): single, multikey, multiquery, multivalue
|
||||
- QA (Question Answering): qa_1, qa_2
|
||||
- CWE (Common Word Extraction)
|
||||
- FWE (Frequent Word Extraction)
|
||||
- VT (Variable Tracking)
|
||||
|
||||
Usage:
|
||||
# Test all datasets with 2 samples each (debug mode)
|
||||
python tests/test_ruler.py --enable-offload --num-samples 2
|
||||
|
||||
# Test specific datasets
|
||||
python tests/test_ruler.py --enable-offload --datasets niah_single_1,qa_1
|
||||
|
||||
# Test all samples in all datasets
|
||||
python tests/test_ruler.py --enable-offload
|
||||
"""
|
||||
|
||||
import os
|
||||
os.environ["NANOVLLM_LOG_LEVEL"] = "INFO"
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
import gc
|
||||
import time
|
||||
import torch
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Tuple, Optional
|
||||
|
||||
from nanovllm import LLM, SamplingParams
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Constants
|
||||
# ============================================================
|
||||
|
||||
DEFAULT_DATA_DIR = Path(__file__).parent / "data/ruler_32k"
|
||||
DEFAULT_MODEL = os.path.expanduser("~/models/Llama-3.1-8B-Instruct")
|
||||
# Note: max_model_len must be > max_input_len to leave room for output tokens
|
||||
# 32k benchmark has inputs up to 32760 tokens, so we need 32768 + 128 = 32896
|
||||
DEFAULT_MAX_MODEL_LEN = 32896
|
||||
DEFAULT_MAX_NEW_TOKENS = 128 # Larger for multi-value tasks
|
||||
|
||||
# Task categories for evaluation
|
||||
NIAH_TASKS = ["niah_single_1", "niah_single_2", "niah_single_3",
|
||||
"niah_multikey_1", "niah_multikey_2", "niah_multikey_3",
|
||||
"niah_multiquery", "niah_multivalue"]
|
||||
QA_TASKS = ["qa_1", "qa_2"]
|
||||
RECALL_TASKS = ["cwe", "fwe", "vt"]
|
||||
|
||||
ALL_TASKS = NIAH_TASKS + QA_TASKS + RECALL_TASKS
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Data Loading
|
||||
# ============================================================
|
||||
|
||||
def load_samples(filepath: Path, indices: Optional[List[int]] = None) -> List[dict]:
|
||||
"""Load samples from a JSONL file."""
|
||||
if not filepath.exists():
|
||||
raise FileNotFoundError(f"Data file not found: {filepath}")
|
||||
|
||||
samples = []
|
||||
with open(filepath) as f:
|
||||
for i, line in enumerate(f):
|
||||
if indices is None or i in indices:
|
||||
sample = json.loads(line)
|
||||
sample["_local_idx"] = i
|
||||
samples.append(sample)
|
||||
return samples
|
||||
|
||||
|
||||
def count_samples(filepath: Path) -> int:
|
||||
"""Count total samples in JSONL file."""
|
||||
with open(filepath) as f:
|
||||
return sum(1 for _ in f)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Evaluation Functions (Following RULER Official Metrics)
|
||||
# Ref: https://github.com/NVIDIA/RULER/blob/main/scripts/eval/synthetic/constants.py
|
||||
# ============================================================
|
||||
|
||||
def string_match_all(output_text: str, expected_list: List[str]) -> float:
|
||||
"""
|
||||
RULER official metric for NIAH, VT, CWE, FWE tasks.
|
||||
|
||||
Formula: sum([1.0 if r.lower() in pred.lower() else 0.0 for r in ref]) / len(ref)
|
||||
|
||||
Returns recall score (0.0 to 1.0): fraction of expected values found in output.
|
||||
"""
|
||||
output_clean = output_text.replace('<|im_end|>', '').replace('\r', ' ').replace('\n', ' ')
|
||||
output_lower = output_clean.lower()
|
||||
|
||||
if not expected_list:
|
||||
return 1.0
|
||||
|
||||
found = sum(1.0 if exp.strip().lower() in output_lower else 0.0 for exp in expected_list)
|
||||
return found / len(expected_list)
|
||||
|
||||
|
||||
def string_match_part(output_text: str, expected_list: List[str]) -> float:
|
||||
"""
|
||||
RULER official metric for QA tasks.
|
||||
|
||||
Formula: max([1.0 if r.lower() in pred.lower() else 0.0 for r in ref])
|
||||
|
||||
Returns 1.0 if ANY expected value is found, 0.0 otherwise.
|
||||
"""
|
||||
output_clean = output_text.replace('<|im_end|>', '').replace('\r', ' ').replace('\n', ' ')
|
||||
output_lower = output_clean.lower()
|
||||
|
||||
if not expected_list:
|
||||
return 1.0
|
||||
|
||||
return max(1.0 if exp.strip().lower() in output_lower else 0.0 for exp in expected_list)
|
||||
|
||||
|
||||
def evaluate_output(output_text: str, expected_outputs: List[str], task_name: str) -> Tuple[bool, float]:
|
||||
"""
|
||||
Evaluate model output using RULER official metrics.
|
||||
|
||||
- QA tasks: string_match_part (any match = full score)
|
||||
- All other tasks: string_match_all (recall-based score)
|
||||
|
||||
Returns (passed, score) where passed = score >= 0.5
|
||||
"""
|
||||
if task_name in QA_TASKS:
|
||||
score = string_match_part(output_text, expected_outputs)
|
||||
else:
|
||||
# NIAH, VT, CWE, FWE all use string_match_all
|
||||
score = string_match_all(output_text, expected_outputs)
|
||||
|
||||
passed = score >= 0.5 # Consider pass if score >= 50%
|
||||
return passed, score
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Test Runner
|
||||
# ============================================================
|
||||
|
||||
def run_task_test(
|
||||
llm: LLM,
|
||||
task_name: str,
|
||||
data_dir: Path,
|
||||
sample_indices: Optional[List[int]] = None,
|
||||
max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
|
||||
verbose: bool = True,
|
||||
) -> Dict:
|
||||
"""
|
||||
Run test for a single RULER task.
|
||||
|
||||
Returns dict with: task, correct, total, score, results
|
||||
"""
|
||||
data_file = data_dir / task_name / "validation.jsonl"
|
||||
samples = load_samples(data_file, sample_indices)
|
||||
|
||||
if verbose:
|
||||
print(f"\n Testing {task_name}: {len(samples)} samples")
|
||||
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0.1,
|
||||
max_tokens=max_new_tokens,
|
||||
)
|
||||
|
||||
correct = 0
|
||||
total_score = 0.0
|
||||
results = []
|
||||
|
||||
for sample in samples:
|
||||
idx = sample.get("index", sample["_local_idx"])
|
||||
prompt = sample["input"]
|
||||
expected = sample["outputs"]
|
||||
|
||||
# Generate
|
||||
outputs = llm.generate([prompt], sampling_params, use_tqdm=False)
|
||||
output_text = outputs[0]["text"]
|
||||
|
||||
# Evaluate
|
||||
passed, score = evaluate_output(output_text, expected, task_name)
|
||||
if passed:
|
||||
correct += 1
|
||||
total_score += score
|
||||
|
||||
results.append({
|
||||
"index": idx,
|
||||
"expected": expected,
|
||||
"output": output_text[:200],
|
||||
"passed": passed,
|
||||
"score": score,
|
||||
})
|
||||
|
||||
if verbose:
|
||||
status = "PASS" if passed else "FAIL"
|
||||
exp_preview = str(expected[0])[:30] if expected else "N/A"
|
||||
out_preview = output_text[:50].replace('\n', ' ')
|
||||
print(f" [{idx}] {status} (score={score:.2f}) exp={exp_preview}... out={out_preview}...")
|
||||
|
||||
avg_score = total_score / len(samples) if samples else 0.0
|
||||
|
||||
return {
|
||||
"task": task_name,
|
||||
"correct": correct,
|
||||
"total": len(samples),
|
||||
"accuracy": correct / len(samples) if samples else 0.0,
|
||||
"avg_score": avg_score,
|
||||
"results": results,
|
||||
}
|
||||
|
||||
|
||||
def run_ruler_benchmark(
|
||||
model_path: str,
|
||||
data_dir: Path,
|
||||
datasets: Optional[List[str]] = None,
|
||||
num_samples: Optional[int] = None,
|
||||
max_model_len: int = DEFAULT_MAX_MODEL_LEN,
|
||||
max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
|
||||
enable_cpu_offload: bool = False,
|
||||
num_gpu_blocks: int = 4,
|
||||
block_size: int = 1024,
|
||||
gpu_utilization: float = 0.9,
|
||||
enforce_eager: bool = True,
|
||||
verbose: bool = True,
|
||||
) -> Dict:
|
||||
"""
|
||||
Run RULER benchmark on multiple tasks.
|
||||
|
||||
Args:
|
||||
model_path: Path to the model
|
||||
data_dir: Directory containing task subdirectories
|
||||
datasets: List of task names to test (None = all)
|
||||
num_samples: Number of samples per task (None = all)
|
||||
...other LLM config params...
|
||||
|
||||
Returns:
|
||||
Dict with overall results and per-task results
|
||||
"""
|
||||
# Determine tasks to run
|
||||
if datasets is None:
|
||||
tasks = [t for t in ALL_TASKS if (data_dir / t / "validation.jsonl").exists()]
|
||||
else:
|
||||
tasks = datasets
|
||||
|
||||
# Sample indices
|
||||
sample_indices = list(range(num_samples)) if num_samples else None
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"RULER Benchmark")
|
||||
print(f"{'='*60}")
|
||||
print(f"Model: {model_path}")
|
||||
print(f"Data dir: {data_dir}")
|
||||
print(f"Tasks: {len(tasks)}")
|
||||
print(f"Samples per task: {num_samples if num_samples else 'all'}")
|
||||
print(f"CPU offload: {enable_cpu_offload}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
# Initialize LLM
|
||||
print("\nInitializing LLM...")
|
||||
llm_kwargs = {
|
||||
"max_model_len": max_model_len,
|
||||
"max_num_batched_tokens": max_model_len,
|
||||
"enforce_eager": enforce_eager,
|
||||
"gpu_memory_utilization": gpu_utilization,
|
||||
"kvcache_block_size": block_size,
|
||||
"enable_cpu_offload": enable_cpu_offload,
|
||||
}
|
||||
if enable_cpu_offload:
|
||||
llm_kwargs["num_gpu_blocks"] = num_gpu_blocks
|
||||
|
||||
llm = LLM(model_path, **llm_kwargs)
|
||||
|
||||
# Run tests
|
||||
start_time = time.time()
|
||||
task_results = []
|
||||
|
||||
for task_name in tasks:
|
||||
result = run_task_test(
|
||||
llm=llm,
|
||||
task_name=task_name,
|
||||
data_dir=data_dir,
|
||||
sample_indices=sample_indices,
|
||||
max_new_tokens=max_new_tokens,
|
||||
verbose=verbose,
|
||||
)
|
||||
task_results.append(result)
|
||||
|
||||
if verbose:
|
||||
print(f" -> {task_name}: {result['correct']}/{result['total']} "
|
||||
f"({result['accuracy']*100:.1f}%) avg_score={result['avg_score']:.3f}")
|
||||
|
||||
total_time = time.time() - start_time
|
||||
|
||||
# Cleanup
|
||||
del llm
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Aggregate results
|
||||
total_correct = sum(r["correct"] for r in task_results)
|
||||
total_samples = sum(r["total"] for r in task_results)
|
||||
overall_accuracy = total_correct / total_samples if total_samples > 0 else 0.0
|
||||
avg_score = sum(r["avg_score"] for r in task_results) / len(task_results) if task_results else 0.0
|
||||
|
||||
# Print summary
|
||||
print(f"\n{'='*60}")
|
||||
print(f"RULER Benchmark Results")
|
||||
print(f"{'='*60}")
|
||||
print(f"\n{'Task':<20} {'Correct':<10} {'Accuracy':<12} {'Avg Score':<12}")
|
||||
print(f"{'-'*54}")
|
||||
for r in task_results:
|
||||
print(f"{r['task']:<20} {r['correct']}/{r['total']:<7} {r['accuracy']*100:>6.1f}% {r['avg_score']:.3f}")
|
||||
print(f"{'-'*54}")
|
||||
print(f"{'TOTAL':<20} {total_correct}/{total_samples:<7} {overall_accuracy*100:>6.1f}% {avg_score:.3f}")
|
||||
print(f"\nTime: {total_time:.1f}s")
|
||||
print(f"{'='*60}\n")
|
||||
|
||||
return {
|
||||
"total_correct": total_correct,
|
||||
"total_samples": total_samples,
|
||||
"overall_accuracy": overall_accuracy,
|
||||
"avg_score": avg_score,
|
||||
"time": total_time,
|
||||
"task_results": task_results,
|
||||
}
|
||||
|
||||
|
||||
# ============================================================
|
||||
# CLI Entry Point
|
||||
# ============================================================
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="RULER benchmark comprehensive test",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
|
||||
parser.add_argument("--model", "-m", type=str, default=DEFAULT_MODEL,
|
||||
help=f"Path to model (default: {DEFAULT_MODEL})")
|
||||
parser.add_argument("--data-dir", type=str, default=str(DEFAULT_DATA_DIR),
|
||||
help=f"Path to data directory (default: {DEFAULT_DATA_DIR})")
|
||||
parser.add_argument("--datasets", type=str, default="",
|
||||
help="Comma-separated list of datasets to test (default: all)")
|
||||
parser.add_argument("--num-samples", type=int, default=0,
|
||||
help="Number of samples per dataset (default: 0 = all)")
|
||||
parser.add_argument("--max-model-len", type=int, default=DEFAULT_MAX_MODEL_LEN,
|
||||
help=f"Maximum model context length (default: {DEFAULT_MAX_MODEL_LEN})")
|
||||
parser.add_argument("--max-new-tokens", type=int, default=DEFAULT_MAX_NEW_TOKENS,
|
||||
help=f"Maximum tokens to generate (default: {DEFAULT_MAX_NEW_TOKENS})")
|
||||
parser.add_argument("--enable-offload", action="store_true",
|
||||
help="Enable CPU offload mode")
|
||||
parser.add_argument("--num-gpu-blocks", type=int, default=4,
|
||||
help="Number of GPU blocks for CPU offload (default: 4)")
|
||||
parser.add_argument("--block-size", type=int, default=1024,
|
||||
help="KV cache block size (default: 1024)")
|
||||
parser.add_argument("--gpu-utilization", type=float, default=0.9,
|
||||
help="GPU memory utilization (default: 0.9)")
|
||||
parser.add_argument("--use-cuda-graph", action="store_true",
|
||||
help="Enable CUDA graph")
|
||||
parser.add_argument("--quiet", "-q", action="store_true",
|
||||
help="Quiet mode")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Parse datasets
|
||||
datasets = args.datasets.split(",") if args.datasets else None
|
||||
num_samples = args.num_samples if args.num_samples > 0 else None
|
||||
|
||||
results = run_ruler_benchmark(
|
||||
model_path=os.path.expanduser(args.model),
|
||||
data_dir=Path(args.data_dir),
|
||||
datasets=datasets,
|
||||
num_samples=num_samples,
|
||||
max_model_len=args.max_model_len,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
enable_cpu_offload=args.enable_offload,
|
||||
num_gpu_blocks=args.num_gpu_blocks,
|
||||
block_size=args.block_size,
|
||||
gpu_utilization=args.gpu_utilization,
|
||||
enforce_eager=not args.use_cuda_graph,
|
||||
verbose=not args.quiet,
|
||||
)
|
||||
|
||||
# Exit code
|
||||
if results["overall_accuracy"] >= 0.5:
|
||||
print("test_ruler: PASSED")
|
||||
else:
|
||||
print(f"test_ruler: FAILED (accuracy={results['overall_accuracy']*100:.1f}%)")
|
||||
exit(1)
|
||||
Reference in New Issue
Block a user