From 2a99ff8ba352b9264b96b7179ddc134a2bb62787 Mon Sep 17 00:00:00 2001 From: Liu Zhiyuan Date: Fri, 12 Jun 2026 11:11:24 +0800 Subject: [PATCH 1/4] Add MVP reasoning benchmark infrastructure --- benches/README.md | 27 + benches/reasoning_bench.py | 339 +++++++++++ benches/reasoning_smoke.jsonl | 3 + benches/test_reasoning_bench.py | 41 ++ inferlets/reasoning-benchmark/Cargo.toml | 15 + inferlets/reasoning-benchmark/Pie.toml | 19 + inferlets/reasoning-benchmark/src/lib.rs | 598 ++++++++++++++++++++ tests/inferlets/run_all.py | 2 + tests/inferlets/test_reasoning_benchmark.py | 36 ++ 9 files changed, 1080 insertions(+) create mode 100644 benches/reasoning_bench.py create mode 100644 benches/reasoning_smoke.jsonl create mode 100644 benches/test_reasoning_bench.py create mode 100644 inferlets/reasoning-benchmark/Cargo.toml create mode 100644 inferlets/reasoning-benchmark/Pie.toml create mode 100644 inferlets/reasoning-benchmark/src/lib.rs create mode 100644 tests/inferlets/test_reasoning_benchmark.py diff --git a/benches/README.md b/benches/README.md index cb4096a8e..d9853d6fd 100644 --- a/benches/README.md +++ b/benches/README.md @@ -73,6 +73,33 @@ python benches/llamacpp_bench.py tput \ Use `--json-out ` to save machine-readable results. +## Reasoning-pattern benchmark + +`reasoning_bench.py` compares Direct, Best-of-N, Tree-of-Thought, and +Graph-of-Thought through one benchmark-specific inferlet. Reference answers +remain in the Python harness and are never sent to the model. + +Build the inferlet: + +```bash +cargo build --manifest-path inferlets/reasoning-benchmark/Cargo.toml \ + --target wasm32-wasip2 --release +``` + +Run the bundled smoke problems: + +```bash +uv --project sdk/python-server run python benches/reasoning_bench.py \ + --driver cuda_native \ + --model Qwen/Qwen3-0.6B \ + --pattern all \ + --json-out .tmp/reasoning-smoke.json +``` + +For GSM8K, pass a JSONL file containing the official `question` and `answer` +fields. The harness extracts the numeric reference after GSM8K's `####` +delimiter. + ## Fairness defaults - vLLM prefix caching is disabled. diff --git a/benches/reasoning_bench.py b/benches/reasoning_bench.py new file mode 100644 index 000000000..c3e6ea648 --- /dev/null +++ b/benches/reasoning_bench.py @@ -0,0 +1,339 @@ +#!/usr/bin/env python3 +"""Benchmark Pie reasoning patterns on GSM8K-format JSONL data. + +The reference answer stays in this harness and is never sent to the inferlet. +Official GSM8K records (`question`, `answer` containing `#### N`) and simple +records (`id`, `question`, `answer`) are both accepted. +""" +from __future__ import annotations + +import argparse +import asyncio +import json +import math +import re +import statistics +import time +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + + +ROOT = Path(__file__).resolve().parent.parent +INFERLET = "reasoning-benchmark" +PATTERNS = ("direct", "best_of_n", "tree_of_thought", "graph_of_thought") +NUMBER_RE = re.compile(r"[-+]?(?:\d[\d,]*\.?\d*|\.\d+)(?:[eE][-+]?\d+)?") + + +@dataclass(frozen=True) +class Problem: + id: str + question: str + answer: str + + +@dataclass +class RunResult: + problem_id: str + pattern: str + repetition: int + correct: bool + oracle_correct: bool + correct_candidates: int + expected_answer: str + predicted_answer: str | None + latency_s: float + output: dict[str, Any] | None + engine_stats_delta: dict[str, int] + error: str | None = None + + +def normalize_number(value: str | int | float | None) -> str | None: + if value is None: + return None + text = str(value).strip().replace(",", "").replace("$", "") + matches = NUMBER_RE.findall(text) + if not matches: + return None + raw = matches[-1].replace(",", "") + try: + number = float(raw) + except ValueError: + return None + if not math.isfinite(number): + return None + if number.is_integer(): + return str(int(number)) + return format(number, ".12g") + + +def reference_answer(raw: Any) -> str: + text = str(raw) + if "####" in text: + text = text.rsplit("####", 1)[1] + normalized = normalize_number(text) + if normalized is None: + raise ValueError(f"reference answer has no numeric value: {raw!r}") + return normalized + + +def load_problems(path: Path, limit: int | None) -> list[Problem]: + problems: list[Problem] = [] + with path.open(encoding="utf-8") as handle: + for line_number, line in enumerate(handle, 1): + if not line.strip(): + continue + record = json.loads(line) + question = str(record["question"]).strip() + if not question: + raise ValueError(f"{path}:{line_number}: empty question") + problems.append( + Problem( + id=str(record.get("id", f"{path.stem}-{line_number}")), + question=question, + answer=reference_answer(record["answer"]), + ) + ) + if limit is not None and len(problems) >= limit: + break + if not problems: + raise ValueError(f"no problems loaded from {path}") + return problems + + +def inferlet_paths() -> tuple[Path, Path, str]: + try: + import tomllib + except ModuleNotFoundError: + import tomli as tomllib + + directory = ROOT / "inferlets" / INFERLET + candidates = [ + directory / "target" / "wasm32-wasip2" / "release" / "reasoning_benchmark.wasm", + directory / "target" / "wasm32-wasip2" / "debug" / "reasoning_benchmark.wasm", + ] + wasm = next((path for path in candidates if path.exists()), None) + if wasm is None: + raise FileNotFoundError( + "reasoning-benchmark wasm is missing; build it with " + "`cargo build --manifest-path inferlets/reasoning-benchmark/Cargo.toml " + "--target wasm32-wasip2 --release`" + ) + manifest = directory / "Pie.toml" + package = tomllib.loads(manifest.read_text(encoding="utf-8"))["package"] + return wasm, manifest, f"{package['name']}@{package['version']}" + + +def build_config(args: argparse.Namespace): + from pie.config import ( + AuthConfig, + Config, + DriverConfig, + ModelConfig, + ServerConfig, + TelemetryConfig, + ) + + device = [part.strip() for part in args.device.split(",")] + options: dict[str, Any] = {} + if args.driver in {"cuda_native", "portable"}: + options["max_num_kv_pages"] = args.kv_pages + if args.driver == "portable" and args.portable_n_gpu_layers is not None: + options["n_gpu_layers"] = args.portable_n_gpu_layers + return Config( + server=ServerConfig(port=0, max_concurrent_processes=1), + auth=AuthConfig(enabled=False), + telemetry=TelemetryConfig(), + models=[ + ModelConfig( + name="default", + hf_repo=args.model, + driver=DriverConfig( + type=args.driver, + device=device, + options=options, + ), + ) + ], + ) + + +async def model_stats(client) -> dict[str, int]: + ok, raw = await client.query("model_status", "") + if not ok: + raise RuntimeError(f"model_status query failed: {raw}") + return { + key: int(value) + for key, value in json.loads(raw).items() + if isinstance(value, (int, float)) + } + + +def stats_delta(before: dict[str, int], after: dict[str, int]) -> dict[str, int]: + return { + key: after.get(key, 0) - before.get(key, 0) + for key in sorted(set(before) | set(after)) + if key.endswith((".total_batches", ".total_tokens_processed")) + } + + +async def run(args: argparse.Namespace) -> list[RunResult]: + from pie.server import Server + from pie_client import Event + + problems = load_problems(Path(args.dataset), args.max_problems) + wasm, manifest, package = inferlet_paths() + patterns = PATTERNS if args.pattern == "all" else (args.pattern,) + results: list[RunResult] = [] + + async with Server(build_config(args)) as server: + client = await server.connect() + await client.install_program(wasm, manifest, force_overwrite=True) + + for problem in problems: + for pattern in patterns: + for repetition in range(args.repetitions): + payload = { + "pattern": pattern, + "question": problem.question, + "num_candidates": args.num_candidates, + "beam_width": args.beam_width, + "max_tokens": args.max_tokens, + "score_tokens": args.score_tokens, + "temperature": args.temperature, + "top_p": args.top_p, + } + before = await model_stats(client) + started = time.perf_counter() + output = None + error = None + try: + process = await client.launch_process(package, input=payload) + while True: + event, message = await asyncio.wait_for( + process.recv(), timeout=args.timeout + ) + if event == Event.Return: + output = json.loads(message) + break + if event == Event.Error: + raise RuntimeError(str(message)) + except Exception as exc: # keep the full experiment running + error = f"{type(exc).__name__}: {exc}" + latency = time.perf_counter() - started + after = await model_stats(client) + predicted = normalize_number( + output.get("final_answer") if output else None + ) + candidate_answers = [ + normalize_number(candidate.get("answer")) + for candidate in (output or {}).get("candidates", []) + ] + correct_candidates = sum( + answer == problem.answer for answer in candidate_answers + ) + result = RunResult( + problem_id=problem.id, + pattern=pattern, + repetition=repetition, + correct=predicted == problem.answer, + oracle_correct=correct_candidates > 0, + correct_candidates=correct_candidates, + expected_answer=problem.answer, + predicted_answer=predicted, + latency_s=latency, + output=output, + engine_stats_delta=stats_delta(before, after), + error=error, + ) + results.append(result) + status = "correct" if result.correct else "wrong" + if error: + status = "error" + print( + f"{problem.id:18} {pattern:18} rep={repetition} " + f"{status:7} answer={predicted!r} latency={latency:.2f}s", + flush=True, + ) + return results + + +def summarize(results: list[RunResult]) -> dict[str, Any]: + summary: dict[str, Any] = {} + for pattern in sorted({result.pattern for result in results}): + rows = [result for result in results if result.pattern == pattern] + valid = [result for result in rows if result.error is None] + latencies = [result.latency_s for result in valid] + generated = [ + int(result.output["stats"]["generated_tokens"]) + for result in valid + if result.output is not None + ] + summary[pattern] = { + "runs": len(rows), + "errors": sum(result.error is not None for result in rows), + "accuracy": ( + sum(result.correct for result in rows) / len(rows) if rows else 0.0 + ), + "oracle_candidate_accuracy": ( + sum(result.oracle_correct for result in rows) / len(rows) + if rows + else 0.0 + ), + "mean_correct_candidates": ( + statistics.fmean(result.correct_candidates for result in rows) + if rows + else 0.0 + ), + "mean_latency_s": statistics.fmean(latencies) if latencies else None, + "mean_generated_tokens": statistics.fmean(generated) if generated else None, + } + return summary + + +def write_results(path: Path, args: argparse.Namespace, results: list[RunResult]) -> None: + artifact = { + "config": vars(args), + "summary": summarize(results), + "runs": [asdict(result) for result in results], + } + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(artifact, indent=2), encoding="utf-8") + + +def parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description="Pie reasoning-pattern benchmark") + p.add_argument("--dataset", default=str(ROOT / "benches" / "reasoning_smoke.jsonl")) + p.add_argument("--pattern", choices=("all", *PATTERNS), default="all") + p.add_argument("--model", default="Qwen/Qwen3-0.6B") + p.add_argument( + "--driver", + choices=("dev", "dummy", "cuda_native", "portable", "vllm", "sglang"), + default="dev", + ) + p.add_argument("--device", default="cuda:0") + p.add_argument("--max-problems", type=int, default=None) + p.add_argument("--repetitions", type=int, default=1) + p.add_argument("--num-candidates", type=int, default=4) + p.add_argument("--beam-width", type=int, default=2) + p.add_argument("--max-tokens", type=int, default=256) + p.add_argument("--score-tokens", type=int, default=16) + p.add_argument("--temperature", type=float, default=0.7) + p.add_argument("--top-p", type=float, default=0.95) + p.add_argument("--timeout", type=float, default=300.0) + p.add_argument("--kv-pages", type=int, default=2048) + p.add_argument("--portable-n-gpu-layers", type=int, default=None) + p.add_argument("--json-out", default=str(ROOT / ".tmp" / "reasoning-benchmark.json")) + return p + + +def main() -> None: + args = parser().parse_args() + results = asyncio.run(run(args)) + summary = summarize(results) + print("\n" + json.dumps(summary, indent=2)) + write_results(Path(args.json_out), args, results) + + +if __name__ == "__main__": + main() diff --git a/benches/reasoning_smoke.jsonl b/benches/reasoning_smoke.jsonl new file mode 100644 index 000000000..d4f80c121 --- /dev/null +++ b/benches/reasoning_smoke.jsonl @@ -0,0 +1,3 @@ +{"id":"smoke-1","question":"Lina has 12 marbles and buys 7 more. She gives 5 marbles away. How many marbles does she have left?","answer":"14"} +{"id":"smoke-2","question":"A train travels 60 miles per hour for 3 hours. How many miles does it travel?","answer":"180"} +{"id":"smoke-3","question":"A box holds 8 rows of 6 pencils. If 9 pencils are removed, how many pencils remain?","answer":"39"} diff --git a/benches/test_reasoning_bench.py b/benches/test_reasoning_bench.py new file mode 100644 index 000000000..c308451e2 --- /dev/null +++ b/benches/test_reasoning_bench.py @@ -0,0 +1,41 @@ +"""Unit tests for the reasoning benchmark's dataset and answer evaluator.""" +import json +import tempfile +import unittest +from pathlib import Path + +from benches.reasoning_bench import load_problems, normalize_number, reference_answer + + +class ReasoningBenchTests(unittest.TestCase): + def test_normalize_number(self): + self.assertEqual(normalize_number("$1,234.00"), "1234") + self.assertEqual(normalize_number("Final Answer: -2.5"), "-2.5") + self.assertIsNone(normalize_number("no numeric answer")) + + def test_gsm8k_reference_answer(self): + self.assertEqual( + reference_answer("Reasoning with intermediate numbers.\n#### 1,250"), + "1250", + ) + + def test_loads_official_gsm8k_shape(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "test.jsonl" + path.write_text( + json.dumps( + { + "question": "How many?", + "answer": "Compute carefully.\n#### 42", + } + ) + + "\n", + encoding="utf-8", + ) + problems = load_problems(path, None) + self.assertEqual(len(problems), 1) + self.assertEqual(problems[0].answer, "42") + + +if __name__ == "__main__": + unittest.main() diff --git a/inferlets/reasoning-benchmark/Cargo.toml b/inferlets/reasoning-benchmark/Cargo.toml new file mode 100644 index 000000000..fb61254fa --- /dev/null +++ b/inferlets/reasoning-benchmark/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "reasoning-benchmark" +version = "0.1.0" +edition = "2024" + +[lib] +crate-type = ["cdylib"] + +[dependencies] +inferlet = { path = "../../sdk/rust/inferlet" } +futures = "0.3" +serde = { version = "1.0", features = ["derive"] } + +[profile.release] +lto = true diff --git a/inferlets/reasoning-benchmark/Pie.toml b/inferlets/reasoning-benchmark/Pie.toml new file mode 100644 index 000000000..a016e69f7 --- /dev/null +++ b/inferlets/reasoning-benchmark/Pie.toml @@ -0,0 +1,19 @@ +[package] +name = "reasoning-benchmark" +version = "0.1.0" +description = "Unified Direct, Best-of-N, Tree-of-Thought, and Graph-of-Thought benchmark inferlet" +authors = ["Pie Team"] + +[runtime] +core = "^0.2.0" +mcp = "^0.2.0" + +[parameters] +pattern = {type = "string", optional = true, description = "direct | best_of_n | tree_of_thought | graph_of_thought"} +question = {type = "string", description = "Math word problem to solve"} +num_candidates = {type = "int", optional = true, description = "Candidate count for branching patterns (default: 4)"} +beam_width = {type = "int", optional = true, description = "Number of candidates refined by tree_of_thought (default: 2)"} +max_tokens = {type = "int", optional = true, description = "Maximum answer tokens per generation (default: 256)"} +score_tokens = {type = "int", optional = true, description = "Maximum tokens per model-scoring call (default: 16)"} +temperature = {type = "float", optional = true, description = "Candidate sampling temperature (default: 0.7)"} +top_p = {type = "float", optional = true, description = "Candidate top-p value (default: 0.95)"} diff --git a/inferlets/reasoning-benchmark/src/lib.rs b/inferlets/reasoning-benchmark/src/lib.rs new file mode 100644 index 000000000..96e39042a --- /dev/null +++ b/inferlets/reasoning-benchmark/src/lib.rs @@ -0,0 +1,598 @@ +//! Unified inferlet for comparing reasoning workflows on math word problems. +//! +//! The inferlet never receives the reference answer. It returns candidates, +//! selections, and execution counters; the external benchmark harness owns +//! deterministic answer evaluation. + +use std::collections::HashMap; +use std::time::Instant; + +use futures::future; +use inferlet::{Context, Result, chat, model::Model, runtime, sample::Sampler}; +use serde::{Deserialize, Serialize}; + +const SYSTEM_PROMPT: &str = "\ +You solve grade-school mathematical word problems carefully. Show concise \ +reasoning and put the final numeric answer on the last line exactly as \ +\"Final Answer: \". Do not use that phrase anywhere else."; + +const SCORE_SYSTEM: &str = "\ +You evaluate candidate solutions to mathematical word problems. Check the \ +problem interpretation and arithmetic. Respond with only one integer from 1 \ +to 10, where 10 means fully correct."; + +const AGGREGATE_SYSTEM: &str = "\ +You synthesize candidate solutions to mathematical word problems. Compare \ +their reasoning, resolve disagreements by recalculating, and return one \ +concise solution. Put the final numeric answer on the last line exactly as \ +\"Final Answer: \"."; + +#[derive(Deserialize)] +struct Input { + #[serde(default = "default_pattern")] + pattern: String, + question: String, + #[serde(default = "default_num_candidates")] + num_candidates: usize, + #[serde(default = "default_beam_width")] + beam_width: usize, + #[serde(default = "default_max_tokens")] + max_tokens: usize, + #[serde(default = "default_score_tokens")] + score_tokens: usize, + #[serde(default = "default_temperature")] + temperature: f32, + #[serde(default = "default_top_p")] + top_p: f32, +} + +fn default_pattern() -> String { + "direct".into() +} +fn default_num_candidates() -> usize { + 4 +} +fn default_beam_width() -> usize { + 2 +} +fn default_max_tokens() -> usize { + 256 +} +fn default_score_tokens() -> usize { + 16 +} +fn default_temperature() -> f32 { + 0.7 +} +fn default_top_p() -> f32 { + 0.95 +} + +#[derive(Clone, Serialize)] +struct Candidate { + id: String, + stage: &'static str, + response: String, + answer: Option, + score: Option, + generated_tokens: usize, + generator_steps: usize, +} + +#[derive(Default, Serialize)] +struct ExecutionStats { + elapsed_ms: u128, + generated_tokens: usize, + generator_steps: usize, + context_forks: usize, + generation_calls: usize, + scoring_calls: usize, +} + +impl ExecutionStats { + fn add_score(&mut self, generated: &Generated) { + self.generated_tokens += generated.tokens; + self.generator_steps += generated.steps; + self.scoring_calls += 1; + } +} + +#[derive(Serialize)] +struct Output { + pattern: String, + final_response: String, + final_answer: Option, + selected_candidate_id: Option, + candidates: Vec, + stats: ExecutionStats, +} + +struct Generated { + text: String, + tokens: usize, + steps: usize, +} + +#[inferlet::main] +async fn main(input: Input) -> Result { + validate(&input)?; + let started = Instant::now(); + let model_name = runtime::models() + .first() + .cloned() + .ok_or("No models available")?; + let model = Model::load(&model_name)?; + + let mut answer_root = Context::new(&model)?; + answer_root.system(SYSTEM_PROMPT); + answer_root.user(&input.question); + answer_root.flush().await?; + + let mut score_root = Context::new(&model)?; + score_root.system(SCORE_SYSTEM); + score_root.flush().await?; + + let mut aggregate_root = Context::new(&model)?; + aggregate_root.system(AGGREGATE_SYSTEM); + aggregate_root.flush().await?; + + let mut output = match input.pattern.to_ascii_lowercase().as_str() { + "direct" => run_direct(&input, &model, &answer_root).await?, + "best_of_n" | "best-of-n" | "self_consistency" => { + run_best_of_n(&input, &model, &answer_root).await? + } + "tree_of_thought" | "tree-of-thought" | "tot" => { + run_tree_of_thought(&input, &model, &answer_root, &score_root).await? + } + "graph_of_thought" | "graph-of-thought" | "got" => { + run_graph_of_thought(&input, &model, &answer_root, &aggregate_root).await? + } + other => { + return Err(format!( + "unknown pattern '{other}': expected direct, best_of_n, \ + tree_of_thought, or graph_of_thought" + )); + } + }; + output.stats.elapsed_ms = started.elapsed().as_millis(); + Ok(output) +} + +fn validate(input: &Input) -> Result<()> { + if input.question.trim().is_empty() { + return Err("question must not be empty".into()); + } + if !(1..=16).contains(&input.num_candidates) { + return Err("num_candidates must be in [1, 16]".into()); + } + if !(1..=input.num_candidates).contains(&input.beam_width) { + return Err("beam_width must be in [1, num_candidates]".into()); + } + if input.max_tokens == 0 || input.score_tokens == 0 { + return Err("token budgets must be at least 1".into()); + } + if !(input.temperature.is_finite() && (0.0..=2.0).contains(&input.temperature)) { + return Err("temperature must be in [0.0, 2.0]".into()); + } + if !(input.top_p.is_finite() && input.top_p > 0.0 && input.top_p <= 1.0) { + return Err("top_p must be in (0.0, 1.0]".into()); + } + Ok(()) +} + +async fn run_direct(input: &Input, model: &Model, root: &Context) -> Result { + let generated = generate_answer( + root.fork()?, + model, + input.max_tokens, + input.temperature, + input.top_p, + None, + ) + .await?; + let candidate = candidate("direct-0", "answer", generated, None); + let mut stats = ExecutionStats { + context_forks: 1, + ..Default::default() + }; + add_candidate_stats(&mut stats, &candidate); + + Ok(Output { + pattern: "direct".into(), + final_response: candidate.response.clone(), + final_answer: candidate.answer.clone(), + selected_candidate_id: Some(candidate.id.clone()), + candidates: vec![candidate], + stats, + }) +} + +async fn run_best_of_n(input: &Input, model: &Model, root: &Context) -> Result { + let generated = generate_candidates(input, model, root, "sample", None).await?; + let candidates: Vec = generated + .into_iter() + .enumerate() + .map(|(idx, generated)| candidate(&format!("sample-{idx}"), "sample", generated, None)) + .collect(); + let selected = majority_vote_index(&candidates).unwrap_or(0); + let chosen = &candidates[selected]; + let mut stats = ExecutionStats { + context_forks: candidates.len(), + ..Default::default() + }; + for candidate in &candidates { + add_candidate_stats(&mut stats, candidate); + } + + Ok(Output { + pattern: "best_of_n".into(), + final_response: chosen.response.clone(), + final_answer: chosen.answer.clone(), + selected_candidate_id: Some(chosen.id.clone()), + candidates, + stats, + }) +} + +async fn run_tree_of_thought( + input: &Input, + model: &Model, + answer_root: &Context, + score_root: &Context, +) -> Result { + let initial = generate_candidates(input, model, answer_root, "initial", None).await?; + let mut candidates: Vec = initial + .into_iter() + .enumerate() + .map(|(idx, generated)| candidate(&format!("initial-{idx}"), "initial", generated, None)) + .collect(); + let mut stats = ExecutionStats { + context_forks: candidates.len(), + ..Default::default() + }; + for candidate in &candidates { + add_candidate_stats(&mut stats, candidate); + } + + score_candidates( + &input.question, + input.score_tokens, + model, + score_root, + &mut candidates, + &mut stats, + ) + .await?; + + let mut ranked: Vec = (0..candidates.len()).collect(); + ranked.sort_by_key(|&idx| std::cmp::Reverse(candidates[idx].score.unwrap_or(0))); + ranked.truncate(input.beam_width); + + let refine_futures = ranked.iter().map(|&idx| { + let previous = candidates[idx].response.clone(); + let prompt = format!( + "Previous candidate solution:\n{previous}\n\n\ + Critique its interpretation and arithmetic, then produce a corrected \ + complete solution to the original problem." + ); + let ctx = answer_root.fork(); + async move { + generate_answer( + ctx?, + model, + input.max_tokens, + input.temperature, + input.top_p, + Some(&prompt), + ) + .await + } + }); + let refined = future::join_all(refine_futures) + .await + .into_iter() + .collect::>>()?; + stats.context_forks += refined.len(); + + let first_refined = candidates.len(); + for (idx, generated) in refined.into_iter().enumerate() { + let candidate = candidate(&format!("refined-{idx}"), "refined", generated, None); + add_candidate_stats(&mut stats, &candidate); + candidates.push(candidate); + } + score_candidates( + &input.question, + input.score_tokens, + model, + score_root, + &mut candidates[first_refined..], + &mut stats, + ) + .await?; + + let selected = (first_refined..candidates.len()) + .max_by_key(|&idx| candidates[idx].score.unwrap_or(0)) + .unwrap_or(0); + let chosen = &candidates[selected]; + + Ok(Output { + pattern: "tree_of_thought".into(), + final_response: chosen.response.clone(), + final_answer: chosen.answer.clone(), + selected_candidate_id: Some(chosen.id.clone()), + candidates, + stats, + }) +} + +async fn run_graph_of_thought( + input: &Input, + model: &Model, + answer_root: &Context, + aggregate_root: &Context, +) -> Result { + let generated = generate_candidates(input, model, answer_root, "proposal", None).await?; + let mut candidates: Vec = generated + .into_iter() + .enumerate() + .map(|(idx, generated)| candidate(&format!("proposal-{idx}"), "proposal", generated, None)) + .collect(); + let mut stats = ExecutionStats { + context_forks: candidates.len() + 1, + ..Default::default() + }; + for candidate in &candidates { + add_candidate_stats(&mut stats, candidate); + } + + let proposals = candidates + .iter() + .enumerate() + .map(|(idx, candidate)| format!("Candidate {}:\n{}", idx + 1, candidate.response)) + .collect::>() + .join("\n\n"); + let prompt = format!( + "Original problem:\n{}\n\nCandidate solutions:\n{}\n\n\ + Synthesize the most reliable final solution.", + input.question, proposals + ); + let generated = generate_answer( + aggregate_root.fork()?, + model, + input.max_tokens, + 0.0, + 1.0, + Some(&prompt), + ) + .await?; + let aggregate = candidate("aggregate-0", "aggregate", generated, None); + add_candidate_stats(&mut stats, &aggregate); + let final_response = aggregate.response.clone(); + let final_answer = aggregate.answer.clone(); + let selected_candidate_id = Some(aggregate.id.clone()); + candidates.push(aggregate); + + Ok(Output { + pattern: "graph_of_thought".into(), + final_response, + final_answer, + selected_candidate_id, + candidates, + stats, + }) +} + +async fn generate_candidates( + input: &Input, + model: &Model, + root: &Context, + _stage: &'static str, + prompt: Option<&str>, +) -> Result> { + let contexts = (0..input.num_candidates) + .map(|_| root.fork()) + .collect::>>()?; + let futures = contexts.into_iter().map(|ctx| { + generate_answer( + ctx, + model, + input.max_tokens, + input.temperature, + input.top_p, + prompt, + ) + }); + future::join_all(futures) + .await + .into_iter() + .collect::>>() +} + +async fn score_candidates( + question: &str, + max_tokens: usize, + model: &Model, + root: &Context, + candidates: &mut [Candidate], + stats: &mut ExecutionStats, +) -> Result<()> { + let futures = candidates.iter().map(|candidate| { + let prompt = format!( + "Problem:\n{question}\n\nCandidate solution:\n{}\n\nScore:", + candidate.response + ); + let ctx = root.fork(); + async move { generate_answer(ctx?, model, max_tokens, 0.0, 1.0, Some(&prompt)).await } + }); + let scores = future::join_all(futures) + .await + .into_iter() + .collect::>>()?; + stats.context_forks += scores.len(); + for (candidate, generated) in candidates.iter_mut().zip(scores) { + candidate.score = parse_score(&generated.text); + stats.add_score(&generated); + } + Ok(()) +} + +async fn generate_answer( + mut ctx: Context, + model: &Model, + max_tokens: usize, + temperature: f32, + top_p: f32, + prompt: Option<&str>, +) -> Result { + if let Some(prompt) = prompt { + ctx.user(prompt); + } + ctx.cue(); + let sampler = if temperature <= 0.0 { + Sampler::Argmax + } else { + Sampler::TopP { + temperature, + p: top_p, + } + }; + let stops = chat::stop_tokens(model); + let mut generator = ctx.generate(sampler).max_tokens(max_tokens).stop(&stops); + let mut decoder = chat::Decoder::new(model); + let mut text = String::new(); + let mut tokens = 0; + let mut steps = 0; + + while let Some(step) = generator.next()? { + let output = step.execute().await?; + tokens += output.tokens.len(); + steps += 1; + match decoder.feed(&output.tokens)? { + chat::Event::Delta(delta) => text.push_str(&delta), + chat::Event::Done(done) => { + text = done; + break; + } + chat::Event::Idle | chat::Event::Interrupt(_) => {} + } + } + Ok(Generated { + text, + tokens, + steps, + }) +} + +fn candidate(id: &str, stage: &'static str, generated: Generated, score: Option) -> Candidate { + Candidate { + id: id.into(), + stage, + answer: extract_final_answer(&generated.text), + response: generated.text, + score, + generated_tokens: generated.tokens, + generator_steps: generated.steps, + } +} + +fn add_candidate_stats(stats: &mut ExecutionStats, candidate: &Candidate) { + stats.generated_tokens += candidate.generated_tokens; + stats.generator_steps += candidate.generator_steps; + stats.generation_calls += 1; +} + +fn majority_vote_index(candidates: &[Candidate]) -> Option { + let mut counts: HashMap<&str, usize> = HashMap::new(); + for answer in candidates + .iter() + .filter_map(|candidate| candidate.answer.as_deref()) + { + *counts.entry(answer).or_default() += 1; + } + candidates + .iter() + .enumerate() + .filter_map(|(idx, candidate)| { + let answer = candidate.answer.as_deref()?; + Some((idx, counts.get(answer).copied().unwrap_or(0))) + }) + .max_by_key(|(idx, count)| (*count, std::cmp::Reverse(*idx))) + .map(|(idx, _)| idx) +} + +fn extract_final_answer(text: &str) -> Option { + let lower = text.to_ascii_lowercase(); + let marker = "final answer"; + let tail = lower + .rfind(marker) + .map(|pos| &text[pos + marker.len()..]) + .unwrap_or(text); + let tail = tail.find(':').map(|pos| &tail[pos + 1..]).unwrap_or(tail); + extract_last_number(tail) +} + +fn extract_last_number(text: &str) -> Option { + let mut tokens = Vec::new(); + let mut current = String::new(); + for ch in text.chars() { + if ch.is_ascii_digit() || matches!(ch, '-' | '+' | '.' | ',' | '/') { + current.push(ch); + } else if !current.is_empty() { + tokens.push(std::mem::take(&mut current)); + } + } + if !current.is_empty() { + tokens.push(current); + } + tokens + .into_iter() + .rev() + .find_map(|token| normalize_number(&token)) +} + +fn normalize_number(raw: &str) -> Option { + let mut value = raw + .trim_matches(|c: char| matches!(c, '+' | '.' | ',' | '/')) + .replace(',', ""); + if value.is_empty() || !value.chars().any(|c| c.is_ascii_digit()) { + return None; + } + if value.ends_with(".0") { + value.truncate(value.len() - 2); + } + Some(value) +} + +fn parse_score(text: &str) -> Option { + text.split(|c: char| !c.is_ascii_digit()) + .filter_map(|part| part.parse::().ok()) + .find(|score| (1..=10).contains(score)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn extracts_marked_numeric_answer() { + assert_eq!( + extract_final_answer("Work.\nFinal Answer: $1,234.0"), + Some("1234".into()) + ); + } + + #[test] + fn vote_prefers_first_candidate_on_tie() { + let make = |id: &str, answer: &str| Candidate { + id: id.into(), + stage: "sample", + response: String::new(), + answer: Some(answer.into()), + score: None, + generated_tokens: 0, + generator_steps: 0, + }; + let candidates = vec![make("a", "1"), make("b", "2")]; + assert_eq!(majority_vote_index(&candidates), Some(0)); + } +} diff --git a/tests/inferlets/run_all.py b/tests/inferlets/run_all.py index 5f6a20778..654c3f181 100644 --- a/tests/inferlets/run_all.py +++ b/tests/inferlets/run_all.py @@ -14,6 +14,7 @@ from test_tree_of_thought import test_tree_of_thought from test_graph_of_thought import test_graph_of_thought from test_recursion_of_thought import test_recursion_of_thought +from test_reasoning_benchmark import test_reasoning_benchmark from test_agent_react import test_agent_react from test_agent_codeact import test_agent_codeact from test_image_fetch import test_image_fetch @@ -50,6 +51,7 @@ test_tree_of_thought, test_graph_of_thought, test_recursion_of_thought, + test_reasoning_benchmark, test_best_of_n, test_knowledge_graph, test_agent_react, diff --git a/tests/inferlets/test_reasoning_benchmark.py b/tests/inferlets/test_reasoning_benchmark.py new file mode 100644 index 000000000..4d98a45ae --- /dev/null +++ b/tests/inferlets/test_reasoning_benchmark.py @@ -0,0 +1,36 @@ +"""E2E interface test for the unified reasoning benchmark inferlet.""" +import json + +from conftest import run_inferlet, run_tests + + +async def test_reasoning_benchmark(client, args): + for pattern in ( + "direct", + "best_of_n", + "tree_of_thought", + "graph_of_thought", + ): + output = await run_inferlet( + client, + "reasoning-benchmark", + { + "pattern": pattern, + "question": "Mia has 7 apples and buys 5 more. How many apples?", + "num_candidates": 2, + "beam_width": 1, + "max_tokens": 48, + "score_tokens": 8, + }, + timeout=args.timeout, + ) + result = json.loads(output) + assert result["pattern"] == pattern + assert isinstance(result["candidates"], list) + assert result["candidates"] + assert result["stats"]["generation_calls"] >= 1 + assert result["stats"]["context_forks"] >= 1 + + +if __name__ == "__main__": + run_tests([test_reasoning_benchmark]) From 267c9011c4f46dd7f31a4c38be33ba449ef7d626 Mon Sep 17 00:00:00 2001 From: LIU-ZHIYUAN-source Date: Thu, 18 Jun 2026 06:00:48 +0000 Subject: [PATCH 2/4] Add CLI mode to reasoning benchmark runner --- benches/reasoning_bench.py | 206 +++++++++++++++++++++++++++++-------- 1 file changed, 163 insertions(+), 43 deletions(-) diff --git a/benches/reasoning_bench.py b/benches/reasoning_bench.py index c3e6ea648..bd80ec172 100644 --- a/benches/reasoning_bench.py +++ b/benches/reasoning_bench.py @@ -13,6 +13,7 @@ import math import re import statistics +import subprocess import time from dataclasses import asdict, dataclass from pathlib import Path @@ -177,7 +178,135 @@ def stats_delta(before: dict[str, int], after: dict[str, int]) -> dict[str, int] } -async def run(args: argparse.Namespace) -> list[RunResult]: +def payload_for(problem: Problem, pattern: str, args: argparse.Namespace) -> dict[str, Any]: + return { + "pattern": pattern, + "question": problem.question, + "num_candidates": args.num_candidates, + "beam_width": args.beam_width, + "max_tokens": args.max_tokens, + "score_tokens": args.score_tokens, + "temperature": args.temperature, + "top_p": args.top_p, + } + + +def result_for( + problem: Problem, + pattern: str, + repetition: int, + latency_s: float, + output: dict[str, Any] | None, + engine_stats_delta: dict[str, int], + error: str | None, +) -> RunResult: + predicted = normalize_number(output.get("final_answer") if output else None) + candidate_answers = [ + normalize_number(candidate.get("answer")) + for candidate in (output or {}).get("candidates", []) + ] + correct_candidates = sum(answer == problem.answer for answer in candidate_answers) + return RunResult( + problem_id=problem.id, + pattern=pattern, + repetition=repetition, + correct=predicted == problem.answer, + oracle_correct=correct_candidates > 0, + correct_candidates=correct_candidates, + expected_answer=problem.answer, + predicted_answer=predicted, + latency_s=latency_s, + output=output, + engine_stats_delta=engine_stats_delta, + error=error, + ) + + +def print_result(result: RunResult) -> None: + status = "correct" if result.correct else "wrong" + if result.error: + status = "error" + print( + f"{result.problem_id:18} {result.pattern:18} rep={result.repetition} " + f"{status:7} answer={result.predicted_answer!r} " + f"latency={result.latency_s:.2f}s", + flush=True, + ) + + +def parse_pie_run_output(stdout: str) -> dict[str, Any]: + for line in reversed(stdout.splitlines()): + line = line.strip() + if not line: + continue + try: + output = json.loads(line) + except json.JSONDecodeError: + continue + if isinstance(output, dict): + return output + raise ValueError("pie run did not emit a JSON object on stdout") + + +def run_cli(args: argparse.Namespace) -> list[RunResult]: + if not args.config: + raise ValueError("--config is required when --execution-mode=cli") + + problems = load_problems(Path(args.dataset), args.max_problems) + wasm, manifest, _package = inferlet_paths() + patterns = PATTERNS if args.pattern == "all" else (args.pattern,) + results: list[RunResult] = [] + + pie_bin = Path(args.pie_bin) + config = Path(args.config) + for problem in problems: + for pattern in patterns: + for repetition in range(args.repetitions): + payload = payload_for(problem, pattern, args) + cmd = [ + str(pie_bin), + "run", + "--path", + str(wasm), + "--manifest", + str(manifest), + "--config", + str(config), + "--quiet", + "--input", + json.dumps(payload, separators=(",", ":")), + ] + started = time.perf_counter() + output = None + error = None + try: + completed = subprocess.run( + cmd, + cwd=ROOT, + check=True, + capture_output=True, + text=True, + timeout=args.timeout, + ) + output = parse_pie_run_output(completed.stdout) + except Exception as exc: # keep the full experiment running + error = f"{type(exc).__name__}: {exc}" + latency_s = time.perf_counter() - started + result = result_for( + problem, + pattern, + repetition, + latency_s, + output, + {}, + error, + ) + results.append(result) + print_result(result) + return results + + +async def run_embedded(args: argparse.Namespace) -> list[RunResult]: from pie.server import Server from pie_client import Event @@ -193,16 +322,7 @@ async def run(args: argparse.Namespace) -> list[RunResult]: for problem in problems: for pattern in patterns: for repetition in range(args.repetitions): - payload = { - "pattern": pattern, - "question": problem.question, - "num_candidates": args.num_candidates, - "beam_width": args.beam_width, - "max_tokens": args.max_tokens, - "score_tokens": args.score_tokens, - "temperature": args.temperature, - "top_p": args.top_p, - } + payload = payload_for(problem, pattern, args) before = await model_stats(client) started = time.perf_counter() output = None @@ -220,44 +340,28 @@ async def run(args: argparse.Namespace) -> list[RunResult]: raise RuntimeError(str(message)) except Exception as exc: # keep the full experiment running error = f"{type(exc).__name__}: {exc}" - latency = time.perf_counter() - started + latency_s = time.perf_counter() - started after = await model_stats(client) - predicted = normalize_number( - output.get("final_answer") if output else None - ) - candidate_answers = [ - normalize_number(candidate.get("answer")) - for candidate in (output or {}).get("candidates", []) - ] - correct_candidates = sum( - answer == problem.answer for answer in candidate_answers - ) - result = RunResult( - problem_id=problem.id, - pattern=pattern, - repetition=repetition, - correct=predicted == problem.answer, - oracle_correct=correct_candidates > 0, - correct_candidates=correct_candidates, - expected_answer=problem.answer, - predicted_answer=predicted, - latency_s=latency, - output=output, - engine_stats_delta=stats_delta(before, after), - error=error, + result = result_for( + problem, + pattern, + repetition, + latency_s, + output, + stats_delta(before, after), + error, ) results.append(result) - status = "correct" if result.correct else "wrong" - if error: - status = "error" - print( - f"{problem.id:18} {pattern:18} rep={repetition} " - f"{status:7} answer={predicted!r} latency={latency:.2f}s", - flush=True, - ) + print_result(result) return results +async def run(args: argparse.Namespace) -> list[RunResult]: + if args.execution_mode == "cli": + return run_cli(args) + return await run_embedded(args) + + def summarize(results: list[RunResult]) -> dict[str, Any]: summary: dict[str, Any] = {} for pattern in sorted({result.pattern for result in results}): @@ -303,6 +407,22 @@ def write_results(path: Path, args: argparse.Namespace, results: list[RunResult] def parser() -> argparse.ArgumentParser: p = argparse.ArgumentParser(description="Pie reasoning-pattern benchmark") + p.add_argument( + "--execution-mode", + choices=("embedded", "cli"), + default="embedded", + help="Use embedded pie.server or shell out to `pie run`.", + ) + p.add_argument( + "--pie-bin", + default="./target/release/pie", + help="Pie binary to use when --execution-mode=cli.", + ) + p.add_argument( + "--config", + default=None, + help="Pie config TOML. Required when --execution-mode=cli.", + ) p.add_argument("--dataset", default=str(ROOT / "benches" / "reasoning_smoke.jsonl")) p.add_argument("--pattern", choices=("all", *PATTERNS), default="all") p.add_argument("--model", default="Qwen/Qwen3-0.6B") From 2b5e0013630dfe8bf254d6a19c529d1606394a06 Mon Sep 17 00:00:00 2001 From: LIU-ZHIYUAN-source Date: Thu, 18 Jun 2026 08:57:47 +0000 Subject: [PATCH 3/4] Add CLI reasoning benchmark mode and Qwen3 no-thinking control --- benches/reasoning_bench.py | 15 +++++++ benches/test_reasoning_bench.py | 36 ++++++++++++++++- inferlets/reasoning-benchmark/src/lib.rs | 51 +++++++++++++++++------- 3 files changed, 87 insertions(+), 15 deletions(-) diff --git a/benches/reasoning_bench.py b/benches/reasoning_bench.py index bd80ec172..dc112de5a 100644 --- a/benches/reasoning_bench.py +++ b/benches/reasoning_bench.py @@ -188,6 +188,7 @@ def payload_for(problem: Problem, pattern: str, args: argparse.Namespace) -> dic "score_tokens": args.score_tokens, "temperature": args.temperature, "top_p": args.top_p, + "thinking": args.thinking, } @@ -440,6 +441,20 @@ def parser() -> argparse.ArgumentParser: p.add_argument("--score-tokens", type=int, default=16) p.add_argument("--temperature", type=float, default=0.7) p.add_argument("--top-p", type=float, default=0.95) + thinking = p.add_mutually_exclusive_group() + thinking.add_argument( + "--thinking", + dest="thinking", + action="store_true", + help="Allow model thinking blocks.", + ) + thinking.add_argument( + "--no-thinking", + dest="thinking", + action="store_false", + help="Use the model/template no-thinking path. This is the default.", + ) + p.set_defaults(thinking=False) p.add_argument("--timeout", type=float, default=300.0) p.add_argument("--kv-pages", type=int, default=2048) p.add_argument("--portable-n-gpu-layers", type=int, default=None) diff --git a/benches/test_reasoning_bench.py b/benches/test_reasoning_bench.py index c308451e2..8404dcb76 100644 --- a/benches/test_reasoning_bench.py +++ b/benches/test_reasoning_bench.py @@ -4,7 +4,15 @@ import unittest from pathlib import Path -from benches.reasoning_bench import load_problems, normalize_number, reference_answer +from argparse import Namespace + +from benches.reasoning_bench import ( + Problem, + load_problems, + normalize_number, + payload_for, + reference_answer, +) class ReasoningBenchTests(unittest.TestCase): @@ -19,6 +27,32 @@ def test_gsm8k_reference_answer(self): "1250", ) + def test_payload_disables_thinking_by_default(self): + args = Namespace( + num_candidates=4, + beam_width=2, + max_tokens=256, + score_tokens=16, + temperature=0.7, + top_p=0.95, + thinking=False, + ) + payload = payload_for(Problem("p1", "How many?", "42"), "direct", args) + self.assertIs(payload["thinking"], False) + + def test_payload_can_enable_thinking(self): + args = Namespace( + num_candidates=4, + beam_width=2, + max_tokens=256, + score_tokens=16, + temperature=0.7, + top_p=0.95, + thinking=True, + ) + payload = payload_for(Problem("p1", "How many?", "42"), "direct", args) + self.assertIs(payload["thinking"], True) + def test_loads_official_gsm8k_shape(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "test.jsonl" diff --git a/inferlets/reasoning-benchmark/src/lib.rs b/inferlets/reasoning-benchmark/src/lib.rs index 96e39042a..11b90ad52 100644 --- a/inferlets/reasoning-benchmark/src/lib.rs +++ b/inferlets/reasoning-benchmark/src/lib.rs @@ -44,6 +44,8 @@ struct Input { temperature: f32, #[serde(default = "default_top_p")] top_p: f32, + #[serde(default)] + thinking: bool, } fn default_pattern() -> String { @@ -187,6 +189,7 @@ async fn run_direct(input: &Input, model: &Model, root: &Context) -> Result Result<()> { - let futures = candidates.iter().map(|candidate| { - let prompt = format!( - "Problem:\n{question}\n\nCandidate solution:\n{}\n\nScore:", - candidate.response - ); - let ctx = root.fork(); - async move { generate_answer(ctx?, model, max_tokens, 0.0, 1.0, Some(&prompt)).await } - }); + let futures = + candidates.iter().map(|candidate| { + let prompt = format!( + "Problem:\n{question}\n\nCandidate solution:\n{}\n\nScore:", + candidate.response + ); + let ctx = root.fork(); + async move { + generate_answer(ctx?, model, max_tokens, 0.0, 1.0, thinking, Some(&prompt)).await + } + }); let scores = future::join_all(futures) .await .into_iter() @@ -442,12 +454,16 @@ async fn generate_answer( max_tokens: usize, temperature: f32, top_p: f32, + thinking: bool, prompt: Option<&str>, ) -> Result { if let Some(prompt) = prompt { ctx.user(prompt); } ctx.cue(); + if !thinking { + ctx.append(&model.tokenizer().encode("\n\n\n\n")); + } let sampler = if temperature <= 0.0 { Sampler::Argmax } else { @@ -522,12 +538,9 @@ fn majority_vote_index(candidates: &[Candidate]) -> Option { fn extract_final_answer(text: &str) -> Option { let lower = text.to_ascii_lowercase(); - let marker = "final answer"; - let tail = lower - .rfind(marker) - .map(|pos| &text[pos + marker.len()..]) - .unwrap_or(text); - let tail = tail.find(':').map(|pos| &tail[pos + 1..]).unwrap_or(tail); + let marker = "final answer:"; + let pos = lower.rfind(marker)?; + let tail = &text[pos + marker.len()..]; extract_last_number(tail) } @@ -581,6 +594,16 @@ mod tests { ); } + #[test] + fn rejects_unmarked_numeric_answer() { + assert_eq!(extract_final_answer("Work.\nThe answer is 1234."), None); + } + + #[test] + fn rejects_final_answer_without_colon() { + assert_eq!(extract_final_answer("Work.\nFinal Answer is 1234."), None); + } + #[test] fn vote_prefers_first_candidate_on_tie() { let make = |id: &str, answer: &str| Candidate { From 48888f8aeed4a3c156f6f0503332f4a9a6a3b061 Mon Sep 17 00:00:00 2001 From: Liu Zhiyuan Date: Tue, 30 Jun 2026 12:41:10 +0800 Subject: [PATCH 4/4] Split reasoning benchmark into separate inferlets --- benches/README.md | 33 +- benches/reasoning_bench.py | 113 +++- benches/test_reasoning_bench.py | 7 + inferlets/reasoning-base/Cargo.toml | 14 + inferlets/reasoning-base/Pie.toml | 17 + inferlets/reasoning-base/src/lib.rs | 117 ++++ inferlets/reasoning-benchmark/Cargo.toml | 3 +- inferlets/reasoning-benchmark/Pie.toml | 1 + inferlets/reasoning-benchmark/src/lib.rs | 621 +---------------- inferlets/reasoning-best-of-n/Cargo.toml | 14 + inferlets/reasoning-best-of-n/Pie.toml | 19 + inferlets/reasoning-best-of-n/src/lib.rs | 10 + inferlets/reasoning-core/Cargo.toml | 9 + inferlets/reasoning-core/src/lib.rs | 626 ++++++++++++++++++ inferlets/reasoning-direct/Cargo.toml | 14 + inferlets/reasoning-direct/Pie.toml | 19 + inferlets/reasoning-direct/src/lib.rs | 10 + .../reasoning-graph-of-thought/Cargo.toml | 14 + inferlets/reasoning-graph-of-thought/Pie.toml | 19 + .../reasoning-graph-of-thought/src/lib.rs | 10 + .../reasoning-tree-of-thought/Cargo.toml | 14 + inferlets/reasoning-tree-of-thought/Pie.toml | 19 + .../reasoning-tree-of-thought/src/lib.rs | 10 + tests/inferlets/test_reasoning_benchmark.py | 53 +- 24 files changed, 1151 insertions(+), 635 deletions(-) create mode 100644 inferlets/reasoning-base/Cargo.toml create mode 100644 inferlets/reasoning-base/Pie.toml create mode 100644 inferlets/reasoning-base/src/lib.rs create mode 100644 inferlets/reasoning-best-of-n/Cargo.toml create mode 100644 inferlets/reasoning-best-of-n/Pie.toml create mode 100644 inferlets/reasoning-best-of-n/src/lib.rs create mode 100644 inferlets/reasoning-core/Cargo.toml create mode 100644 inferlets/reasoning-core/src/lib.rs create mode 100644 inferlets/reasoning-direct/Cargo.toml create mode 100644 inferlets/reasoning-direct/Pie.toml create mode 100644 inferlets/reasoning-direct/src/lib.rs create mode 100644 inferlets/reasoning-graph-of-thought/Cargo.toml create mode 100644 inferlets/reasoning-graph-of-thought/Pie.toml create mode 100644 inferlets/reasoning-graph-of-thought/src/lib.rs create mode 100644 inferlets/reasoning-tree-of-thought/Cargo.toml create mode 100644 inferlets/reasoning-tree-of-thought/Pie.toml create mode 100644 inferlets/reasoning-tree-of-thought/src/lib.rs diff --git a/benches/README.md b/benches/README.md index d9853d6fd..902068cef 100644 --- a/benches/README.md +++ b/benches/README.md @@ -76,14 +76,19 @@ Use `--json-out ` to save machine-readable results. ## Reasoning-pattern benchmark `reasoning_bench.py` compares Direct, Best-of-N, Tree-of-Thought, and -Graph-of-Thought through one benchmark-specific inferlet. Reference answers -remain in the Python harness and are never sent to the model. +Graph-of-Thought through method-specific benchmark inferlets by default. +Reference answers remain in the Python harness and are never sent to the +model. The older `reasoning-benchmark` inferlet remains available as a +prototype/reference with `--inferlet-mode prototype`. -Build the inferlet: +Build the base and method-specific inferlets: ```bash -cargo build --manifest-path inferlets/reasoning-benchmark/Cargo.toml \ - --target wasm32-wasip2 --release +cargo build --manifest-path inferlets/reasoning-base/Cargo.toml --target wasm32-wasip2 --release +cargo build --manifest-path inferlets/reasoning-direct/Cargo.toml --target wasm32-wasip2 --release +cargo build --manifest-path inferlets/reasoning-best-of-n/Cargo.toml --target wasm32-wasip2 --release +cargo build --manifest-path inferlets/reasoning-tree-of-thought/Cargo.toml --target wasm32-wasip2 --release +cargo build --manifest-path inferlets/reasoning-graph-of-thought/Cargo.toml --target wasm32-wasip2 --release ``` Run the bundled smoke problems: @@ -100,6 +105,24 @@ For GSM8K, pass a JSONL file containing the official `question` and `answer` fields. The harness extracts the numeric reference after GSM8K's `####` delimiter. +`inferlets/reasoning-base` is the minimal prompt-to-completion inferlet. It +does not extract answers, score candidates, or implement a scaling method. The +method-specific benchmark inferlets share one internal implementation crate, +`inferlets/reasoning-core`, so they preserve a common comparison schema: +`pattern`, `final_response`, `final_answer`, `selected_candidate_id`, +`candidates`, and `stats`. Use `--inferlet-dir PATTERN=PATH` to point one +method at an external/user-submitted inferlet with the same comparison schema. + +RatioThink's ToT implementation is a useful reference for the next refinement +pass. The pieces worth adopting first are explicit beam/frontier bookkeeping, +separate `ok`/`incomplete`/`error` node states, observable scorer failures, and +final-answer selection from the deepest surviving beam rather than blindly +trusting the last generated level. Its source also records an important Pie +runtime lesson: in the current host path, sibling decode from one inferlet does +not reliably coalesce into larger GPU batches, so a memory-frugal sequential +or coupled expansion strategy can be the right baseline until forward-pass +deferral improves. + ## Fairness defaults - vLLM prefix caching is disabled. diff --git a/benches/reasoning_bench.py b/benches/reasoning_bench.py index dc112de5a..d98dcffd8 100644 --- a/benches/reasoning_bench.py +++ b/benches/reasoning_bench.py @@ -21,13 +21,21 @@ ROOT = Path(__file__).resolve().parent.parent -INFERLET = "reasoning-benchmark" PATTERNS = ("direct", "best_of_n", "tree_of_thought", "graph_of_thought") +PROTOTYPE_INFERLET = "reasoning-benchmark" +PATTERN_INFERLETS = { + "direct": "reasoning-direct", + "best_of_n": "reasoning-best-of-n", + "tree_of_thought": "reasoning-tree-of-thought", + "graph_of_thought": "reasoning-graph-of-thought", +} NUMBER_RE = re.compile(r"[-+]?(?:\d[\d,]*\.?\d*|\.\d+)(?:[eE][-+]?\d+)?") @dataclass(frozen=True) class Problem: + """One benchmark example with a normalized reference answer.""" + id: str question: str answer: str @@ -35,6 +43,8 @@ class Problem: @dataclass class RunResult: + """Per-pattern result kept in the JSON artifact for later analysis.""" + problem_id: str pattern: str repetition: int @@ -50,6 +60,13 @@ class RunResult: def normalize_number(value: str | int | float | None) -> str | None: + """Return the last numeric value in a model/reference string. + + The evaluator intentionally ignores formatting differences such as commas, + dollar signs, and integer-looking floats so runs can be scored + deterministically without sending the reference answer to the inferlet. + """ + if value is None: return None text = str(value).strip().replace(",", "").replace("$", "") @@ -69,6 +86,8 @@ def normalize_number(value: str | int | float | None) -> str | None: def reference_answer(raw: Any) -> str: + """Extract a normalized numeric answer from GSM8K or simple JSONL records.""" + text = str(raw) if "####" in text: text = text.rsplit("####", 1)[1] @@ -79,6 +98,8 @@ def reference_answer(raw: Any) -> str: def load_problems(path: Path, limit: int | None) -> list[Problem]: + """Load benchmark problems while validating the fields used for scoring.""" + problems: list[Problem] = [] with path.open(encoding="utf-8") as handle: for line_number, line in enumerate(handle, 1): @@ -102,22 +123,45 @@ def load_problems(path: Path, limit: int | None) -> list[Problem]: return problems -def inferlet_paths() -> tuple[Path, Path, str]: +def inferlet_overrides(values: list[str] | None) -> dict[str, Path]: + """Parse PATTERN=PATH overrides for user-submitted inferlet directories.""" + + overrides: dict[str, Path] = {} + for value in values or []: + if "=" not in value: + raise ValueError("--inferlet-dir must be formatted as pattern=path") + pattern, raw_path = value.split("=", 1) + if pattern not in PATTERNS: + raise ValueError(f"unknown inferlet override pattern: {pattern!r}") + overrides[pattern] = Path(raw_path) + return overrides + + +def inferlet_paths(pattern: str, args: argparse.Namespace) -> tuple[Path, Path, str]: + """Resolve the wasm, manifest, and package id for a benchmark pattern.""" + try: import tomllib except ModuleNotFoundError: import tomli as tomllib - directory = ROOT / "inferlets" / INFERLET + if args.inferlet_mode == "prototype": + name = PROTOTYPE_INFERLET + else: + name = PATTERN_INFERLETS[pattern] + directory = inferlet_overrides(args.inferlet_dir).get( + pattern, ROOT / "inferlets" / name + ) + wasm_stem = name.replace("-", "_") candidates = [ - directory / "target" / "wasm32-wasip2" / "release" / "reasoning_benchmark.wasm", - directory / "target" / "wasm32-wasip2" / "debug" / "reasoning_benchmark.wasm", + directory / "target" / "wasm32-wasip2" / "release" / f"{wasm_stem}.wasm", + directory / "target" / "wasm32-wasip2" / "debug" / f"{wasm_stem}.wasm", ] wasm = next((path for path in candidates if path.exists()), None) if wasm is None: raise FileNotFoundError( - "reasoning-benchmark wasm is missing; build it with " - "`cargo build --manifest-path inferlets/reasoning-benchmark/Cargo.toml " + f"{name} wasm is missing; build it with " + f"`cargo build --manifest-path {directory / 'Cargo.toml'} " "--target wasm32-wasip2 --release`" ) manifest = directory / "Pie.toml" @@ -126,6 +170,8 @@ def inferlet_paths() -> tuple[Path, Path, str]: def build_config(args: argparse.Namespace): + """Build an embedded Pie server config for one-model benchmark runs.""" + from pie.config import ( AuthConfig, Config, @@ -160,6 +206,8 @@ def build_config(args: argparse.Namespace): async def model_stats(client) -> dict[str, int]: + """Fetch numeric model counters used to compute coarse engine deltas.""" + ok, raw = await client.query("model_status", "") if not ok: raise RuntimeError(f"model_status query failed: {raw}") @@ -171,6 +219,8 @@ async def model_stats(client) -> dict[str, int]: def stats_delta(before: dict[str, int], after: dict[str, int]) -> dict[str, int]: + """Keep only stable token/batch counters from the model status snapshots.""" + return { key: after.get(key, 0) - before.get(key, 0) for key in sorted(set(before) | set(after)) @@ -179,6 +229,8 @@ def stats_delta(before: dict[str, int], after: dict[str, int]) -> dict[str, int] def payload_for(problem: Problem, pattern: str, args: argparse.Namespace) -> dict[str, Any]: + """Create the inferlet input while keeping the reference answer outside.""" + return { "pattern": pattern, "question": problem.question, @@ -201,6 +253,8 @@ def result_for( engine_stats_delta: dict[str, int], error: str | None, ) -> RunResult: + """Score one inferlet output with deterministic answer and oracle metrics.""" + predicted = normalize_number(output.get("final_answer") if output else None) candidate_answers = [ normalize_number(candidate.get("answer")) @@ -224,6 +278,8 @@ def result_for( def print_result(result: RunResult) -> None: + """Print a compact progress line during long benchmark runs.""" + status = "correct" if result.correct else "wrong" if result.error: status = "error" @@ -236,6 +292,8 @@ def print_result(result: RunResult) -> None: def parse_pie_run_output(stdout: str) -> dict[str, Any]: + """Read the final JSON object from `pie run` stdout.""" + for line in reversed(stdout.splitlines()): line = line.strip() if not line: @@ -250,11 +308,12 @@ def parse_pie_run_output(stdout: str) -> dict[str, Any]: def run_cli(args: argparse.Namespace) -> list[RunResult]: + """Run each problem/pattern through `pie run` for deployment parity tests.""" + if not args.config: raise ValueError("--config is required when --execution-mode=cli") problems = load_problems(Path(args.dataset), args.max_problems) - wasm, manifest, _package = inferlet_paths() patterns = PATTERNS if args.pattern == "all" else (args.pattern,) results: list[RunResult] = [] @@ -264,6 +323,7 @@ def run_cli(args: argparse.Namespace) -> list[RunResult]: for pattern in patterns: for repetition in range(args.repetitions): payload = payload_for(problem, pattern, args) + wasm, manifest, _package = inferlet_paths(pattern, args) cmd = [ str(pie_bin), "run", @@ -308,17 +368,22 @@ def run_cli(args: argparse.Namespace) -> list[RunResult]: async def run_embedded(args: argparse.Namespace) -> list[RunResult]: + """Run all requested patterns against one embedded server instance.""" + from pie.server import Server from pie_client import Event problems = load_problems(Path(args.dataset), args.max_problems) - wasm, manifest, package = inferlet_paths() patterns = PATTERNS if args.pattern == "all" else (args.pattern,) results: list[RunResult] = [] async with Server(build_config(args)) as server: client = await server.connect() - await client.install_program(wasm, manifest, force_overwrite=True) + packages: dict[str, str] = {} + for pattern in patterns: + wasm, manifest, package = inferlet_paths(pattern, args) + await client.install_program(wasm, manifest, force_overwrite=True) + packages[pattern] = package for problem in problems: for pattern in patterns: @@ -329,7 +394,9 @@ async def run_embedded(args: argparse.Namespace) -> list[RunResult]: output = None error = None try: - process = await client.launch_process(package, input=payload) + process = await client.launch_process( + packages[pattern], input=payload + ) while True: event, message = await asyncio.wait_for( process.recv(), timeout=args.timeout @@ -358,12 +425,16 @@ async def run_embedded(args: argparse.Namespace) -> list[RunResult]: async def run(args: argparse.Namespace) -> list[RunResult]: + """Dispatch to the selected execution mode.""" + if args.execution_mode == "cli": return run_cli(args) return await run_embedded(args) def summarize(results: list[RunResult]) -> dict[str, Any]: + """Aggregate accuracy, oracle-candidate, latency, and token metrics.""" + summary: dict[str, Any] = {} for pattern in sorted({result.pattern for result in results}): rows = [result for result in results if result.pattern == pattern] @@ -397,6 +468,8 @@ def summarize(results: list[RunResult]) -> dict[str, Any]: def write_results(path: Path, args: argparse.Namespace, results: list[RunResult]) -> None: + """Persist the full run artifact including config, summary, and raw rows.""" + artifact = { "config": vars(args), "summary": summarize(results), @@ -407,6 +480,8 @@ def write_results(path: Path, args: argparse.Namespace, results: list[RunResult] def parser() -> argparse.ArgumentParser: + """Define CLI options for local smoke checks and RunPod experiments.""" + p = argparse.ArgumentParser(description="Pie reasoning-pattern benchmark") p.add_argument( "--execution-mode", @@ -426,6 +501,22 @@ def parser() -> argparse.ArgumentParser: ) p.add_argument("--dataset", default=str(ROOT / "benches" / "reasoning_smoke.jsonl")) p.add_argument("--pattern", choices=("all", *PATTERNS), default="all") + p.add_argument( + "--inferlet-mode", + choices=("separate", "prototype"), + default="separate", + help=( + "Use method-specific reasoning-* inferlets or the legacy " + "reasoning-benchmark pattern switch." + ), + ) + p.add_argument( + "--inferlet-dir", + action="append", + default=[], + metavar="PATTERN=PATH", + help="Override a method-specific inferlet directory.", + ) p.add_argument("--model", default="Qwen/Qwen3-0.6B") p.add_argument( "--driver", diff --git a/benches/test_reasoning_bench.py b/benches/test_reasoning_bench.py index 8404dcb76..ace516496 100644 --- a/benches/test_reasoning_bench.py +++ b/benches/test_reasoning_bench.py @@ -8,6 +8,7 @@ from benches.reasoning_bench import ( Problem, + PATTERN_INFERLETS, load_problems, normalize_number, payload_for, @@ -53,6 +54,12 @@ def test_payload_can_enable_thinking(self): payload = payload_for(Problem("p1", "How many?", "42"), "direct", args) self.assertIs(payload["thinking"], True) + def test_has_separate_inferlet_for_each_pattern(self): + self.assertEqual( + set(PATTERN_INFERLETS), + {"direct", "best_of_n", "tree_of_thought", "graph_of_thought"}, + ) + def test_loads_official_gsm8k_shape(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "test.jsonl" diff --git a/inferlets/reasoning-base/Cargo.toml b/inferlets/reasoning-base/Cargo.toml new file mode 100644 index 000000000..ac0fca79f --- /dev/null +++ b/inferlets/reasoning-base/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "reasoning-base" +version = "0.1.0" +edition = "2024" + +[lib] +crate-type = ["cdylib"] + +[dependencies] +inferlet = { path = "../../sdk/rust/inferlet" } +serde = { version = "1.0", features = ["derive"] } + +[profile.release] +lto = true diff --git a/inferlets/reasoning-base/Pie.toml b/inferlets/reasoning-base/Pie.toml new file mode 100644 index 000000000..a578a27d4 --- /dev/null +++ b/inferlets/reasoning-base/Pie.toml @@ -0,0 +1,17 @@ +[package] +name = "reasoning-base" +version = "0.1.0" +description = "Base prompt-to-completion inferlet for reasoning benchmarks" +authors = ["Pie Team"] + +[runtime] +core = "^0.2.0" +mcp = "^0.2.0" + +[parameters] +prompt = {type = "string", description = "Prompt to complete"} +system = {type = "string", optional = true, description = "System message for the completion model"} +max_tokens = {type = "int", optional = true, description = "Maximum completion tokens (default: 256)"} +temperature = {type = "float", optional = true, description = "Sampling temperature (default: 0.7)"} +top_p = {type = "float", optional = true, description = "Top-p value (default: 0.95)"} +thinking = {type = "bool", optional = true, description = "Allow model thinking blocks (default: false)"} diff --git a/inferlets/reasoning-base/src/lib.rs b/inferlets/reasoning-base/src/lib.rs new file mode 100644 index 000000000..fdee80822 --- /dev/null +++ b/inferlets/reasoning-base/src/lib.rs @@ -0,0 +1,117 @@ +//! Base prompt-to-completion inferlet for reasoning benchmark experiments. +//! +//! This inferlet intentionally does not implement Direct, Best-of-N, ToT, GoT, +//! answer extraction, scoring, or aggregation. It is the smallest reusable +//! surface that user-submitted reasoning inferlets can build around. + +use inferlet::{Context, Result, chat, model::Model, runtime, sample::Sampler}; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize)] +struct Input { + prompt: String, + #[serde(default = "default_system")] + system: String, + #[serde(default = "default_max_tokens")] + max_tokens: usize, + #[serde(default = "default_temperature")] + temperature: f32, + #[serde(default = "default_top_p")] + top_p: f32, + #[serde(default)] + thinking: bool, +} + +fn default_system() -> String { + "You are a helpful, careful assistant.".into() +} + +fn default_max_tokens() -> usize { + 256 +} + +fn default_temperature() -> f32 { + 0.7 +} + +fn default_top_p() -> f32 { + 0.95 +} + +#[derive(Serialize)] +struct Output { + completion: String, + stats: GenerationStats, +} + +#[derive(Default, Serialize)] +struct GenerationStats { + generated_tokens: usize, + generator_steps: usize, +} + +#[inferlet::main] +async fn main(input: Input) -> Result { + validate(&input)?; + + let model_name = runtime::models() + .first() + .cloned() + .ok_or("No models available")?; + let model = Model::load(&model_name)?; + + let mut ctx = Context::new(&model)?; + ctx.system(&input.system).user(&input.prompt).cue(); + if !input.thinking { + ctx.append(&model.tokenizer().encode("\n\n\n\n")); + } + + let sampler = if input.temperature <= 0.0 { + Sampler::Argmax + } else { + Sampler::TopP { + temperature: input.temperature, + p: input.top_p, + } + }; + let stops = chat::stop_tokens(&model); + let mut generator = ctx + .generate(sampler) + .max_tokens(input.max_tokens) + .stop(&stops); + let mut decoder = chat::Decoder::new(&model); + let mut completion = String::new(); + let mut stats = GenerationStats::default(); + + while let Some(step) = generator.next()? { + let output = step.execute().await?; + stats.generated_tokens += output.tokens.len(); + stats.generator_steps += 1; + match decoder.feed(&output.tokens)? { + chat::Event::Delta(delta) => completion.push_str(&delta), + chat::Event::Done(done) => { + completion = done; + break; + } + chat::Event::Idle | chat::Event::Interrupt(_) => {} + } + } + + Ok(Output { completion, stats }) +} + +fn validate(input: &Input) -> Result<()> { + if input.prompt.trim().is_empty() { + return Err("prompt must not be empty".into()); + } + if input.max_tokens == 0 { + return Err("max_tokens must be at least 1".into()); + } + if !(input.temperature.is_finite() && (0.0..=2.0).contains(&input.temperature)) { + return Err("temperature must be in [0.0, 2.0]".into()); + } + if !(input.top_p.is_finite() && input.top_p > 0.0 && input.top_p <= 1.0) { + return Err("top_p must be in (0.0, 1.0]".into()); + } + Ok(()) +} diff --git a/inferlets/reasoning-benchmark/Cargo.toml b/inferlets/reasoning-benchmark/Cargo.toml index fb61254fa..889ea90c8 100644 --- a/inferlets/reasoning-benchmark/Cargo.toml +++ b/inferlets/reasoning-benchmark/Cargo.toml @@ -8,8 +8,7 @@ crate-type = ["cdylib"] [dependencies] inferlet = { path = "../../sdk/rust/inferlet" } -futures = "0.3" -serde = { version = "1.0", features = ["derive"] } +reasoning-core = { path = "../reasoning-core" } [profile.release] lto = true diff --git a/inferlets/reasoning-benchmark/Pie.toml b/inferlets/reasoning-benchmark/Pie.toml index a016e69f7..ff09a28c7 100644 --- a/inferlets/reasoning-benchmark/Pie.toml +++ b/inferlets/reasoning-benchmark/Pie.toml @@ -17,3 +17,4 @@ max_tokens = {type = "int", optional = true, description = "Maximum answer token score_tokens = {type = "int", optional = true, description = "Maximum tokens per model-scoring call (default: 16)"} temperature = {type = "float", optional = true, description = "Candidate sampling temperature (default: 0.7)"} top_p = {type = "float", optional = true, description = "Candidate top-p value (default: 0.95)"} +thinking = {type = "bool", optional = true, description = "Allow model thinking blocks (default: false)"} diff --git a/inferlets/reasoning-benchmark/src/lib.rs b/inferlets/reasoning-benchmark/src/lib.rs index 11b90ad52..f44b82691 100644 --- a/inferlets/reasoning-benchmark/src/lib.rs +++ b/inferlets/reasoning-benchmark/src/lib.rs @@ -1,621 +1,12 @@ -//! Unified inferlet for comparing reasoning workflows on math word problems. +//! Compatibility inferlet for comparing reasoning workflows with a pattern switch. //! -//! The inferlet never receives the reference answer. It returns candidates, -//! selections, and execution counters; the external benchmark harness owns -//! deterministic answer evaluation. +//! New benchmark runs can target the method-specific `reasoning-*` inferlets. +//! This wrapper remains as a prototype/reference surface. -use std::collections::HashMap; -use std::time::Instant; - -use futures::future; -use inferlet::{Context, Result, chat, model::Model, runtime, sample::Sampler}; -use serde::{Deserialize, Serialize}; - -const SYSTEM_PROMPT: &str = "\ -You solve grade-school mathematical word problems carefully. Show concise \ -reasoning and put the final numeric answer on the last line exactly as \ -\"Final Answer: \". Do not use that phrase anywhere else."; - -const SCORE_SYSTEM: &str = "\ -You evaluate candidate solutions to mathematical word problems. Check the \ -problem interpretation and arithmetic. Respond with only one integer from 1 \ -to 10, where 10 means fully correct."; - -const AGGREGATE_SYSTEM: &str = "\ -You synthesize candidate solutions to mathematical word problems. Compare \ -their reasoning, resolve disagreements by recalculating, and return one \ -concise solution. Put the final numeric answer on the last line exactly as \ -\"Final Answer: \"."; - -#[derive(Deserialize)] -struct Input { - #[serde(default = "default_pattern")] - pattern: String, - question: String, - #[serde(default = "default_num_candidates")] - num_candidates: usize, - #[serde(default = "default_beam_width")] - beam_width: usize, - #[serde(default = "default_max_tokens")] - max_tokens: usize, - #[serde(default = "default_score_tokens")] - score_tokens: usize, - #[serde(default = "default_temperature")] - temperature: f32, - #[serde(default = "default_top_p")] - top_p: f32, - #[serde(default)] - thinking: bool, -} - -fn default_pattern() -> String { - "direct".into() -} -fn default_num_candidates() -> usize { - 4 -} -fn default_beam_width() -> usize { - 2 -} -fn default_max_tokens() -> usize { - 256 -} -fn default_score_tokens() -> usize { - 16 -} -fn default_temperature() -> f32 { - 0.7 -} -fn default_top_p() -> f32 { - 0.95 -} - -#[derive(Clone, Serialize)] -struct Candidate { - id: String, - stage: &'static str, - response: String, - answer: Option, - score: Option, - generated_tokens: usize, - generator_steps: usize, -} - -#[derive(Default, Serialize)] -struct ExecutionStats { - elapsed_ms: u128, - generated_tokens: usize, - generator_steps: usize, - context_forks: usize, - generation_calls: usize, - scoring_calls: usize, -} - -impl ExecutionStats { - fn add_score(&mut self, generated: &Generated) { - self.generated_tokens += generated.tokens; - self.generator_steps += generated.steps; - self.scoring_calls += 1; - } -} - -#[derive(Serialize)] -struct Output { - pattern: String, - final_response: String, - final_answer: Option, - selected_candidate_id: Option, - candidates: Vec, - stats: ExecutionStats, -} - -struct Generated { - text: String, - tokens: usize, - steps: usize, -} +use inferlet::Result; +use reasoning_core::{Input, Output}; #[inferlet::main] async fn main(input: Input) -> Result { - validate(&input)?; - let started = Instant::now(); - let model_name = runtime::models() - .first() - .cloned() - .ok_or("No models available")?; - let model = Model::load(&model_name)?; - - let mut answer_root = Context::new(&model)?; - answer_root.system(SYSTEM_PROMPT); - answer_root.user(&input.question); - answer_root.flush().await?; - - let mut score_root = Context::new(&model)?; - score_root.system(SCORE_SYSTEM); - score_root.flush().await?; - - let mut aggregate_root = Context::new(&model)?; - aggregate_root.system(AGGREGATE_SYSTEM); - aggregate_root.flush().await?; - - let mut output = match input.pattern.to_ascii_lowercase().as_str() { - "direct" => run_direct(&input, &model, &answer_root).await?, - "best_of_n" | "best-of-n" | "self_consistency" => { - run_best_of_n(&input, &model, &answer_root).await? - } - "tree_of_thought" | "tree-of-thought" | "tot" => { - run_tree_of_thought(&input, &model, &answer_root, &score_root).await? - } - "graph_of_thought" | "graph-of-thought" | "got" => { - run_graph_of_thought(&input, &model, &answer_root, &aggregate_root).await? - } - other => { - return Err(format!( - "unknown pattern '{other}': expected direct, best_of_n, \ - tree_of_thought, or graph_of_thought" - )); - } - }; - output.stats.elapsed_ms = started.elapsed().as_millis(); - Ok(output) -} - -fn validate(input: &Input) -> Result<()> { - if input.question.trim().is_empty() { - return Err("question must not be empty".into()); - } - if !(1..=16).contains(&input.num_candidates) { - return Err("num_candidates must be in [1, 16]".into()); - } - if !(1..=input.num_candidates).contains(&input.beam_width) { - return Err("beam_width must be in [1, num_candidates]".into()); - } - if input.max_tokens == 0 || input.score_tokens == 0 { - return Err("token budgets must be at least 1".into()); - } - if !(input.temperature.is_finite() && (0.0..=2.0).contains(&input.temperature)) { - return Err("temperature must be in [0.0, 2.0]".into()); - } - if !(input.top_p.is_finite() && input.top_p > 0.0 && input.top_p <= 1.0) { - return Err("top_p must be in (0.0, 1.0]".into()); - } - Ok(()) -} - -async fn run_direct(input: &Input, model: &Model, root: &Context) -> Result { - let generated = generate_answer( - root.fork()?, - model, - input.max_tokens, - input.temperature, - input.top_p, - input.thinking, - None, - ) - .await?; - let candidate = candidate("direct-0", "answer", generated, None); - let mut stats = ExecutionStats { - context_forks: 1, - ..Default::default() - }; - add_candidate_stats(&mut stats, &candidate); - - Ok(Output { - pattern: "direct".into(), - final_response: candidate.response.clone(), - final_answer: candidate.answer.clone(), - selected_candidate_id: Some(candidate.id.clone()), - candidates: vec![candidate], - stats, - }) -} - -async fn run_best_of_n(input: &Input, model: &Model, root: &Context) -> Result { - let generated = generate_candidates(input, model, root, "sample", None).await?; - let candidates: Vec = generated - .into_iter() - .enumerate() - .map(|(idx, generated)| candidate(&format!("sample-{idx}"), "sample", generated, None)) - .collect(); - let selected = majority_vote_index(&candidates).unwrap_or(0); - let chosen = &candidates[selected]; - let mut stats = ExecutionStats { - context_forks: candidates.len(), - ..Default::default() - }; - for candidate in &candidates { - add_candidate_stats(&mut stats, candidate); - } - - Ok(Output { - pattern: "best_of_n".into(), - final_response: chosen.response.clone(), - final_answer: chosen.answer.clone(), - selected_candidate_id: Some(chosen.id.clone()), - candidates, - stats, - }) -} - -async fn run_tree_of_thought( - input: &Input, - model: &Model, - answer_root: &Context, - score_root: &Context, -) -> Result { - let initial = generate_candidates(input, model, answer_root, "initial", None).await?; - let mut candidates: Vec = initial - .into_iter() - .enumerate() - .map(|(idx, generated)| candidate(&format!("initial-{idx}"), "initial", generated, None)) - .collect(); - let mut stats = ExecutionStats { - context_forks: candidates.len(), - ..Default::default() - }; - for candidate in &candidates { - add_candidate_stats(&mut stats, candidate); - } - - score_candidates( - &input.question, - input.thinking, - input.score_tokens, - model, - score_root, - &mut candidates, - &mut stats, - ) - .await?; - - let mut ranked: Vec = (0..candidates.len()).collect(); - ranked.sort_by_key(|&idx| std::cmp::Reverse(candidates[idx].score.unwrap_or(0))); - ranked.truncate(input.beam_width); - - let refine_futures = ranked.iter().map(|&idx| { - let previous = candidates[idx].response.clone(); - let prompt = format!( - "Previous candidate solution:\n{previous}\n\n\ - Critique its interpretation and arithmetic, then produce a corrected \ - complete solution to the original problem." - ); - let ctx = answer_root.fork(); - async move { - generate_answer( - ctx?, - model, - input.max_tokens, - input.temperature, - input.top_p, - input.thinking, - Some(&prompt), - ) - .await - } - }); - let refined = future::join_all(refine_futures) - .await - .into_iter() - .collect::>>()?; - stats.context_forks += refined.len(); - - let first_refined = candidates.len(); - for (idx, generated) in refined.into_iter().enumerate() { - let candidate = candidate(&format!("refined-{idx}"), "refined", generated, None); - add_candidate_stats(&mut stats, &candidate); - candidates.push(candidate); - } - score_candidates( - &input.question, - input.thinking, - input.score_tokens, - model, - score_root, - &mut candidates[first_refined..], - &mut stats, - ) - .await?; - - let selected = (first_refined..candidates.len()) - .max_by_key(|&idx| candidates[idx].score.unwrap_or(0)) - .unwrap_or(0); - let chosen = &candidates[selected]; - - Ok(Output { - pattern: "tree_of_thought".into(), - final_response: chosen.response.clone(), - final_answer: chosen.answer.clone(), - selected_candidate_id: Some(chosen.id.clone()), - candidates, - stats, - }) -} - -async fn run_graph_of_thought( - input: &Input, - model: &Model, - answer_root: &Context, - aggregate_root: &Context, -) -> Result { - let generated = generate_candidates(input, model, answer_root, "proposal", None).await?; - let mut candidates: Vec = generated - .into_iter() - .enumerate() - .map(|(idx, generated)| candidate(&format!("proposal-{idx}"), "proposal", generated, None)) - .collect(); - let mut stats = ExecutionStats { - context_forks: candidates.len() + 1, - ..Default::default() - }; - for candidate in &candidates { - add_candidate_stats(&mut stats, candidate); - } - - let proposals = candidates - .iter() - .enumerate() - .map(|(idx, candidate)| format!("Candidate {}:\n{}", idx + 1, candidate.response)) - .collect::>() - .join("\n\n"); - let prompt = format!( - "Original problem:\n{}\n\nCandidate solutions:\n{}\n\n\ - Synthesize the most reliable final solution.", - input.question, proposals - ); - let generated = generate_answer( - aggregate_root.fork()?, - model, - input.max_tokens, - 0.0, - 1.0, - input.thinking, - Some(&prompt), - ) - .await?; - let aggregate = candidate("aggregate-0", "aggregate", generated, None); - add_candidate_stats(&mut stats, &aggregate); - let final_response = aggregate.response.clone(); - let final_answer = aggregate.answer.clone(); - let selected_candidate_id = Some(aggregate.id.clone()); - candidates.push(aggregate); - - Ok(Output { - pattern: "graph_of_thought".into(), - final_response, - final_answer, - selected_candidate_id, - candidates, - stats, - }) -} - -async fn generate_candidates( - input: &Input, - model: &Model, - root: &Context, - _stage: &'static str, - prompt: Option<&str>, -) -> Result> { - let contexts = (0..input.num_candidates) - .map(|_| root.fork()) - .collect::>>()?; - let futures = contexts.into_iter().map(|ctx| { - generate_answer( - ctx, - model, - input.max_tokens, - input.temperature, - input.top_p, - input.thinking, - prompt, - ) - }); - future::join_all(futures) - .await - .into_iter() - .collect::>>() -} - -async fn score_candidates( - question: &str, - thinking: bool, - max_tokens: usize, - model: &Model, - root: &Context, - candidates: &mut [Candidate], - stats: &mut ExecutionStats, -) -> Result<()> { - let futures = - candidates.iter().map(|candidate| { - let prompt = format!( - "Problem:\n{question}\n\nCandidate solution:\n{}\n\nScore:", - candidate.response - ); - let ctx = root.fork(); - async move { - generate_answer(ctx?, model, max_tokens, 0.0, 1.0, thinking, Some(&prompt)).await - } - }); - let scores = future::join_all(futures) - .await - .into_iter() - .collect::>>()?; - stats.context_forks += scores.len(); - for (candidate, generated) in candidates.iter_mut().zip(scores) { - candidate.score = parse_score(&generated.text); - stats.add_score(&generated); - } - Ok(()) -} - -async fn generate_answer( - mut ctx: Context, - model: &Model, - max_tokens: usize, - temperature: f32, - top_p: f32, - thinking: bool, - prompt: Option<&str>, -) -> Result { - if let Some(prompt) = prompt { - ctx.user(prompt); - } - ctx.cue(); - if !thinking { - ctx.append(&model.tokenizer().encode("\n\n\n\n")); - } - let sampler = if temperature <= 0.0 { - Sampler::Argmax - } else { - Sampler::TopP { - temperature, - p: top_p, - } - }; - let stops = chat::stop_tokens(model); - let mut generator = ctx.generate(sampler).max_tokens(max_tokens).stop(&stops); - let mut decoder = chat::Decoder::new(model); - let mut text = String::new(); - let mut tokens = 0; - let mut steps = 0; - - while let Some(step) = generator.next()? { - let output = step.execute().await?; - tokens += output.tokens.len(); - steps += 1; - match decoder.feed(&output.tokens)? { - chat::Event::Delta(delta) => text.push_str(&delta), - chat::Event::Done(done) => { - text = done; - break; - } - chat::Event::Idle | chat::Event::Interrupt(_) => {} - } - } - Ok(Generated { - text, - tokens, - steps, - }) -} - -fn candidate(id: &str, stage: &'static str, generated: Generated, score: Option) -> Candidate { - Candidate { - id: id.into(), - stage, - answer: extract_final_answer(&generated.text), - response: generated.text, - score, - generated_tokens: generated.tokens, - generator_steps: generated.steps, - } -} - -fn add_candidate_stats(stats: &mut ExecutionStats, candidate: &Candidate) { - stats.generated_tokens += candidate.generated_tokens; - stats.generator_steps += candidate.generator_steps; - stats.generation_calls += 1; -} - -fn majority_vote_index(candidates: &[Candidate]) -> Option { - let mut counts: HashMap<&str, usize> = HashMap::new(); - for answer in candidates - .iter() - .filter_map(|candidate| candidate.answer.as_deref()) - { - *counts.entry(answer).or_default() += 1; - } - candidates - .iter() - .enumerate() - .filter_map(|(idx, candidate)| { - let answer = candidate.answer.as_deref()?; - Some((idx, counts.get(answer).copied().unwrap_or(0))) - }) - .max_by_key(|(idx, count)| (*count, std::cmp::Reverse(*idx))) - .map(|(idx, _)| idx) -} - -fn extract_final_answer(text: &str) -> Option { - let lower = text.to_ascii_lowercase(); - let marker = "final answer:"; - let pos = lower.rfind(marker)?; - let tail = &text[pos + marker.len()..]; - extract_last_number(tail) -} - -fn extract_last_number(text: &str) -> Option { - let mut tokens = Vec::new(); - let mut current = String::new(); - for ch in text.chars() { - if ch.is_ascii_digit() || matches!(ch, '-' | '+' | '.' | ',' | '/') { - current.push(ch); - } else if !current.is_empty() { - tokens.push(std::mem::take(&mut current)); - } - } - if !current.is_empty() { - tokens.push(current); - } - tokens - .into_iter() - .rev() - .find_map(|token| normalize_number(&token)) -} - -fn normalize_number(raw: &str) -> Option { - let mut value = raw - .trim_matches(|c: char| matches!(c, '+' | '.' | ',' | '/')) - .replace(',', ""); - if value.is_empty() || !value.chars().any(|c| c.is_ascii_digit()) { - return None; - } - if value.ends_with(".0") { - value.truncate(value.len() - 2); - } - Some(value) -} - -fn parse_score(text: &str) -> Option { - text.split(|c: char| !c.is_ascii_digit()) - .filter_map(|part| part.parse::().ok()) - .find(|score| (1..=10).contains(score)) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn extracts_marked_numeric_answer() { - assert_eq!( - extract_final_answer("Work.\nFinal Answer: $1,234.0"), - Some("1234".into()) - ); - } - - #[test] - fn rejects_unmarked_numeric_answer() { - assert_eq!(extract_final_answer("Work.\nThe answer is 1234."), None); - } - - #[test] - fn rejects_final_answer_without_colon() { - assert_eq!(extract_final_answer("Work.\nFinal Answer is 1234."), None); - } - - #[test] - fn vote_prefers_first_candidate_on_tie() { - let make = |id: &str, answer: &str| Candidate { - id: id.into(), - stage: "sample", - response: String::new(), - answer: Some(answer.into()), - score: None, - generated_tokens: 0, - generator_steps: 0, - }; - let candidates = vec![make("a", "1"), make("b", "2")]; - assert_eq!(majority_vote_index(&candidates), Some(0)); - } + reasoning_core::run(input).await } diff --git a/inferlets/reasoning-best-of-n/Cargo.toml b/inferlets/reasoning-best-of-n/Cargo.toml new file mode 100644 index 000000000..3d7ccc158 --- /dev/null +++ b/inferlets/reasoning-best-of-n/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "reasoning-best-of-n" +version = "0.1.0" +edition = "2024" + +[lib] +crate-type = ["cdylib"] + +[dependencies] +inferlet = { path = "../../sdk/rust/inferlet" } +reasoning-core = { path = "../reasoning-core" } + +[profile.release] +lto = true diff --git a/inferlets/reasoning-best-of-n/Pie.toml b/inferlets/reasoning-best-of-n/Pie.toml new file mode 100644 index 000000000..49497c46c --- /dev/null +++ b/inferlets/reasoning-best-of-n/Pie.toml @@ -0,0 +1,19 @@ +[package] +name = "reasoning-best-of-n" +version = "0.1.0" +description = "Best-of-N/self-consistency inferlet for reasoning benchmarks" +authors = ["Pie Team"] + +[runtime] +core = "^0.2.0" +mcp = "^0.2.0" + +[parameters] +question = {type = "string", description = "Math word problem to solve"} +num_candidates = {type = "int", optional = true, description = "Candidate count (default: 4)"} +beam_width = {type = "int", optional = true, description = "Accepted for runner compatibility"} +max_tokens = {type = "int", optional = true, description = "Maximum answer tokens per generation (default: 256)"} +score_tokens = {type = "int", optional = true, description = "Accepted for runner compatibility"} +temperature = {type = "float", optional = true, description = "Candidate sampling temperature (default: 0.7)"} +top_p = {type = "float", optional = true, description = "Candidate top-p value (default: 0.95)"} +thinking = {type = "bool", optional = true, description = "Allow model thinking blocks (default: false)"} diff --git a/inferlets/reasoning-best-of-n/src/lib.rs b/inferlets/reasoning-best-of-n/src/lib.rs new file mode 100644 index 000000000..ae21b8e01 --- /dev/null +++ b/inferlets/reasoning-best-of-n/src/lib.rs @@ -0,0 +1,10 @@ +//! Method-isolated Best-of-N inferlet for reasoning benchmarks. + +use inferlet::Result; +use reasoning_core::{Input, Output}; + +#[inferlet::main] +async fn main(mut input: Input) -> Result { + input.force_pattern("best_of_n"); + reasoning_core::run(input).await +} diff --git a/inferlets/reasoning-core/Cargo.toml b/inferlets/reasoning-core/Cargo.toml new file mode 100644 index 000000000..57e8cc0a9 --- /dev/null +++ b/inferlets/reasoning-core/Cargo.toml @@ -0,0 +1,9 @@ +[package] +name = "reasoning-core" +version = "0.1.0" +edition = "2024" + +[dependencies] +inferlet = { path = "../../sdk/rust/inferlet" } +futures = "0.3" +serde = { version = "1.0", features = ["derive"] } diff --git a/inferlets/reasoning-core/src/lib.rs b/inferlets/reasoning-core/src/lib.rs new file mode 100644 index 000000000..af5026b6d --- /dev/null +++ b/inferlets/reasoning-core/src/lib.rs @@ -0,0 +1,626 @@ +//! Shared implementation for benchmark-focused reasoning inferlets. +//! +//! The inferlet never receives the reference answer. It returns candidates, +//! selections, and execution counters; the external benchmark harness owns +//! deterministic answer evaluation. + +use std::collections::HashMap; +use std::time::Instant; + +use futures::future; +use inferlet::{Context, Result, chat, model::Model, runtime, sample::Sampler}; +use serde::{Deserialize, Serialize}; + +const SYSTEM_PROMPT: &str = "\ +You solve grade-school mathematical word problems carefully. Show concise \ +reasoning and put the final numeric answer on the last line exactly as \ +\"Final Answer: \". Do not use that phrase anywhere else."; + +const SCORE_SYSTEM: &str = "\ +You evaluate candidate solutions to mathematical word problems. Check the \ +problem interpretation and arithmetic. Respond with only one integer from 1 \ +to 10, where 10 means fully correct."; + +const AGGREGATE_SYSTEM: &str = "\ +You synthesize candidate solutions to mathematical word problems. Compare \ +their reasoning, resolve disagreements by recalculating, and return one \ +concise solution. Put the final numeric answer on the last line exactly as \ +\"Final Answer: \"."; + +#[derive(Deserialize)] +pub struct Input { + #[serde(default = "default_pattern")] + pattern: String, + question: String, + #[serde(default = "default_num_candidates")] + num_candidates: usize, + #[serde(default = "default_beam_width")] + beam_width: usize, + #[serde(default = "default_max_tokens")] + max_tokens: usize, + #[serde(default = "default_score_tokens")] + score_tokens: usize, + #[serde(default = "default_temperature")] + temperature: f32, + #[serde(default = "default_top_p")] + top_p: f32, + #[serde(default)] + thinking: bool, +} + +impl Input { + pub fn force_pattern(&mut self, pattern: &'static str) { + self.pattern = pattern.into(); + } +} + +fn default_pattern() -> String { + "direct".into() +} +fn default_num_candidates() -> usize { + 4 +} +fn default_beam_width() -> usize { + 2 +} +fn default_max_tokens() -> usize { + 256 +} +fn default_score_tokens() -> usize { + 16 +} +fn default_temperature() -> f32 { + 0.7 +} +fn default_top_p() -> f32 { + 0.95 +} + +#[derive(Clone, Serialize)] +struct Candidate { + id: String, + stage: &'static str, + response: String, + answer: Option, + score: Option, + generated_tokens: usize, + generator_steps: usize, +} + +#[derive(Default, Serialize)] +struct ExecutionStats { + elapsed_ms: u128, + generated_tokens: usize, + generator_steps: usize, + context_forks: usize, + generation_calls: usize, + scoring_calls: usize, +} + +impl ExecutionStats { + fn add_score(&mut self, generated: &Generated) { + self.generated_tokens += generated.tokens; + self.generator_steps += generated.steps; + self.scoring_calls += 1; + } +} + +#[derive(Serialize)] +pub struct Output { + pattern: String, + final_response: String, + final_answer: Option, + selected_candidate_id: Option, + candidates: Vec, + stats: ExecutionStats, +} + +struct Generated { + text: String, + tokens: usize, + steps: usize, +} + +pub async fn run(input: Input) -> Result { + validate(&input)?; + let started = Instant::now(); + let model_name = runtime::models() + .first() + .cloned() + .ok_or("No models available")?; + let model = Model::load(&model_name)?; + + let mut answer_root = Context::new(&model)?; + answer_root.system(SYSTEM_PROMPT); + answer_root.user(&input.question); + answer_root.flush().await?; + + let mut score_root = Context::new(&model)?; + score_root.system(SCORE_SYSTEM); + score_root.flush().await?; + + let mut aggregate_root = Context::new(&model)?; + aggregate_root.system(AGGREGATE_SYSTEM); + aggregate_root.flush().await?; + + let mut output = match input.pattern.to_ascii_lowercase().as_str() { + "direct" => run_direct(&input, &model, &answer_root).await?, + "best_of_n" | "best-of-n" | "self_consistency" => { + run_best_of_n(&input, &model, &answer_root).await? + } + "tree_of_thought" | "tree-of-thought" | "tot" => { + run_tree_of_thought(&input, &model, &answer_root, &score_root).await? + } + "graph_of_thought" | "graph-of-thought" | "got" => { + run_graph_of_thought(&input, &model, &answer_root, &aggregate_root).await? + } + other => { + return Err(format!( + "unknown pattern '{other}': expected direct, best_of_n, \ + tree_of_thought, or graph_of_thought" + )); + } + }; + output.stats.elapsed_ms = started.elapsed().as_millis(); + Ok(output) +} + +fn validate(input: &Input) -> Result<()> { + if input.question.trim().is_empty() { + return Err("question must not be empty".into()); + } + if !(1..=16).contains(&input.num_candidates) { + return Err("num_candidates must be in [1, 16]".into()); + } + if !(1..=input.num_candidates).contains(&input.beam_width) { + return Err("beam_width must be in [1, num_candidates]".into()); + } + if input.max_tokens == 0 || input.score_tokens == 0 { + return Err("token budgets must be at least 1".into()); + } + if !(input.temperature.is_finite() && (0.0..=2.0).contains(&input.temperature)) { + return Err("temperature must be in [0.0, 2.0]".into()); + } + if !(input.top_p.is_finite() && input.top_p > 0.0 && input.top_p <= 1.0) { + return Err("top_p must be in (0.0, 1.0]".into()); + } + Ok(()) +} + +async fn run_direct(input: &Input, model: &Model, root: &Context) -> Result { + let generated = generate_answer( + root.fork()?, + model, + input.max_tokens, + input.temperature, + input.top_p, + input.thinking, + None, + ) + .await?; + let candidate = candidate("direct-0", "answer", generated, None); + let mut stats = ExecutionStats { + context_forks: 1, + ..Default::default() + }; + add_candidate_stats(&mut stats, &candidate); + + Ok(Output { + pattern: "direct".into(), + final_response: candidate.response.clone(), + final_answer: candidate.answer.clone(), + selected_candidate_id: Some(candidate.id.clone()), + candidates: vec![candidate], + stats, + }) +} + +async fn run_best_of_n(input: &Input, model: &Model, root: &Context) -> Result { + let generated = generate_candidates(input, model, root, "sample", None).await?; + let candidates: Vec = generated + .into_iter() + .enumerate() + .map(|(idx, generated)| candidate(&format!("sample-{idx}"), "sample", generated, None)) + .collect(); + let selected = majority_vote_index(&candidates).unwrap_or(0); + let chosen = &candidates[selected]; + let mut stats = ExecutionStats { + context_forks: candidates.len(), + ..Default::default() + }; + for candidate in &candidates { + add_candidate_stats(&mut stats, candidate); + } + + Ok(Output { + pattern: "best_of_n".into(), + final_response: chosen.response.clone(), + final_answer: chosen.answer.clone(), + selected_candidate_id: Some(chosen.id.clone()), + candidates, + stats, + }) +} + +async fn run_tree_of_thought( + input: &Input, + model: &Model, + answer_root: &Context, + score_root: &Context, +) -> Result { + let initial = generate_candidates(input, model, answer_root, "initial", None).await?; + let mut candidates: Vec = initial + .into_iter() + .enumerate() + .map(|(idx, generated)| candidate(&format!("initial-{idx}"), "initial", generated, None)) + .collect(); + let mut stats = ExecutionStats { + context_forks: candidates.len(), + ..Default::default() + }; + for candidate in &candidates { + add_candidate_stats(&mut stats, candidate); + } + + score_candidates( + &input.question, + input.thinking, + input.score_tokens, + model, + score_root, + &mut candidates, + &mut stats, + ) + .await?; + + let mut ranked: Vec = (0..candidates.len()).collect(); + ranked.sort_by_key(|&idx| std::cmp::Reverse(candidates[idx].score.unwrap_or(0))); + ranked.truncate(input.beam_width); + + let refine_futures = ranked.iter().map(|&idx| { + let previous = candidates[idx].response.clone(); + let prompt = format!( + "Previous candidate solution:\n{previous}\n\n\ + Critique its interpretation and arithmetic, then produce a corrected \ + complete solution to the original problem." + ); + let ctx = answer_root.fork(); + async move { + generate_answer( + ctx?, + model, + input.max_tokens, + input.temperature, + input.top_p, + input.thinking, + Some(&prompt), + ) + .await + } + }); + let refined = future::join_all(refine_futures) + .await + .into_iter() + .collect::>>()?; + stats.context_forks += refined.len(); + + let first_refined = candidates.len(); + for (idx, generated) in refined.into_iter().enumerate() { + let candidate = candidate(&format!("refined-{idx}"), "refined", generated, None); + add_candidate_stats(&mut stats, &candidate); + candidates.push(candidate); + } + score_candidates( + &input.question, + input.thinking, + input.score_tokens, + model, + score_root, + &mut candidates[first_refined..], + &mut stats, + ) + .await?; + + let selected = (first_refined..candidates.len()) + .max_by_key(|&idx| candidates[idx].score.unwrap_or(0)) + .unwrap_or(0); + let chosen = &candidates[selected]; + + Ok(Output { + pattern: "tree_of_thought".into(), + final_response: chosen.response.clone(), + final_answer: chosen.answer.clone(), + selected_candidate_id: Some(chosen.id.clone()), + candidates, + stats, + }) +} + +async fn run_graph_of_thought( + input: &Input, + model: &Model, + answer_root: &Context, + aggregate_root: &Context, +) -> Result { + let generated = generate_candidates(input, model, answer_root, "proposal", None).await?; + let mut candidates: Vec = generated + .into_iter() + .enumerate() + .map(|(idx, generated)| candidate(&format!("proposal-{idx}"), "proposal", generated, None)) + .collect(); + let mut stats = ExecutionStats { + context_forks: candidates.len() + 1, + ..Default::default() + }; + for candidate in &candidates { + add_candidate_stats(&mut stats, candidate); + } + + let proposals = candidates + .iter() + .enumerate() + .map(|(idx, candidate)| format!("Candidate {}:\n{}", idx + 1, candidate.response)) + .collect::>() + .join("\n\n"); + let prompt = format!( + "Original problem:\n{}\n\nCandidate solutions:\n{}\n\n\ + Synthesize the most reliable final solution.", + input.question, proposals + ); + let generated = generate_answer( + aggregate_root.fork()?, + model, + input.max_tokens, + 0.0, + 1.0, + input.thinking, + Some(&prompt), + ) + .await?; + let aggregate = candidate("aggregate-0", "aggregate", generated, None); + add_candidate_stats(&mut stats, &aggregate); + let final_response = aggregate.response.clone(); + let final_answer = aggregate.answer.clone(); + let selected_candidate_id = Some(aggregate.id.clone()); + candidates.push(aggregate); + + Ok(Output { + pattern: "graph_of_thought".into(), + final_response, + final_answer, + selected_candidate_id, + candidates, + stats, + }) +} + +async fn generate_candidates( + input: &Input, + model: &Model, + root: &Context, + _stage: &'static str, + prompt: Option<&str>, +) -> Result> { + let contexts = (0..input.num_candidates) + .map(|_| root.fork()) + .collect::>>()?; + let futures = contexts.into_iter().map(|ctx| { + generate_answer( + ctx, + model, + input.max_tokens, + input.temperature, + input.top_p, + input.thinking, + prompt, + ) + }); + future::join_all(futures) + .await + .into_iter() + .collect::>>() +} + +async fn score_candidates( + question: &str, + thinking: bool, + max_tokens: usize, + model: &Model, + root: &Context, + candidates: &mut [Candidate], + stats: &mut ExecutionStats, +) -> Result<()> { + let futures = + candidates.iter().map(|candidate| { + let prompt = format!( + "Problem:\n{question}\n\nCandidate solution:\n{}\n\nScore:", + candidate.response + ); + let ctx = root.fork(); + async move { + generate_answer(ctx?, model, max_tokens, 0.0, 1.0, thinking, Some(&prompt)).await + } + }); + let scores = future::join_all(futures) + .await + .into_iter() + .collect::>>()?; + stats.context_forks += scores.len(); + for (candidate, generated) in candidates.iter_mut().zip(scores) { + candidate.score = parse_score(&generated.text); + stats.add_score(&generated); + } + Ok(()) +} + +async fn generate_answer( + mut ctx: Context, + model: &Model, + max_tokens: usize, + temperature: f32, + top_p: f32, + thinking: bool, + prompt: Option<&str>, +) -> Result { + if let Some(prompt) = prompt { + ctx.user(prompt); + } + ctx.cue(); + if !thinking { + ctx.append(&model.tokenizer().encode("\n\n\n\n")); + } + let sampler = if temperature <= 0.0 { + Sampler::Argmax + } else { + Sampler::TopP { + temperature, + p: top_p, + } + }; + let stops = chat::stop_tokens(model); + let mut generator = ctx.generate(sampler).max_tokens(max_tokens).stop(&stops); + let mut decoder = chat::Decoder::new(model); + let mut text = String::new(); + let mut tokens = 0; + let mut steps = 0; + + while let Some(step) = generator.next()? { + let output = step.execute().await?; + tokens += output.tokens.len(); + steps += 1; + match decoder.feed(&output.tokens)? { + chat::Event::Delta(delta) => text.push_str(&delta), + chat::Event::Done(done) => { + text = done; + break; + } + chat::Event::Idle | chat::Event::Interrupt(_) => {} + } + } + Ok(Generated { + text, + tokens, + steps, + }) +} + +fn candidate(id: &str, stage: &'static str, generated: Generated, score: Option) -> Candidate { + Candidate { + id: id.into(), + stage, + answer: extract_final_answer(&generated.text), + response: generated.text, + score, + generated_tokens: generated.tokens, + generator_steps: generated.steps, + } +} + +fn add_candidate_stats(stats: &mut ExecutionStats, candidate: &Candidate) { + stats.generated_tokens += candidate.generated_tokens; + stats.generator_steps += candidate.generator_steps; + stats.generation_calls += 1; +} + +fn majority_vote_index(candidates: &[Candidate]) -> Option { + let mut counts: HashMap<&str, usize> = HashMap::new(); + for answer in candidates + .iter() + .filter_map(|candidate| candidate.answer.as_deref()) + { + *counts.entry(answer).or_default() += 1; + } + candidates + .iter() + .enumerate() + .filter_map(|(idx, candidate)| { + let answer = candidate.answer.as_deref()?; + Some((idx, counts.get(answer).copied().unwrap_or(0))) + }) + .max_by_key(|(idx, count)| (*count, std::cmp::Reverse(*idx))) + .map(|(idx, _)| idx) +} + +fn extract_final_answer(text: &str) -> Option { + let lower = text.to_ascii_lowercase(); + let marker = "final answer:"; + let pos = lower.rfind(marker)?; + let tail = &text[pos + marker.len()..]; + extract_last_number(tail) +} + +fn extract_last_number(text: &str) -> Option { + let mut tokens = Vec::new(); + let mut current = String::new(); + for ch in text.chars() { + if ch.is_ascii_digit() || matches!(ch, '-' | '+' | '.' | ',' | '/') { + current.push(ch); + } else if !current.is_empty() { + tokens.push(std::mem::take(&mut current)); + } + } + if !current.is_empty() { + tokens.push(current); + } + tokens + .into_iter() + .rev() + .find_map(|token| normalize_number(&token)) +} + +fn normalize_number(raw: &str) -> Option { + let mut value = raw + .trim_matches(|c: char| matches!(c, '+' | '.' | ',' | '/')) + .replace(',', ""); + if value.is_empty() || !value.chars().any(|c| c.is_ascii_digit()) { + return None; + } + if value.ends_with(".0") { + value.truncate(value.len() - 2); + } + Some(value) +} + +fn parse_score(text: &str) -> Option { + text.split(|c: char| !c.is_ascii_digit()) + .filter_map(|part| part.parse::().ok()) + .find(|score| (1..=10).contains(score)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn extracts_marked_numeric_answer() { + assert_eq!( + extract_final_answer("Work.\nFinal Answer: $1,234.0"), + Some("1234".into()) + ); + } + + #[test] + fn rejects_unmarked_numeric_answer() { + assert_eq!(extract_final_answer("Work.\nThe answer is 1234."), None); + } + + #[test] + fn rejects_final_answer_without_colon() { + assert_eq!(extract_final_answer("Work.\nFinal Answer is 1234."), None); + } + + #[test] + fn vote_prefers_first_candidate_on_tie() { + let make = |id: &str, answer: &str| Candidate { + id: id.into(), + stage: "sample", + response: String::new(), + answer: Some(answer.into()), + score: None, + generated_tokens: 0, + generator_steps: 0, + }; + let candidates = vec![make("a", "1"), make("b", "2")]; + assert_eq!(majority_vote_index(&candidates), Some(0)); + } +} diff --git a/inferlets/reasoning-direct/Cargo.toml b/inferlets/reasoning-direct/Cargo.toml new file mode 100644 index 000000000..369049516 --- /dev/null +++ b/inferlets/reasoning-direct/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "reasoning-direct" +version = "0.1.0" +edition = "2024" + +[lib] +crate-type = ["cdylib"] + +[dependencies] +inferlet = { path = "../../sdk/rust/inferlet" } +reasoning-core = { path = "../reasoning-core" } + +[profile.release] +lto = true diff --git a/inferlets/reasoning-direct/Pie.toml b/inferlets/reasoning-direct/Pie.toml new file mode 100644 index 000000000..f56c0acd4 --- /dev/null +++ b/inferlets/reasoning-direct/Pie.toml @@ -0,0 +1,19 @@ +[package] +name = "reasoning-direct" +version = "0.1.0" +description = "Direct text-completion baseline for reasoning benchmarks" +authors = ["Pie Team"] + +[runtime] +core = "^0.2.0" +mcp = "^0.2.0" + +[parameters] +question = {type = "string", description = "Math word problem to solve"} +num_candidates = {type = "int", optional = true, description = "Ignored by direct mode; accepted for runner compatibility"} +beam_width = {type = "int", optional = true, description = "Ignored by direct mode; accepted for runner compatibility"} +max_tokens = {type = "int", optional = true, description = "Maximum answer tokens per generation (default: 256)"} +score_tokens = {type = "int", optional = true, description = "Ignored by direct mode; accepted for runner compatibility"} +temperature = {type = "float", optional = true, description = "Sampling temperature (default: 0.7)"} +top_p = {type = "float", optional = true, description = "Candidate top-p value (default: 0.95)"} +thinking = {type = "bool", optional = true, description = "Allow model thinking blocks (default: false)"} diff --git a/inferlets/reasoning-direct/src/lib.rs b/inferlets/reasoning-direct/src/lib.rs new file mode 100644 index 000000000..82ad7413b --- /dev/null +++ b/inferlets/reasoning-direct/src/lib.rs @@ -0,0 +1,10 @@ +//! Method-isolated Direct baseline inferlet for reasoning benchmarks. + +use inferlet::Result; +use reasoning_core::{Input, Output}; + +#[inferlet::main] +async fn main(mut input: Input) -> Result { + input.force_pattern("direct"); + reasoning_core::run(input).await +} diff --git a/inferlets/reasoning-graph-of-thought/Cargo.toml b/inferlets/reasoning-graph-of-thought/Cargo.toml new file mode 100644 index 000000000..c9dab8380 --- /dev/null +++ b/inferlets/reasoning-graph-of-thought/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "reasoning-graph-of-thought" +version = "0.1.0" +edition = "2024" + +[lib] +crate-type = ["cdylib"] + +[dependencies] +inferlet = { path = "../../sdk/rust/inferlet" } +reasoning-core = { path = "../reasoning-core" } + +[profile.release] +lto = true diff --git a/inferlets/reasoning-graph-of-thought/Pie.toml b/inferlets/reasoning-graph-of-thought/Pie.toml new file mode 100644 index 000000000..6943e0012 --- /dev/null +++ b/inferlets/reasoning-graph-of-thought/Pie.toml @@ -0,0 +1,19 @@ +[package] +name = "reasoning-graph-of-thought" +version = "0.1.0" +description = "Graph-of-Thought aggregation inferlet for reasoning benchmarks" +authors = ["Pie Team"] + +[runtime] +core = "^0.2.0" +mcp = "^0.2.0" + +[parameters] +question = {type = "string", description = "Math word problem to solve"} +num_candidates = {type = "int", optional = true, description = "Proposal count (default: 4)"} +beam_width = {type = "int", optional = true, description = "Accepted for runner compatibility"} +max_tokens = {type = "int", optional = true, description = "Maximum answer tokens per generation (default: 256)"} +score_tokens = {type = "int", optional = true, description = "Accepted for runner compatibility"} +temperature = {type = "float", optional = true, description = "Proposal sampling temperature (default: 0.7)"} +top_p = {type = "float", optional = true, description = "Proposal top-p value (default: 0.95)"} +thinking = {type = "bool", optional = true, description = "Allow model thinking blocks (default: false)"} diff --git a/inferlets/reasoning-graph-of-thought/src/lib.rs b/inferlets/reasoning-graph-of-thought/src/lib.rs new file mode 100644 index 000000000..e9066526e --- /dev/null +++ b/inferlets/reasoning-graph-of-thought/src/lib.rs @@ -0,0 +1,10 @@ +//! Method-isolated Graph-of-Thought inferlet for reasoning benchmarks. + +use inferlet::Result; +use reasoning_core::{Input, Output}; + +#[inferlet::main] +async fn main(mut input: Input) -> Result { + input.force_pattern("graph_of_thought"); + reasoning_core::run(input).await +} diff --git a/inferlets/reasoning-tree-of-thought/Cargo.toml b/inferlets/reasoning-tree-of-thought/Cargo.toml new file mode 100644 index 000000000..493b20d08 --- /dev/null +++ b/inferlets/reasoning-tree-of-thought/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "reasoning-tree-of-thought" +version = "0.1.0" +edition = "2024" + +[lib] +crate-type = ["cdylib"] + +[dependencies] +inferlet = { path = "../../sdk/rust/inferlet" } +reasoning-core = { path = "../reasoning-core" } + +[profile.release] +lto = true diff --git a/inferlets/reasoning-tree-of-thought/Pie.toml b/inferlets/reasoning-tree-of-thought/Pie.toml new file mode 100644 index 000000000..f07fa752d --- /dev/null +++ b/inferlets/reasoning-tree-of-thought/Pie.toml @@ -0,0 +1,19 @@ +[package] +name = "reasoning-tree-of-thought" +version = "0.1.0" +description = "Tree-of-Thought inferlet for reasoning benchmarks" +authors = ["Pie Team"] + +[runtime] +core = "^0.2.0" +mcp = "^0.2.0" + +[parameters] +question = {type = "string", description = "Math word problem to solve"} +num_candidates = {type = "int", optional = true, description = "Initial candidate count (default: 4)"} +beam_width = {type = "int", optional = true, description = "Number of candidates to refine (default: 2)"} +max_tokens = {type = "int", optional = true, description = "Maximum answer tokens per generation (default: 256)"} +score_tokens = {type = "int", optional = true, description = "Maximum tokens per model-scoring call (default: 16)"} +temperature = {type = "float", optional = true, description = "Candidate sampling temperature (default: 0.7)"} +top_p = {type = "float", optional = true, description = "Candidate top-p value (default: 0.95)"} +thinking = {type = "bool", optional = true, description = "Allow model thinking blocks (default: false)"} diff --git a/inferlets/reasoning-tree-of-thought/src/lib.rs b/inferlets/reasoning-tree-of-thought/src/lib.rs new file mode 100644 index 000000000..ba607963b --- /dev/null +++ b/inferlets/reasoning-tree-of-thought/src/lib.rs @@ -0,0 +1,10 @@ +//! Method-isolated Tree-of-Thought inferlet for reasoning benchmarks. + +use inferlet::Result; +use reasoning_core::{Input, Output}; + +#[inferlet::main] +async fn main(mut input: Input) -> Result { + input.force_pattern("tree_of_thought"); + reasoning_core::run(input).await +} diff --git a/tests/inferlets/test_reasoning_benchmark.py b/tests/inferlets/test_reasoning_benchmark.py index 4d98a45ae..eefd8876b 100644 --- a/tests/inferlets/test_reasoning_benchmark.py +++ b/tests/inferlets/test_reasoning_benchmark.py @@ -1,9 +1,17 @@ -"""E2E interface test for the unified reasoning benchmark inferlet.""" +"""E2E interface tests for benchmark reasoning inferlets.""" import json from conftest import run_inferlet, run_tests +REASONING_INFERLETS = { + "reasoning-direct": "direct", + "reasoning-best-of-n": "best_of_n", + "reasoning-tree-of-thought": "tree_of_thought", + "reasoning-graph-of-thought": "graph_of_thought", +} + + async def test_reasoning_benchmark(client, args): for pattern in ( "direct", @@ -32,5 +40,46 @@ async def test_reasoning_benchmark(client, args): assert result["stats"]["context_forks"] >= 1 +async def test_reasoning_base(client, args): + output = await run_inferlet( + client, + "reasoning-base", + { + "prompt": "Say the word ready.", + "max_tokens": 16, + }, + timeout=args.timeout, + ) + result = json.loads(output) + assert isinstance(result["completion"], str) + assert result["stats"]["generator_steps"] >= 0 + assert result["stats"]["generated_tokens"] >= 0 + + +async def test_separate_reasoning_inferlets(client, args): + for inferlet, pattern in REASONING_INFERLETS.items(): + output = await run_inferlet( + client, + inferlet, + { + "pattern": "direct", + "question": "Mia has 7 apples and buys 5 more. How many apples?", + "num_candidates": 2, + "beam_width": 1, + "max_tokens": 48, + "score_tokens": 8, + }, + timeout=args.timeout, + ) + result = json.loads(output) + assert result["pattern"] == pattern + assert isinstance(result["candidates"], list) + assert result["candidates"] + assert result["stats"]["generation_calls"] >= 1 + assert result["stats"]["context_forks"] >= 1 + + if __name__ == "__main__": - run_tests([test_reasoning_benchmark]) + run_tests( + [test_reasoning_benchmark, test_reasoning_base, test_separate_reasoning_inferlets] + )