604 lines
20 KiB
Python
604 lines
20 KiB
Python
"""
|
|
Honcho OOLONG benchmark runner.
|
|
|
|
Evaluates long-context reasoning and aggregation on:
|
|
- OOLONG-synth: synthetic ICL aggregation tasks
|
|
- OOLONG-real: D&D transcript aggregation tasks
|
|
"""
|
|
|
|
import argparse
|
|
import sys
|
|
import time
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any, TypedDict, cast
|
|
|
|
from dotenv import load_dotenv
|
|
from honcho.api_types import MessageCreateParams
|
|
from honcho.session import Session, SessionPeerConfig
|
|
|
|
from src.config import settings
|
|
|
|
from .oolong_common import (
|
|
calculate_context_length,
|
|
calculate_task_statistics,
|
|
calculate_timing_statistics,
|
|
format_duration,
|
|
load_oolong_real_dataset,
|
|
load_oolong_synth_dataset,
|
|
parse_real_answer,
|
|
parse_real_context_messages,
|
|
parse_synth_answer,
|
|
parse_synth_context_messages,
|
|
score_real_response,
|
|
score_synth_response,
|
|
write_json_summary,
|
|
)
|
|
from .runner_common import (
|
|
BaseRunner,
|
|
ItemContext,
|
|
RunnerConfig,
|
|
add_common_arguments,
|
|
validate_common_arguments,
|
|
)
|
|
|
|
load_dotenv()
|
|
|
|
CONTEXT_SIZE_MAP: dict[str, int] = {
|
|
"1K": 1024,
|
|
"2K": 2 * 1024,
|
|
"4K": 4 * 1024,
|
|
"8K": 8 * 1024,
|
|
"16K": 16 * 1024,
|
|
"32K": 32 * 1024,
|
|
"64K": 64 * 1024,
|
|
"128K": 128 * 1024,
|
|
"256K": 256 * 1024,
|
|
"512K": 512 * 1024,
|
|
"1M": 1024 * 1024,
|
|
"2M": 2 * 1024 * 1024,
|
|
"4M": 4 * 1024 * 1024,
|
|
}
|
|
|
|
|
|
def parse_context_size(size_str: str) -> int:
|
|
"""Parse a context-size string into an exact token count."""
|
|
normalized = size_str.strip().upper()
|
|
|
|
if normalized.isdigit():
|
|
value = int(normalized)
|
|
if value <= 0:
|
|
raise ValueError("Context size must be positive")
|
|
return value
|
|
|
|
if normalized in CONTEXT_SIZE_MAP:
|
|
return CONTEXT_SIZE_MAP[normalized]
|
|
|
|
valid_sizes = ", ".join(CONTEXT_SIZE_MAP)
|
|
raise ValueError(
|
|
f"Invalid context size '{size_str}'. Use one of [{valid_sizes}] or a positive integer token count."
|
|
)
|
|
|
|
|
|
class QueryResult(TypedDict):
|
|
"""Query execution result for one OOLONG example."""
|
|
|
|
question: str
|
|
expected_answer: str
|
|
actual_response: str
|
|
score: float
|
|
context_length_tokens: int
|
|
|
|
|
|
class TestResult(TypedDict):
|
|
"""Single OOLONG example result."""
|
|
|
|
question_id: str
|
|
context_window_id: str
|
|
task_group: str
|
|
dataset: str
|
|
answer_type: str
|
|
passed: bool
|
|
score: float
|
|
error: str | None
|
|
start_time: float
|
|
end_time: float
|
|
duration_seconds: float
|
|
query_executed: QueryResult | None
|
|
output_lines: list[str]
|
|
|
|
|
|
class OolongRunner(BaseRunner[TestResult]):
|
|
"""Execute OOLONG benchmark examples through the shared runner framework."""
|
|
|
|
variant: str
|
|
data_dir: Path
|
|
split: str
|
|
merge_sessions: bool
|
|
max_examples: int | None
|
|
min_context_len: int | None
|
|
max_context_len: int | None
|
|
context_window_id: str | None
|
|
use_labels: bool
|
|
|
|
def __init__(
|
|
self,
|
|
config: RunnerConfig,
|
|
variant: str,
|
|
data_dir: Path,
|
|
split: str,
|
|
merge_sessions: bool,
|
|
max_examples: int | None = None,
|
|
min_context_len: int | None = None,
|
|
max_context_len: int | None = None,
|
|
context_window_id: str | None = None,
|
|
use_labels: bool = False,
|
|
):
|
|
self.variant = variant
|
|
self.data_dir = data_dir
|
|
self.split = split
|
|
self.merge_sessions = merge_sessions
|
|
self.max_examples = max_examples
|
|
self.min_context_len = min_context_len
|
|
self.max_context_len = max_context_len
|
|
self.context_window_id = context_window_id
|
|
self.use_labels = use_labels
|
|
super().__init__(config)
|
|
|
|
def get_metrics_prefix(self) -> str:
|
|
return "oolong"
|
|
|
|
def load_items(self) -> list[Any]:
|
|
if self.variant == "synth":
|
|
dataset = load_oolong_synth_dataset(
|
|
split=self.split,
|
|
data_dir=self.data_dir,
|
|
max_context_len=self.max_context_len,
|
|
min_context_len=self.min_context_len,
|
|
max_examples=self.max_examples,
|
|
context_window_id=self.context_window_id,
|
|
)
|
|
else:
|
|
dataset = load_oolong_real_dataset(
|
|
split=self.split,
|
|
data_dir=self.data_dir,
|
|
max_context_len=self.max_context_len,
|
|
min_context_len=self.min_context_len,
|
|
max_examples=self.max_examples,
|
|
context_window_id=self.context_window_id,
|
|
)
|
|
return [dataset[i] for i in range(len(dataset))]
|
|
|
|
def get_workspace_id(self, item: Any) -> str:
|
|
return f"oolong_{self.variant}_{item['id']}"
|
|
|
|
def get_session_id(self, item: Any, workspace_id: str) -> str:
|
|
return f"{workspace_id}_session"
|
|
|
|
async def setup_peers(self, ctx: ItemContext, item: Any) -> None:
|
|
ctx.peers["user"] = await ctx.honcho_client.aio.peer(id="user")
|
|
|
|
async def setup_session(self, ctx: ItemContext, item: Any) -> None:
|
|
if not self.merge_sessions:
|
|
ctx.session = None
|
|
return
|
|
|
|
user_peer = ctx.peers["user"]
|
|
ctx.session = await ctx.honcho_client.aio.session(
|
|
id=ctx.session_id, configuration=self._get_session_configuration()
|
|
)
|
|
await ctx.session.aio.add_peers(
|
|
[(user_peer, SessionPeerConfig(observe_me=True, observe_others=False))]
|
|
)
|
|
|
|
async def _add_messages_to_session(
|
|
self, session: Session, user_peer: Any, messages: list[dict[str, Any]]
|
|
) -> None:
|
|
honcho_messages: list[MessageCreateParams] = []
|
|
for msg in messages:
|
|
honcho_messages.append(
|
|
user_peer.message(
|
|
content=msg["content"],
|
|
metadata=msg.get("metadata"),
|
|
)
|
|
)
|
|
|
|
for i in range(0, len(honcho_messages), 100):
|
|
batch = honcho_messages[i : i + 100]
|
|
await session.aio.add_messages(batch)
|
|
|
|
async def ingest_messages(self, ctx: ItemContext, item: Any) -> int:
|
|
context_text = item["context_window_text"]
|
|
if self.variant == "synth":
|
|
if self.use_labels:
|
|
context_text = item.get("context_window_text_with_labels", context_text)
|
|
messages = parse_synth_context_messages(context_text)
|
|
else:
|
|
messages = parse_real_context_messages(context_text)
|
|
|
|
user_peer = ctx.peers["user"]
|
|
|
|
if self.merge_sessions:
|
|
if ctx.session is None:
|
|
raise ValueError("Merged mode requires a configured session")
|
|
await self._add_messages_to_session(ctx.session, user_peer, messages)
|
|
return len(messages)
|
|
|
|
chunk_size = 200
|
|
session_ids: list[str] = []
|
|
for idx, start in enumerate(range(0, len(messages), chunk_size)):
|
|
chunk = messages[start : start + chunk_size]
|
|
session_id = f"{ctx.workspace_id}_session_{idx + 1}"
|
|
session = await ctx.honcho_client.aio.session(
|
|
id=session_id, configuration=self._get_session_configuration()
|
|
)
|
|
await session.aio.add_peers(
|
|
[(user_peer, SessionPeerConfig(observe_me=True, observe_others=False))]
|
|
)
|
|
await self._add_messages_to_session(session, user_peer, chunk)
|
|
session_ids.append(session_id)
|
|
|
|
ctx.peers["_session_ids"] = session_ids
|
|
return len(messages)
|
|
|
|
def get_dream_observers(self, item: Any) -> list[str]:
|
|
return ["user"]
|
|
|
|
def get_dream_session_ids(self, ctx: ItemContext, _item: Any) -> list[str]:
|
|
if self.merge_sessions:
|
|
return [ctx.session_id]
|
|
|
|
session_ids = ctx.peers.get("_session_ids")
|
|
if not isinstance(session_ids, list) or not session_ids:
|
|
raise ValueError(
|
|
"Non-merged OOLONG mode requires at least one chunk session ID for dreams"
|
|
)
|
|
|
|
session_ids_typed = cast(list[object], session_ids)
|
|
cleaned_session_ids: list[str] = []
|
|
for maybe_session_id in session_ids_typed:
|
|
if isinstance(maybe_session_id, str) and maybe_session_id:
|
|
cleaned_session_ids.append(maybe_session_id)
|
|
if not cleaned_session_ids:
|
|
raise ValueError(
|
|
"Non-merged OOLONG mode has no valid chunk session IDs for dreams"
|
|
)
|
|
return cleaned_session_ids
|
|
|
|
async def execute_questions(self, ctx: ItemContext, item: Any) -> TestResult:
|
|
start_time = time.time()
|
|
question_id = item["id"]
|
|
context_window_id = item["context_window_id"]
|
|
question = item["question"]
|
|
answer_str = item["answer"]
|
|
|
|
if self.variant == "synth":
|
|
task_group = item.get("task_group", "unknown")
|
|
dataset_name = item.get("dataset", "oolong-synth")
|
|
answer_type = item.get("answer_type", "unknown")
|
|
gold_answer = parse_synth_answer(answer_str)
|
|
else:
|
|
task_group = item.get("question_type", "unknown")
|
|
dataset_name = "oolong-real"
|
|
answer_type = "varied"
|
|
gold_answer = parse_real_answer(answer_str)
|
|
|
|
context_text = item["context_window_text"]
|
|
if self.variant == "synth" and self.use_labels:
|
|
context_text = item.get("context_window_text_with_labels", context_text)
|
|
context_length = calculate_context_length(context_text)
|
|
|
|
result: TestResult = {
|
|
"question_id": question_id,
|
|
"context_window_id": context_window_id,
|
|
"task_group": task_group,
|
|
"dataset": dataset_name,
|
|
"answer_type": answer_type,
|
|
"passed": False,
|
|
"score": 0.0,
|
|
"error": None,
|
|
"start_time": start_time,
|
|
"end_time": 0.0,
|
|
"duration_seconds": 0.0,
|
|
"query_executed": None,
|
|
"output_lines": [],
|
|
}
|
|
|
|
user_peer = ctx.peers["user"]
|
|
try:
|
|
chat_kwargs: dict[str, Any] = {}
|
|
if self.config.reasoning_level:
|
|
chat_kwargs["reasoning_level"] = self.config.reasoning_level
|
|
if self.merge_sessions and ctx.session is not None:
|
|
chat_kwargs["session"] = ctx.session
|
|
|
|
response = await user_peer.aio.chat(question, **chat_kwargs)
|
|
actual_response = response if isinstance(response, str) else ""
|
|
|
|
if self.variant == "synth":
|
|
score = score_synth_response(gold_answer, actual_response, answer_type)
|
|
else:
|
|
score = score_real_response(gold_answer, actual_response)
|
|
|
|
result["query_executed"] = QueryResult(
|
|
question=question,
|
|
expected_answer=str(gold_answer),
|
|
actual_response=actual_response,
|
|
score=score,
|
|
context_length_tokens=context_length,
|
|
)
|
|
result["score"] = score
|
|
result["passed"] = score >= 0.99
|
|
result["output_lines"] = [
|
|
f"Question: {question}",
|
|
f"Expected: {gold_answer}",
|
|
f"Score: {score:.3f}",
|
|
]
|
|
except Exception as e:
|
|
result["error"] = str(e)
|
|
result["query_executed"] = QueryResult(
|
|
question=question,
|
|
expected_answer=str(gold_answer),
|
|
actual_response=f"ERROR: {e}",
|
|
score=0.0,
|
|
context_length_tokens=context_length,
|
|
)
|
|
|
|
result["end_time"] = time.time()
|
|
result["duration_seconds"] = result["end_time"] - result["start_time"]
|
|
return result
|
|
|
|
def print_summary(self, results: list[TestResult], total_duration: float) -> None:
|
|
total_examples = len(results)
|
|
perfect_scores = sum(1 for r in results if r["score"] >= 0.99)
|
|
average_score = (
|
|
sum(result["score"] for result in results) / total_examples
|
|
if total_examples
|
|
else 0.0
|
|
)
|
|
|
|
print(f"\n{'=' * 80}")
|
|
print(f"OOLONG-{self.variant.upper()} BENCHMARK SUMMARY")
|
|
print(f"{'=' * 80}")
|
|
print(f"Total examples: {total_examples}")
|
|
print(f"Average score: {average_score:.3f}")
|
|
perfect_rate = (
|
|
(perfect_scores / total_examples) * 100 if total_examples else 0.0
|
|
)
|
|
print(f"Perfect scores (>=0.99): {perfect_scores} ({perfect_rate:.1f}%)")
|
|
print(f"Total test time: {format_duration(total_duration)}")
|
|
|
|
task_stats = calculate_task_statistics(results)
|
|
if task_stats:
|
|
print("\nTask group statistics:")
|
|
for task_name, stats in sorted(task_stats.items()):
|
|
print(
|
|
f" {task_name}: avg={stats['average_score']:.3f}, perfect={stats['perfect_score_rate']:.1f}% ({stats['total']})"
|
|
)
|
|
print(f"{'=' * 80}")
|
|
|
|
def generate_output(self, results: list[TestResult], total_duration: float) -> None:
|
|
total_examples = len(results)
|
|
perfect_scores = sum(1 for r in results if r["score"] >= 0.99)
|
|
average_score = (
|
|
sum(result["score"] for result in results) / total_examples
|
|
if total_examples
|
|
else 0.0
|
|
)
|
|
task_stats = calculate_task_statistics(results)
|
|
timing_stats = calculate_timing_statistics(results, total_duration)
|
|
|
|
summary: dict[str, Any] = {
|
|
"metadata": {
|
|
"benchmark": "oolong",
|
|
"variant": self.variant,
|
|
"split": self.split,
|
|
"data_dir": str(self.data_dir),
|
|
"execution_timestamp": datetime.now().isoformat(),
|
|
"runner_version": "2.0.0",
|
|
"base_api_port": self.config.base_api_port,
|
|
"pool_size": self.config.pool_size,
|
|
"timeout_seconds": self.config.timeout_seconds,
|
|
"merge_sessions": self.merge_sessions,
|
|
"labels": self.use_labels,
|
|
"reasoning_level": self.config.reasoning_level,
|
|
"deriver_settings": settings.DERIVER.model_dump(),
|
|
"dialectic_settings": settings.DIALECTIC.model_dump(),
|
|
"dream_settings": settings.DREAM.model_dump(),
|
|
},
|
|
"summary_statistics": {
|
|
"total_examples": total_examples,
|
|
"perfect_scores": perfect_scores,
|
|
"perfect_score_rate": perfect_scores / total_examples
|
|
if total_examples
|
|
else 0.0,
|
|
"average_score": average_score,
|
|
"statistics_by_task_group": task_stats,
|
|
},
|
|
"timing": timing_stats,
|
|
"detailed_results": [
|
|
{
|
|
"question_id": result["question_id"],
|
|
"context_window_id": result["context_window_id"],
|
|
"task_group": result["task_group"],
|
|
"dataset": result["dataset"],
|
|
"answer_type": result["answer_type"],
|
|
"score": result["score"],
|
|
"passed": result["passed"],
|
|
"error": result["error"],
|
|
"duration_seconds": result["duration_seconds"],
|
|
"query_executed": result["query_executed"],
|
|
}
|
|
for result in results
|
|
],
|
|
}
|
|
|
|
if self.config.json_output:
|
|
output_file = self.config.json_output
|
|
else:
|
|
output_file = Path(
|
|
f"tests/bench/eval_results/oolong_{self.variant}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
|
|
)
|
|
|
|
write_json_summary(summary, output_file)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(
|
|
description="Run OOLONG benchmark tests against a Honcho instance",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
Examples:
|
|
%(prog)s --variant synth --data-dir /datasets/oolong-synth
|
|
%(prog)s --variant real --data-dir /datasets/oolong-real --split validation
|
|
%(prog)s --variant synth --data-dir /datasets/oolong-synth --context-size 16K
|
|
%(prog)s --variant synth --data-dir /datasets/oolong-synth --no-merge-sessions
|
|
""",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--variant",
|
|
type=str,
|
|
default="synth",
|
|
choices=["synth", "real"],
|
|
help="Which OOLONG variant to run (default: synth)",
|
|
)
|
|
parser.add_argument(
|
|
"--data-dir",
|
|
type=Path,
|
|
required=True,
|
|
help="Path to the dataset directory",
|
|
)
|
|
parser.add_argument(
|
|
"--split",
|
|
type=str,
|
|
default="test",
|
|
choices=["test", "validation"],
|
|
help="Dataset split to use (default: test)",
|
|
)
|
|
parser.add_argument(
|
|
"--max-examples",
|
|
type=int,
|
|
default=None,
|
|
help="Maximum number of examples to run (default: all)",
|
|
)
|
|
parser.add_argument(
|
|
"--context-size",
|
|
type=str,
|
|
default=None,
|
|
help=(
|
|
"Context-size bucket cap, e.g. 8K, 16K, 1M, or exact token count like 16384. "
|
|
"Sets --max-context-len; --min-context-len is only kept when explicitly passed."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--max-context-len",
|
|
type=int,
|
|
default=None,
|
|
help="Maximum context length in tokens",
|
|
)
|
|
parser.add_argument(
|
|
"--labels",
|
|
action="store_true",
|
|
default=False,
|
|
help="Use context_window_text_with_labels for synth examples",
|
|
)
|
|
parser.add_argument(
|
|
"--min-context-len",
|
|
type=int,
|
|
default=1024,
|
|
help="Minimum context length in tokens (default: 1024, upstream behavior)",
|
|
)
|
|
parser.add_argument(
|
|
"--context-window-id",
|
|
type=str,
|
|
default=None,
|
|
help="Run only examples with this context_window_id",
|
|
)
|
|
parser.add_argument(
|
|
"--no-merge-sessions",
|
|
action="store_false",
|
|
dest="merge_sessions",
|
|
default=True,
|
|
help="Store context across multiple sessions instead of a merged session",
|
|
)
|
|
|
|
add_common_arguments(parser)
|
|
args = parser.parse_args()
|
|
min_context_len_explicit = "--min-context-len" in sys.argv
|
|
|
|
error = validate_common_arguments(args)
|
|
if error:
|
|
print(error)
|
|
return 1
|
|
|
|
if args.use_get_context:
|
|
print("Error: --use-get-context is not supported by the OOLONG runner")
|
|
return 1
|
|
|
|
if not args.data_dir.exists():
|
|
print(f"Error: data directory does not exist: {args.data_dir}")
|
|
return 1
|
|
|
|
if args.max_examples is not None and args.max_examples <= 0:
|
|
print(f"Error: max examples must be positive, got {args.max_examples}")
|
|
return 1
|
|
|
|
if args.context_size:
|
|
try:
|
|
exact_size = parse_context_size(args.context_size)
|
|
except ValueError as e:
|
|
print(f"Error: {e}")
|
|
return 1
|
|
|
|
# OOLONG-style behavior: context-size is a bucket cap. Preserve
|
|
# an explicit lower bound only when the user provides one.
|
|
args.max_context_len = exact_size
|
|
if not min_context_len_explicit:
|
|
args.min_context_len = None
|
|
print(
|
|
f"Using context-size cap: <= {exact_size} tokens"
|
|
+ (
|
|
f" (min: {args.min_context_len})"
|
|
if args.min_context_len is not None
|
|
else ""
|
|
)
|
|
)
|
|
|
|
if args.min_context_len is not None and args.min_context_len < 0:
|
|
print(f"Error: min context len must be >= 0, got {args.min_context_len}")
|
|
return 1
|
|
|
|
if args.max_context_len is not None and args.max_context_len <= 0:
|
|
print(f"Error: max context len must be positive, got {args.max_context_len}")
|
|
return 1
|
|
|
|
if (
|
|
args.max_context_len is not None
|
|
and args.min_context_len is not None
|
|
and args.max_context_len < args.min_context_len
|
|
):
|
|
print(
|
|
f"Error: max context len must be >= min context len ({args.max_context_len} < {args.min_context_len})"
|
|
)
|
|
return 1
|
|
|
|
config = RunnerConfig.from_args(args, default_timeout=600)
|
|
runner = OolongRunner(
|
|
config=config,
|
|
variant=args.variant,
|
|
data_dir=args.data_dir,
|
|
split=args.split,
|
|
merge_sessions=args.merge_sessions,
|
|
max_examples=args.max_examples,
|
|
min_context_len=args.min_context_len,
|
|
max_context_len=args.max_context_len,
|
|
context_window_id=args.context_window_id,
|
|
use_labels=args.labels,
|
|
)
|
|
return runner.run_and_summarize()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
exit(main())
|