Skip to content

Commit 916c803

Browse files
authored
Merge pull request #14 from ShaneIsley/claude/rlm-testing-framework-byfLB
Claude/rlm testing framework byf lb
2 parents 8aed4f6 + fa63a3a commit 916c803

12 files changed

Lines changed: 2799 additions & 0 deletions

File tree

‎benchmarks/__init__.py‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
"""
2+
RLM Benchmarks Framework
3+
4+
A minimal, extensible framework for evaluating Recursive Language Models
5+
against the benchmarks from the original RLM paper (arXiv:2512.24601).
6+
7+
Benchmark tasks:
8+
- S-NIAH: Single-Needle-in-a-Haystack (context-insensitive retrieval)
9+
- BrowseComp-Plus: Compositional multi-hop QA over document corpora
10+
- OOLONG: Semantic aggregation over long contexts
11+
- OOLONG-Pairs: Pairwise combinatorial aggregation
12+
13+
Usage:
14+
from benchmarks import BenchmarkRunner, OolongBenchmark
15+
16+
runner = BenchmarkRunner(backend="openai", model="gpt-5-mini")
17+
results = runner.run(OolongBenchmark(subset="toy_dnd", num_samples=10))
18+
"""
19+
20+
from benchmarks.base import Benchmark, BenchmarkResult, BenchmarkSample
21+
from benchmarks.metrics import Metrics
22+
from benchmarks.results import ExperimentConfig, ExperimentRecord, ResultsStore
23+
from benchmarks.runner import BenchmarkRunner, compare_methods
24+
from benchmarks.tasks import (
25+
BrowseCompPlusBenchmark,
26+
NIAHBenchmark,
27+
OolongBenchmark,
28+
OolongPairsBenchmark,
29+
)
30+
31+
__all__ = [
32+
# Base classes
33+
"Benchmark",
34+
"BenchmarkResult",
35+
"BenchmarkSample",
36+
# Runner
37+
"BenchmarkRunner",
38+
"compare_methods",
39+
# Metrics
40+
"Metrics",
41+
# Results storage
42+
"ResultsStore",
43+
"ExperimentConfig",
44+
"ExperimentRecord",
45+
# Benchmark tasks
46+
"NIAHBenchmark",
47+
"OolongBenchmark",
48+
"OolongPairsBenchmark",
49+
"BrowseCompPlusBenchmark",
50+
]

‎benchmarks/base.py‎

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,174 @@
1+
"""
2+
Base classes for RLM benchmarks.
3+
4+
Provides abstract interfaces that all benchmark implementations must follow,
5+
enabling consistent evaluation and extensibility.
6+
"""
7+
8+
from abc import ABC, abstractmethod
9+
from collections.abc import Iterator
10+
from dataclasses import dataclass, field
11+
from typing import Any
12+
13+
14+
@dataclass
15+
class BenchmarkSample:
16+
"""A single benchmark sample with context, question, and expected answer."""
17+
18+
id: str
19+
context: str
20+
question: str
21+
expected_answer: str | list[str]
22+
metadata: dict[str, Any] = field(default_factory=dict)
23+
24+
@property
25+
def context_tokens_approx(self) -> int:
26+
"""Approximate token count (rough estimate: ~4 chars per token)."""
27+
return len(self.context) // 4
28+
29+
30+
@dataclass
31+
class SampleResult:
32+
"""Result for a single benchmark sample."""
33+
34+
sample_id: str
35+
prediction: str
36+
expected: str | list[str]
37+
is_correct: bool
38+
metrics: dict[str, float] = field(default_factory=dict)
39+
iterations: int = 0
40+
total_tokens: int = 0
41+
execution_time_ms: float = 0.0
42+
error: str | None = None
43+
44+
45+
@dataclass
46+
class BenchmarkResult:
47+
"""Aggregated results for a complete benchmark run."""
48+
49+
benchmark_name: str
50+
method: str # "rlm", "direct", "summarize", "rag", etc.
51+
model: str
52+
sample_results: list[SampleResult] = field(default_factory=list)
53+
config: dict[str, Any] = field(default_factory=dict)
54+
55+
@property
56+
def accuracy(self) -> float:
57+
"""Fraction of correct predictions."""
58+
if not self.sample_results:
59+
return 0.0
60+
correct = sum(1 for r in self.sample_results if r.is_correct)
61+
return correct / len(self.sample_results)
62+
63+
@property
64+
def mean_f1(self) -> float:
65+
"""Mean F1 score across samples."""
66+
f1_scores = [r.metrics.get("f1", 0.0) for r in self.sample_results]
67+
return sum(f1_scores) / len(f1_scores) if f1_scores else 0.0
68+
69+
@property
70+
def total_tokens(self) -> int:
71+
"""Total tokens used across all samples."""
72+
return sum(r.total_tokens for r in self.sample_results)
73+
74+
@property
75+
def mean_iterations(self) -> float:
76+
"""Mean number of RLM iterations."""
77+
iters = [r.iterations for r in self.sample_results]
78+
return sum(iters) / len(iters) if iters else 0.0
79+
80+
@property
81+
def mean_execution_time_ms(self) -> float:
82+
"""Mean execution time per sample in milliseconds."""
83+
times = [r.execution_time_ms for r in self.sample_results]
84+
return sum(times) / len(times) if times else 0.0
85+
86+
@property
87+
def error_rate(self) -> float:
88+
"""Fraction of samples that resulted in errors."""
89+
errors = sum(1 for r in self.sample_results if r.error is not None)
90+
return errors / len(self.sample_results) if self.sample_results else 0.0
91+
92+
def summary(self) -> dict[str, Any]:
93+
"""Return summary statistics."""
94+
return {
95+
"benchmark": self.benchmark_name,
96+
"method": self.method,
97+
"model": self.model,
98+
"num_samples": len(self.sample_results),
99+
"accuracy": self.accuracy,
100+
"mean_f1": self.mean_f1,
101+
"total_tokens": self.total_tokens,
102+
"mean_iterations": self.mean_iterations,
103+
"mean_execution_time_ms": self.mean_execution_time_ms,
104+
"error_rate": self.error_rate,
105+
}
106+
107+
108+
class Benchmark(ABC):
109+
"""Abstract base class for all benchmarks.
110+
111+
To create a new benchmark:
112+
1. Subclass Benchmark
113+
2. Implement name, description, load_samples(), and evaluate()
114+
3. Optionally override default_metrics() for custom evaluation
115+
116+
Example:
117+
class MyBenchmark(Benchmark):
118+
@property
119+
def name(self) -> str:
120+
return "my-benchmark"
121+
122+
def load_samples(self, num_samples: int | None = None) -> Iterator[BenchmarkSample]:
123+
# Load from dataset, file, or generate samples
124+
yield BenchmarkSample(...)
125+
126+
def evaluate(self, prediction: str, expected: str | list[str]) -> dict[str, float]:
127+
# Return metrics dict with at least "correct" and "f1"
128+
return {"correct": 1.0 if prediction == expected else 0.0, "f1": ...}
129+
"""
130+
131+
@property
132+
@abstractmethod
133+
def name(self) -> str:
134+
"""Unique identifier for this benchmark."""
135+
pass
136+
137+
@property
138+
def description(self) -> str:
139+
"""Human-readable description of the benchmark."""
140+
return ""
141+
142+
@abstractmethod
143+
def load_samples(
144+
self, num_samples: int | None = None, seed: int | None = None
145+
) -> Iterator[BenchmarkSample]:
146+
"""Load benchmark samples.
147+
148+
Args:
149+
num_samples: Maximum number of samples to load. None for all.
150+
seed: Random seed for reproducible sampling.
151+
152+
Yields:
153+
BenchmarkSample instances.
154+
"""
155+
pass
156+
157+
@abstractmethod
158+
def evaluate(self, prediction: str, expected: str | list[str]) -> dict[str, float]:
159+
"""Evaluate a prediction against expected answer(s).
160+
161+
Args:
162+
prediction: Model's prediction string.
163+
expected: Expected answer(s). Can be a single string or list of valid answers.
164+
165+
Returns:
166+
Dictionary with at least:
167+
- "correct": 1.0 if correct, 0.0 otherwise
168+
- "f1": F1 score between prediction and expected
169+
"""
170+
pass
171+
172+
def default_metrics(self) -> list[str]:
173+
"""List of metric names this benchmark produces."""
174+
return ["correct", "f1"]

0 commit comments

Comments
 (0)