mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-27 21:05:34 +00:00
orchestrator: SFT data generation + verification pipeline
Rejection-sampling data gen for the router: run anonymized ToolOrchestra rollouts over a task pool, verify correctness (judge with backoff, math normalization, essay/code gates), and keep only correct, well-formed, actually-routing trajectories. Includes the clean-gate that strips self-answering and control-token garble, cost-aware reward, dataset loaders, unified trajectory serialization, and naming (single source of truth for dataset names). ~76% of rollouts are dropped by design.
This commit is contained in:
Andrew Park
@@ -0,0 +1,399 @@
|
||||
#!/usr/bin/env python
|
||||
"""Build orchestrator SFT cold-start data by **base Qwen3-8B self-sampling** (v1).
|
||||
|
||||
The orchestrator *is* the local Qwen3-8B. We roll the base (un-trained) model out
|
||||
over the reasoning-SFT task pool (``load_sft_tasks``: GeneralThought + OpenThoughts
|
||||
code/math/science), verify each trajectory's final answer against the gold
|
||||
(``verify.make_verifier`` — math/code checkers + Gemini judge fallback), and keep
|
||||
the cheapest-correct trajectory per task. Those passing trajectories become the
|
||||
``conversations`` JSONL the SFT trainer consumes — i.e. the model learns from its
|
||||
own successful rollouts (rejection-sampling cold-start, STaR-style).
|
||||
|
||||
This is **v1** (base self-sampling): point ``--orchestrator-endpoint`` at the vLLM
|
||||
server hosting the *base* ``Qwen/Qwen3-8B``.
|
||||
|
||||
For **v2**, train on the v1 data, serve the v1 checkpoint with vLLM, then re-run
|
||||
this exact driver with ``--orchestrator-endpoint`` pointed at the v1 checkpoint's
|
||||
endpoint, ``--orchestrator-model`` set to the checkpoint name, and
|
||||
``--samples-per-task 8`` — the (now stronger) policy self-samples a v2 dataset.
|
||||
|
||||
The math/coder specialist tools are only wired when ``--math-endpoint`` /
|
||||
``--coder-endpoint`` are passed (those are separately-served vLLM specialists);
|
||||
otherwise they're omitted so the orchestrator only sees the tools it can actually
|
||||
call.
|
||||
|
||||
Example (v1, base self-sampling):
|
||||
.venv/bin/python scripts/orchestrator/build_orchestrator_sft.py \
|
||||
--orchestrator-label qwen \
|
||||
--orchestrator-endpoint http://localhost:8001/v1 \
|
||||
--orchestrator-model qwen3-8b \
|
||||
--samples-per-task 8 --max-keep-per-task 1
|
||||
|
||||
Example (v2, re-point at the v1 checkpoint):
|
||||
.venv/bin/python scripts/orchestrator/build_orchestrator_sft.py \
|
||||
--orchestrator-label qwen \
|
||||
--orchestrator-endpoint http://localhost:8010/v1 \
|
||||
--orchestrator-model orchestrator-sft-v1 \
|
||||
--samples-per-task 8
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog
|
||||
from openjarvis.agents.hybrid.toolorchestra.rollout import run_unified_rollout
|
||||
from openjarvis.agents.hybrid.toolorchestra.unified import (
|
||||
make_call_orchestrator,
|
||||
make_dispatch,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import (
|
||||
load_sft_tasks,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.naming import raw_dir_name
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
|
||||
generate_sft_dataset,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.verify import make_verifier
|
||||
|
||||
|
||||
def main(argv: Optional[list[str]] = None) -> int:
|
||||
p = argparse.ArgumentParser(
|
||||
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
|
||||
)
|
||||
# --out is a TAG, not a path. The run dir is always
|
||||
# <data-root>/raw/{label}-{month}-{day}-{year}-{hhmm}{am|pm}[-{tag}]/ and the
|
||||
# file inside is always data.jsonl — see the run_dir block below. Use --out
|
||||
# only to distinguish two runs of the same orchestrator (e.g. --out balanced50).
|
||||
p.add_argument(
|
||||
"--out",
|
||||
default=None,
|
||||
help="Optional tag appended to the run dir name (not a path).",
|
||||
)
|
||||
# The orchestrator == the local Qwen3-8B self-sampling over its own rollouts.
|
||||
p.add_argument(
|
||||
"--orchestrator-endpoint",
|
||||
default="http://localhost:8001/v1",
|
||||
help="OpenAI-compatible vLLM base URL serving the orchestrator.",
|
||||
)
|
||||
p.add_argument("--orchestrator-model", default="qwen3-8b")
|
||||
p.add_argument("--orchestrator-api-key", default="EMPTY")
|
||||
# Provenance stamps: when set, every written record carries these fields so
|
||||
# the JSONL is self-identifying once pooled across model families (gemma vs
|
||||
# qwen). Default None -> no stamping (existing qwen runs unaffected).
|
||||
p.add_argument(
|
||||
"--gen-model",
|
||||
default=None,
|
||||
help="Full HF id of the generating orchestrator, stamped as "
|
||||
"record['gen_model'] (e.g. google/gemma-4-26B-A4B-it).",
|
||||
)
|
||||
p.add_argument(
|
||||
"--orchestrator-label",
|
||||
default=None,
|
||||
help="Short label stamped as record['orchestrator_model'] (e.g. gemma-4-26b).",
|
||||
)
|
||||
# Local OSS model endpoints; repeatable "model_id=base_url" (e.g.
|
||||
# "Qwen/Qwen3.5-9B=http://localhost:8001/v1"). Unmapped local models are
|
||||
# still listed but served unconfigured (base_url=None).
|
||||
p.add_argument(
|
||||
"--local-endpoint",
|
||||
action="append",
|
||||
default=[],
|
||||
metavar="MODEL_ID=URL",
|
||||
help="Local model id -> vLLM base URL (repeatable).",
|
||||
)
|
||||
p.add_argument(
|
||||
"--max-tasks",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Cap on tasks (default: all of load_sft_tasks()).",
|
||||
)
|
||||
p.add_argument(
|
||||
"--skip-task-ids-from",
|
||||
action="append",
|
||||
default=[],
|
||||
metavar="GLOB",
|
||||
help="Glob of prior data.jsonl files; skip task_ids already "
|
||||
"generated there so this run only does unseen prompts "
|
||||
"(resume). Repeatable; applied before sharding.",
|
||||
)
|
||||
p.add_argument("--samples-per-task", type=int, default=8)
|
||||
p.add_argument("--max-keep-per-task", type=int, default=1)
|
||||
p.add_argument("--max-turns", type=int, default=8)
|
||||
p.add_argument("--temperature", type=float, default=0.7)
|
||||
# Orchestrator completion cap per turn. The library default (4096) intermittently
|
||||
# truncates the final answer mid-sentence on longer reasoning turns; bump it so
|
||||
# the trace ends cleanly. Raise further (e.g. 16384) if traces still cut off.
|
||||
p.add_argument(
|
||||
"--max-tokens",
|
||||
type=int,
|
||||
default=8192,
|
||||
help="Per-turn completion cap for the orchestrator (default 8192).",
|
||||
)
|
||||
# Number of tasks rolled out concurrently against vLLM. Each in-flight task
|
||||
# issues its samples sequentially, so ~concurrency requests hit the server at
|
||||
# once. ~30-50 is comfortable on 1xL40S (prefix caching + continuous batching).
|
||||
p.add_argument(
|
||||
"--concurrency",
|
||||
type=int,
|
||||
default=32,
|
||||
help="Tasks rolled out in parallel (default 32).",
|
||||
)
|
||||
# Balanced-smoke pull: cap//4 from each of GeneralThought + OpenThoughts
|
||||
# code/math/science instead of the GeneralThought-only fast cap path. Use for a
|
||||
# representative smoke; the real run omits --max-tasks for the full 8K balanced set.
|
||||
p.add_argument(
|
||||
"--balanced",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="Draw an EVEN cross-domain sample (GeneralThought + "
|
||||
"OpenThoughts code/math/science). Default ON — pass "
|
||||
"--no-balanced for the old GeneralThought-only skew.",
|
||||
)
|
||||
# Data-parallel sharding: split the (deterministic, seed-42) task list across
|
||||
# N independent driver processes, each pointed at its own orchestrator vLLM
|
||||
# replica. Shard i takes tasks[i::N] (a strided slice, so every shard stays
|
||||
# domain-balanced). Each writes its own --out file; concatenate the shard
|
||||
# JSONLs afterward. Lets the GPU-bound orchestrator scale across idle GPUs.
|
||||
p.add_argument(
|
||||
"--shard-index",
|
||||
type=int,
|
||||
default=0,
|
||||
help="This shard's index in [0, shard-count).",
|
||||
)
|
||||
p.add_argument(
|
||||
"--shard-count",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Total number of shards (default 1 = no sharding).",
|
||||
)
|
||||
# Throughput knob: stop sampling a task as soon as max-keep-per-task passing
|
||||
# trajectories are found, instead of always running all --samples-per-task.
|
||||
# ~3-4x faster when most tasks solve early, but keeps the *first* passers
|
||||
# rather than the *cheapest* of N (drops the cost-optimisation signal).
|
||||
p.add_argument(
|
||||
"--stop-at-keep",
|
||||
action="store_true",
|
||||
help="Short-circuit a task once max-keep passing samples found.",
|
||||
)
|
||||
# Anonymize model experts (opaque random labels, uniform description, no cost
|
||||
# line, shuffled order) so the policy can't route on a model's name/position/
|
||||
# cost. The anon->real map is saved per record (metrics.anon_map) for analysis.
|
||||
p.add_argument(
|
||||
"--anonymize-experts",
|
||||
action="store_true",
|
||||
help="Hide expert identity (random labels + shuffle) when routing.",
|
||||
)
|
||||
# By default we keep EVERY rolled-out trajectory (correct + incorrect), each
|
||||
# tagged with ``correct`` (verifier verdict) and ``kept`` (the cheapest-correct
|
||||
# sample the rejection sampler would pick). Lets you compute accuracy / inspect
|
||||
# failures from the JSONL directly. --rejection-only restores the old
|
||||
# drop-the-failures behaviour (only cheapest-correct written).
|
||||
p.add_argument(
|
||||
"--rejection-only",
|
||||
action="store_true",
|
||||
help="Drop incorrect rollouts; write only the cheapest-correct "
|
||||
"sample per task (the original behaviour).",
|
||||
)
|
||||
# Rollouts call shell_exec / file_write with model-chosen relative paths
|
||||
# (e.g. ``solution.py``), which otherwise land in the repo root. Run them
|
||||
# from a throwaway scratch dir so generated files never dirty the tree.
|
||||
# gitignored via the existing ``scratch/`` rule.
|
||||
p.add_argument(
|
||||
"--scratch-dir",
|
||||
default="scratch/sft-rollouts",
|
||||
help="CWD for rollouts; stray tool-written files go here.",
|
||||
)
|
||||
args = p.parse_args(argv)
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
||||
# Quiet the per-request HTTP / dataset-stream spam so the run log stays
|
||||
# readable (rollout calls and dataset shards otherwise flood it).
|
||||
for _noisy in (
|
||||
"httpx",
|
||||
"httpcore",
|
||||
"urllib3",
|
||||
"datasets",
|
||||
"fsspec",
|
||||
"huggingface_hub",
|
||||
"openai",
|
||||
):
|
||||
logging.getLogger(_noisy).setLevel(logging.WARNING)
|
||||
|
||||
# Every run lands in its OWN folder under <data-root>/raw/, sharing its stamp
|
||||
# with every file later carved from it — only the stage word differs:
|
||||
#
|
||||
# raw/qwen-july-7-2026-0553pm/data.jsonl <- every rollout, incl. failures
|
||||
# sft/qwen-clean-july-7-2026-0553pm.jsonl <- reject-sampled from it
|
||||
#
|
||||
# so any curated file traces back to its generation run by eye. raw/ sits ABOVE
|
||||
# the sft/rl fork on purpose: SFT keeps only correct+clean rows, but GRPO needs
|
||||
# the failures, and both read the same pool. Naming lives in sft_data.naming.
|
||||
# --out is an optional extra TAG, not a path; the file inside is always data.jsonl.
|
||||
#
|
||||
# OJ_DATA_ROOT keeps the data OUT of the git checkout (repo-relative default so
|
||||
# a fresh clone still works); this workspace points it at the experiments tree.
|
||||
data_root = Path(os.getenv("OJ_DATA_ROOT", "data/orchestrator"))
|
||||
prefix = args.orchestrator_label or "orch"
|
||||
tag = Path(args.out).stem if args.out else ""
|
||||
base = data_root / "raw" / raw_dir_name(prefix, tag=tag)
|
||||
run_dir = base.resolve()
|
||||
n = 1
|
||||
while run_dir.exists(): # same name+minute (parallel shards) -> disambiguate
|
||||
n += 1
|
||||
run_dir = base.with_name(f"{base.name}-{n}").resolve()
|
||||
label = run_dir.name
|
||||
out_p = run_dir / "data.jsonl"
|
||||
lock_p = out_p.with_suffix(out_p.suffix + ".lock")
|
||||
out_p.parent.mkdir(parents=True, exist_ok=True)
|
||||
lock_p.write_text(f"{time.strftime('%m-%d-%I%M%p').lower()} pid={os.getpid()}")
|
||||
args.out = str(out_p)
|
||||
import atexit
|
||||
|
||||
atexit.register(lambda: lock_p.exists() and lock_p.unlink())
|
||||
logging.info("Run dir: %s", run_dir)
|
||||
|
||||
# Enrich Braintrust rollout traces with run-level provenance (no-op-safe;
|
||||
# tracing.run_context() reads these). setdefault so an explicit env wins.
|
||||
os.environ.setdefault("OJ_RUN_LABEL", label)
|
||||
os.environ.setdefault("OJ_GEN_MODEL", args.gen_model or args.orchestrator_model)
|
||||
os.environ.setdefault("OJ_RUN_STAGE", "smoke" if args.max_tasks else "prod")
|
||||
os.environ["OJ_CFG_TEMPERATURE"] = str(args.temperature)
|
||||
os.environ["OJ_CFG_MAX_TURNS"] = str(args.max_turns)
|
||||
os.environ["OJ_CFG_ANONYMIZE"] = str(bool(args.anonymize_experts))
|
||||
os.environ["OJ_CFG_REJECTION_ONLY"] = str(bool(args.rejection_only))
|
||||
|
||||
# Sandbox the rollouts: chdir into a scratch dir so any file a rollout writes
|
||||
# (shell_exec redirects, file_write with a relative path) lands there instead
|
||||
# of the repo root. (--out is already absolute via run_dir.resolve() above.)
|
||||
scratch = Path(args.scratch_dir).resolve()
|
||||
scratch.mkdir(parents=True, exist_ok=True)
|
||||
os.chdir(scratch)
|
||||
logging.info("Rollout scratch dir (cwd): %s", scratch)
|
||||
|
||||
local_endpoints = {}
|
||||
for item in args.local_endpoint:
|
||||
model_id, _, url = item.partition("=")
|
||||
if model_id and url:
|
||||
local_endpoints[model_id] = url
|
||||
tools = orchestrator_catalog(local_endpoints=local_endpoints or None)
|
||||
# Drop the no-op ``think`` scratchpad tool: it adds turns/cost and is a
|
||||
# harness-leakage vector (the model narrates the routing rule inside a think
|
||||
# call, which the reasoning-scrub can't reach). The model still reasons in
|
||||
# its own <think> blocks — it just can't emit reasoning AS a tool call.
|
||||
tools = [t for t in tools if t.name != "think"]
|
||||
logging.info("Tool catalog (%d): %s", len(tools), [t.name for t in tools])
|
||||
|
||||
call_orch = make_call_orchestrator(
|
||||
args.orchestrator_model,
|
||||
base_url=args.orchestrator_endpoint,
|
||||
api_key=args.orchestrator_api_key,
|
||||
temperature=args.temperature,
|
||||
max_tokens=args.max_tokens,
|
||||
)
|
||||
dispatch = make_dispatch({})
|
||||
|
||||
# Build the provenance stamp once (None if neither flag given).
|
||||
record_extra = {}
|
||||
if args.gen_model:
|
||||
record_extra["gen_model"] = args.gen_model
|
||||
if args.orchestrator_label:
|
||||
record_extra["orchestrator_model"] = args.orchestrator_label
|
||||
record_extra = record_extra or None
|
||||
|
||||
def rollout_fn(task):
|
||||
try:
|
||||
return run_unified_rollout(
|
||||
task.instruction,
|
||||
tools,
|
||||
call_orchestrator=call_orch,
|
||||
dispatch=dispatch,
|
||||
max_turns=args.max_turns,
|
||||
anonymize=args.anonymize_experts,
|
||||
)
|
||||
except Exception as exc: # network/key failures shouldn't kill the run
|
||||
logging.warning("rollout failed for %s: %s", task.task_id, exc)
|
||||
return None
|
||||
|
||||
# cap= caps the task count for smoke runs; --balanced makes that small sample
|
||||
# cross-domain (GeneralThought + OpenThoughts code/math/science) instead of the
|
||||
# GeneralThought-only fast path. The real run omits --max-tasks (full 8K balanced).
|
||||
tasks = load_sft_tasks(cap=args.max_tasks, balanced=args.balanced)
|
||||
# Resume on unseen prompts: drop task_ids already generated in prior runs
|
||||
# (per-track seen set via --skip-task-ids-from). Filter BEFORE sharding so the
|
||||
# remaining unseen tasks stride evenly and never re-cover finished prompts.
|
||||
if args.skip_task_ids_from:
|
||||
import glob as _glob
|
||||
|
||||
skip_ids = set()
|
||||
for pattern in args.skip_task_ids_from:
|
||||
for fp in _glob.glob(pattern):
|
||||
try:
|
||||
with open(fp) as fh:
|
||||
for line in fh:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
tid = json.loads(line).get("task_id")
|
||||
except Exception:
|
||||
continue
|
||||
if tid:
|
||||
skip_ids.add(tid)
|
||||
except OSError:
|
||||
continue
|
||||
if skip_ids:
|
||||
n_before = len(tasks)
|
||||
tasks = [t for t in tasks if t.task_id not in skip_ids]
|
||||
logging.info(
|
||||
"Resume: %d done task_ids -> skipping; %d of %d tasks remain",
|
||||
len(skip_ids),
|
||||
len(tasks),
|
||||
n_before,
|
||||
)
|
||||
if args.shard_count > 1:
|
||||
n_all = len(tasks)
|
||||
tasks = tasks[args.shard_index :: args.shard_count]
|
||||
logging.info(
|
||||
"Shard %d/%d -> %d of %d tasks",
|
||||
args.shard_index,
|
||||
args.shard_count,
|
||||
len(tasks),
|
||||
n_all,
|
||||
)
|
||||
from collections import Counter as _Counter
|
||||
|
||||
_dom = _Counter(getattr(t, "domain", "unknown") for t in tasks)
|
||||
logging.info(
|
||||
"Loaded %d SFT tasks | balanced=%s | domain split: %s",
|
||||
len(tasks),
|
||||
args.balanced,
|
||||
", ".join(f"{d}={n}" for d, n in sorted(_dom.items())),
|
||||
)
|
||||
|
||||
stats = generate_sft_dataset(
|
||||
args.out,
|
||||
tasks=tasks,
|
||||
tools=tools,
|
||||
rollout_fn=rollout_fn,
|
||||
verify_fn=make_verifier(),
|
||||
samples_per_task=args.samples_per_task,
|
||||
max_keep_per_task=args.max_keep_per_task,
|
||||
reward_fn=lambda r: -r.cost_usd, # cheapest-correct gets highest reward
|
||||
concurrency=args.concurrency,
|
||||
stop_at_keep=args.stop_at_keep,
|
||||
keep_all=not args.rejection_only,
|
||||
record_extra=record_extra,
|
||||
)
|
||||
print(json.dumps(stats, indent=2))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,29 @@
|
||||
"""SFT-data generation for the orchestrator cold-start.
|
||||
|
||||
Execution-grounded rejection sampling: for each reasoning task (GeneralThought +
|
||||
OpenThoughts, via :func:`~...sft_data.datasets.load_sft_tasks`), roll out a
|
||||
teacher orchestrator over the unified tool catalog N times, verify each
|
||||
trajectory, keep the cheapest passing one(s), and serialize them into the
|
||||
``<tool_call>`` ``conversations`` JSONL that
|
||||
:class:`~openjarvis.learning.intelligence.orchestrator.sft_trainer.OrchestratorSFTDataset`
|
||||
loads directly.
|
||||
|
||||
Pipeline::
|
||||
|
||||
reasoning task -> N teacher rollouts -> verify -> keep cheapest passing
|
||||
-> conversations JSONL
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
|
||||
generate_sft_dataset,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import (
|
||||
trajectory_to_record,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"generate_sft_dataset",
|
||||
"trajectory_to_record",
|
||||
]
|
||||
@@ -0,0 +1,419 @@
|
||||
"""Loaders for the orchestrator reasoning-SFT task sources.
|
||||
|
||||
Two HuggingFace reasoning datasets give us verifiable question/answer tasks for
|
||||
the small orchestrator's cold-start SFT and the GRPO prompt pool:
|
||||
|
||||
- ``natolambert/GeneralThought-430K-filtered`` — open reasoning traces scraped
|
||||
from gr.inc. Each row carries a ``question``, a short ``reference_answer``
|
||||
(the gold), a long ``model_answer`` (R1's full solution), and provenance in
|
||||
``question_source`` / ``task``. There is **no** clean single "category" field;
|
||||
we derive a coarse ``domain`` (math / medical / code / chat / misc) from
|
||||
``question_source`` so the verifier can pick the right checker.
|
||||
|
||||
- ``open-thoughts/OpenThoughts3-1.2M`` — OpenThoughts3 distillation set. Each
|
||||
row has a ``domain`` in {code, math, science}, a ``source``, a ``difficulty``,
|
||||
and a ``conversations`` list ``[{from: human, value}, {from: gpt, value}]``.
|
||||
The human turn is the question; the gpt turn is a ``<think>…</think>`` trace
|
||||
followed by the final solution, from which we extract the gold answer.
|
||||
|
||||
Mirrors the style of the sibling ``hotpotqa.py`` / ``toolscale.py`` loaders: a
|
||||
plain dataclass plus loaders that accept a ``source=`` iterable override so the
|
||||
normalization path is exercised offline with no network. Network imports of
|
||||
``datasets`` are kept lazy (inside the function).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Iterable, Iterator, List, Optional
|
||||
|
||||
GENERALTHOUGHT_ID = "natolambert/GeneralThought-430K-filtered"
|
||||
OPENTHOUGHTS_ID = "open-thoughts/OpenThoughts3-1.2M"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Task:
|
||||
task_id: str
|
||||
question: str
|
||||
answer: str
|
||||
domain: str # coarse area: math / code / science / medical / chat / misc
|
||||
difficulty: str = "" # OpenThoughts difficulty tier (GeneralThought has none)
|
||||
dataset: str = "" # source dataset: GeneralThought / OpenThoughts3
|
||||
subsector: str = "" # fine source: NuminaMath / NHSQA / TACO / glaive / ...
|
||||
|
||||
@property
|
||||
def instruction(self) -> str:
|
||||
return self.question
|
||||
|
||||
|
||||
# --- GeneralThought -----------------------------------------------------------
|
||||
|
||||
# question_source -> coarse domain. GeneralThought has no literal category field;
|
||||
# the source string is the most reliable signal (NuminaMath is math, NHSQA is
|
||||
# medical, TACO/glaive are code, oasst1/lmsys are open chat, everything else misc).
|
||||
_SOURCE_DOMAIN = (
|
||||
("numina", "math"),
|
||||
("math", "math"),
|
||||
("nhsqa", "medical"),
|
||||
("medical", "medical"),
|
||||
("medicine", "medical"),
|
||||
("taco", "code"),
|
||||
("code", "code"),
|
||||
("glaive", "code"),
|
||||
("oasst", "chat"),
|
||||
("lmsys", "chat"),
|
||||
("chat", "chat"),
|
||||
)
|
||||
|
||||
|
||||
def _generalthought_domain(row: Dict[str, Any]) -> str:
|
||||
src = str(row.get("question_source") or "").lower()
|
||||
task = str(row.get("task") or "").lower()
|
||||
hay = f"{src} {task}"
|
||||
for needle, dom in _SOURCE_DOMAIN:
|
||||
if needle in hay:
|
||||
return dom
|
||||
return "misc"
|
||||
|
||||
|
||||
def _normalize_generalthought(row: Dict[str, Any], *, index: int = 0) -> Optional[Task]:
|
||||
question = str(row.get("question") or "").strip()
|
||||
# Prefer the short reference answer (exact-match-able); fall back to the
|
||||
# long model_answer when no reference is present.
|
||||
answer = str(row.get("reference_answer") or "").strip()
|
||||
if not answer:
|
||||
answer = str(row.get("model_answer") or "").strip()
|
||||
if not question or not answer:
|
||||
return None
|
||||
domain = _generalthought_domain(row)
|
||||
# Same unwinnable-gold guards as the OpenThoughts path — GeneralThought also
|
||||
# carries NuminaMath proof problems whose "answer" is just a QED marker, and
|
||||
# the math verifier has no judge to fall back on. (Audit 2026-07-12: these
|
||||
# were slipping through because the filter only lived on the other loader.)
|
||||
if _gold_is_unusable(answer, domain):
|
||||
return None
|
||||
task_id = str(row.get("question_id") or f"generalthought-{index}")
|
||||
return Task(
|
||||
task_id=task_id,
|
||||
question=question,
|
||||
answer=answer,
|
||||
domain=domain,
|
||||
difficulty=str(row.get("difficulty") or "").strip(), # usually absent for GT
|
||||
dataset="GeneralThought",
|
||||
subsector=str(row.get("question_source") or "").strip(),
|
||||
)
|
||||
|
||||
|
||||
# Domains whose questions are *lookup*-shaped, not *reasoning*-shaped ("What is
|
||||
# idiopathic anaphylaxis?"). The orchestrator correctly answers these with a
|
||||
# single web_search and never delegates to a model expert — so the trajectory is
|
||||
# rejected by the clean gate ("NEVER routed to a model expert") and the rollout is
|
||||
# wasted. Audit (2026-07-12) on the 213-row July train split:
|
||||
#
|
||||
# medical 41/57 never routed (72%) | math 0/48 (0%)
|
||||
# chat 2/3 never routed (67%) | code 0/9 (0%)
|
||||
#
|
||||
# i.e. ~25% of every generation batch was spent on tasks that cannot teach routing.
|
||||
# Excluded by default; pass ``exclude_domains=()`` to restore the old mix.
|
||||
LOOKUP_DOMAINS = frozenset({"medical", "chat"})
|
||||
|
||||
|
||||
def load_generalthought(
|
||||
*,
|
||||
n: int = 2000,
|
||||
seed: int = 42,
|
||||
source: Optional[Iterable[Dict[str, Any]]] = None,
|
||||
buffer: Optional[int] = None,
|
||||
exclude_domains: Iterable[str] = LOOKUP_DOMAINS,
|
||||
) -> Iterator[Task]:
|
||||
"""Yield up to ``n`` GeneralThought tasks, randomly mixed across categories.
|
||||
|
||||
``source`` overrides the HF stream with an iterable of raw row dicts (tests).
|
||||
We over-stream a buffer and shuffle it so the yielded tasks are a random mix
|
||||
of the source's categories rather than a single leading shard. ``buffer``
|
||||
overrides the shuffle-buffer size; pass a small value (e.g. for smoke runs)
|
||||
to avoid streaming the default 6000-row floor.
|
||||
|
||||
``exclude_domains`` drops lookup-shaped questions (see ``LOOKUP_DOMAINS``);
|
||||
the buffer is grown to compensate so we still yield ``n`` tasks.
|
||||
"""
|
||||
drop = {d.lower() for d in exclude_domains}
|
||||
if source is None:
|
||||
from datasets import load_dataset # lazy: optional dep / network
|
||||
|
||||
source = load_dataset(GENERALTHOUGHT_ID, split="train", streaming=True)
|
||||
|
||||
rng = random.Random(seed)
|
||||
# Buffer a generous multiple of n so the shuffle actually mixes categories,
|
||||
# but stay bounded for the multi-hundred-K streams.
|
||||
buf_cap = buffer if buffer is not None else max(n * 6, 6000)
|
||||
buf: List[Task] = []
|
||||
seen = 0
|
||||
for i, row in enumerate(source):
|
||||
task = _normalize_generalthought(dict(row), index=i)
|
||||
if task is None:
|
||||
continue
|
||||
seen += 1
|
||||
if task.domain in drop:
|
||||
continue
|
||||
buf.append(task)
|
||||
if len(buf) >= buf_cap:
|
||||
break
|
||||
# Excluded domains are a large slice of GeneralThought, so cap how far we
|
||||
# read rather than streaming the whole shard hunting for survivors.
|
||||
if seen >= buf_cap * 6:
|
||||
break
|
||||
rng.shuffle(buf)
|
||||
for task in buf[:n]:
|
||||
yield task
|
||||
|
||||
|
||||
# --- OpenThoughts3 ------------------------------------------------------------
|
||||
|
||||
_BOXED_RE = re.compile(r"\\boxed\{")
|
||||
|
||||
|
||||
def _extract_boxed(text: str) -> Optional[str]:
|
||||
"""Return the contents of the last ``\\boxed{...}`` in ``text`` (brace-balanced)."""
|
||||
last = None
|
||||
for m in _BOXED_RE.finditer(text):
|
||||
start = m.end()
|
||||
depth = 1
|
||||
i = start
|
||||
while i < len(text) and depth > 0:
|
||||
c = text[i]
|
||||
if c == "{":
|
||||
depth += 1
|
||||
elif c == "}":
|
||||
depth -= 1
|
||||
i += 1
|
||||
if depth == 0:
|
||||
last = text[start : i - 1].strip()
|
||||
return last
|
||||
|
||||
|
||||
def _conversation_qa(conversations: Any) -> Optional[tuple[str, str]]:
|
||||
"""Pull (question, gold_answer) from an OpenThoughts ``conversations`` list."""
|
||||
if not isinstance(conversations, list):
|
||||
return None
|
||||
question = ""
|
||||
full = ""
|
||||
for turn in conversations:
|
||||
if not isinstance(turn, dict):
|
||||
continue
|
||||
who = str(turn.get("from") or "").lower()
|
||||
val = str(turn.get("value") or "")
|
||||
if who in ("human", "user") and not question:
|
||||
question = val.strip()
|
||||
elif who in ("gpt", "assistant"):
|
||||
full = val
|
||||
if not question or not full:
|
||||
return None
|
||||
# Drop the <think>…</think> trace; keep the post-reasoning solution.
|
||||
visible = re.sub(r"(?s)<think>.*?</think>", "", full).strip()
|
||||
# …and an UNCLOSED <think> (the source row was truncated mid-thought). The
|
||||
# regex above needs a closing tag, so without this the strip silently no-ops
|
||||
# and — when there's no \boxed{} either — `answer` falls through to the whole
|
||||
# reasoning dump. That gold is unusable: math is verified by string/numeric
|
||||
# match with no judge, so `_math_equal("19", "<think> Okay, so...")` is always
|
||||
# False and the task becomes unwinnable no matter what the model answers.
|
||||
if "<think>" in visible:
|
||||
visible = visible.split("<think>", 1)[0].strip()
|
||||
body = visible or full
|
||||
answer = _extract_boxed(body) or _extract_boxed(full) or body.strip()
|
||||
if not answer:
|
||||
return None
|
||||
return question, answer
|
||||
|
||||
|
||||
# Sources that cannot produce a usable training example, dropped at load time so
|
||||
# we never spend a rollout on them (audit 2026-07-12, 100-task Haiku batch):
|
||||
#
|
||||
# stackexchange_codegolf 0/13 correct (0%)
|
||||
#
|
||||
# Code-golf asks for the SHORTEST program, which is adversarial to both the model
|
||||
# and any verifier — there's no canonical answer to match and no judge can score
|
||||
# "is this minimal". By contrast normal code is fine (nvidia/OpenCodeReasoning
|
||||
# scored 7/12 = 58%), so this drops the pathological source, not the domain.
|
||||
UNUSABLE_SOURCES = ("codegolf",)
|
||||
|
||||
# A gold that is just a QED / proof marker — there's no value to match against.
|
||||
_PROOF_GOLD_RE = re.compile(
|
||||
r"^\s*\\?\(?\s*(\\blacksquare|\\qed|\\square|QED)\s*\\?\)?\s*$", re.IGNORECASE
|
||||
)
|
||||
|
||||
|
||||
def _gold_is_unusable(answer: str, domain: str) -> bool:
|
||||
"""Would this task be unwinnable no matter what the model answers?
|
||||
|
||||
Skip it at LOAD time — a rollout spent here is guaranteed waste.
|
||||
|
||||
* ``<think>`` in the gold means answer-extraction *failed* and we kept the raw
|
||||
reasoning trace. Unwinnable in EVERY domain (audit 2026-07-13: 12/90 rollouts,
|
||||
correct rate 2/12 = 17% — vs 65-68% on a good gold). Note this is about
|
||||
extraction failing, NOT about length: a *long* gold is fine, because the LLM
|
||||
judge copes with it (code 67%, science 67%).
|
||||
* A proof marker (``\\(\\blacksquare\\)``) is a QED symbol, not an answer.
|
||||
* A long prose gold is unusable for MATH specifically, because math is verified
|
||||
by string/numeric match with no judge fallback.
|
||||
"""
|
||||
if "<think>" in answer:
|
||||
return True
|
||||
if _PROOF_GOLD_RE.search(answer):
|
||||
return True
|
||||
if domain == "math" and len(answer) > 300:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _normalize_openthoughts(row: Dict[str, Any], *, index: int = 0) -> Optional[Task]:
|
||||
source = str(row.get("source") or "").lower()
|
||||
if any(s in source for s in UNUSABLE_SOURCES):
|
||||
return None
|
||||
qa = _conversation_qa(row.get("conversations"))
|
||||
if qa is None:
|
||||
return None
|
||||
question, answer = qa
|
||||
if not question or not answer:
|
||||
return None
|
||||
domain = str(row.get("domain") or "unknown").strip().lower() or "unknown"
|
||||
if _gold_is_unusable(answer, domain):
|
||||
return None
|
||||
# PROOF tasks: the "answer" is a QED marker, not a value ("\(\blacksquare\)").
|
||||
# There is nothing to match against — the real answer is the proof itself, and
|
||||
# the math path has no judge. Unwinnable by construction; skip.
|
||||
if _PROOF_GOLD_RE.search(answer):
|
||||
return None
|
||||
task_id = (
|
||||
str(row.get("id") or row.get("source") or f"openthoughts-{index}") + f"-{index}"
|
||||
)
|
||||
return Task(
|
||||
task_id=task_id,
|
||||
question=question,
|
||||
answer=answer,
|
||||
domain=domain,
|
||||
difficulty=str(row.get("difficulty") or "").strip(),
|
||||
dataset="OpenThoughts3",
|
||||
subsector=str(row.get("source") or "").strip(),
|
||||
)
|
||||
|
||||
|
||||
# The 1.2M set ships as 120 uniform parquet shards with no domain in the file
|
||||
# names, but the rows are front-ordered by domain. Probing one row per shard puts
|
||||
# code in shards ~0-30, math ~32-104, science ~105-119. Streaming the single
|
||||
# combined split therefore has to read the entire code (+ math) prefix before it
|
||||
# reaches math/science — minutes of wasted I/O. Instead we stream each domain from
|
||||
# shards *inside* its own region. A few shards (~10K rows each) cover any quota.
|
||||
_OT_SHARD_FMT = "data/train-{:05d}-of-00120.parquet"
|
||||
_OT_DOMAIN_SHARDS: Dict[str, List[int]] = {
|
||||
"code": [0, 1, 2, 3],
|
||||
"math": [45, 46, 47, 48],
|
||||
"science": [112, 113, 114, 115, 116, 117, 118, 119],
|
||||
}
|
||||
|
||||
|
||||
def load_openthoughts(
|
||||
*,
|
||||
n_code: int = 2000,
|
||||
n_math: int = 2000,
|
||||
n_science: int = 2000,
|
||||
seed: int = 42,
|
||||
source: Optional[Iterable[Dict[str, Any]]] = None,
|
||||
) -> Iterator[Task]:
|
||||
"""Yield OpenThoughts3 tasks balanced across code / math / science.
|
||||
|
||||
``source`` overrides the stream with raw row dicts (tests): rows are routed
|
||||
into per-domain buffers, stopping once every quota is met. With no override we
|
||||
stream each domain from its own shard region (see ``_OT_DOMAIN_SHARDS``) so a
|
||||
balanced pull never has to read the entire front-ordered code/math prefix.
|
||||
"""
|
||||
quotas = {"code": n_code, "math": n_math, "science": n_science}
|
||||
bufs: Dict[str, List[Task]] = {k: [] for k in quotas}
|
||||
|
||||
if source is not None:
|
||||
for i, row in enumerate(source):
|
||||
task = _normalize_openthoughts(dict(row), index=i)
|
||||
if task is None:
|
||||
continue
|
||||
dom = task.domain
|
||||
if dom in bufs and len(bufs[dom]) < quotas[dom]:
|
||||
bufs[dom].append(task)
|
||||
if all(len(bufs[k]) >= quotas[k] for k in quotas):
|
||||
break
|
||||
else:
|
||||
from datasets import load_dataset # lazy: optional dep / network
|
||||
|
||||
idx = 0
|
||||
for dom, quota in quotas.items():
|
||||
if quota <= 0:
|
||||
continue
|
||||
files = [_OT_SHARD_FMT.format(s) for s in _OT_DOMAIN_SHARDS[dom]]
|
||||
ds = load_dataset(
|
||||
OPENTHOUGHTS_ID, data_files=files, split="train", streaming=True
|
||||
)
|
||||
for row in ds:
|
||||
task = _normalize_openthoughts(dict(row), index=idx)
|
||||
idx += 1
|
||||
if task is None or task.domain != dom:
|
||||
continue
|
||||
bufs[dom].append(task)
|
||||
if len(bufs[dom]) >= quota:
|
||||
break
|
||||
|
||||
rng = random.Random(seed)
|
||||
out: List[Task] = []
|
||||
for k in quotas:
|
||||
out.extend(bufs[k])
|
||||
rng.shuffle(out)
|
||||
for task in out:
|
||||
yield task
|
||||
|
||||
|
||||
# --- combined SFT / GRPO sets -------------------------------------------------
|
||||
|
||||
|
||||
def load_sft_tasks(
|
||||
*, seed: int = 42, cap: Optional[int] = None, balanced: bool = True
|
||||
) -> List[Task]:
|
||||
"""The 8K cold-start SFT set: 2K GeneralThought + 2K code + 2K math + 2K science.
|
||||
|
||||
``cap`` (smoke runs): when set, draw ~``cap`` tasks total.
|
||||
- default (``balanced=False``): GeneralThought only with a small stream
|
||||
buffer — fastest, but the domain mix is whatever GeneralThought yields
|
||||
(math/code/medical/chat), so a given sample can skew (e.g. all-medical).
|
||||
- ``balanced=True``: ~cap/4 from each of GeneralThought + OpenThoughts
|
||||
code/math/science, so the smoke is representative of the real run. Now
|
||||
cheap because OpenThoughts streams from per-domain shards (no front-prefix).
|
||||
"""
|
||||
if cap is not None and cap > 0:
|
||||
if balanced:
|
||||
per = max(cap // 4, 1)
|
||||
tasks = list(load_generalthought(n=per, seed=seed, buffer=max(per * 6, 64)))
|
||||
tasks.extend(
|
||||
load_openthoughts(n_code=per, n_math=per, n_science=per, seed=seed)
|
||||
)
|
||||
rng = random.Random(seed)
|
||||
rng.shuffle(tasks)
|
||||
return tasks[:cap]
|
||||
buf = max(cap * 4, 64)
|
||||
tasks = list(load_generalthought(n=cap, seed=seed, buffer=buf))
|
||||
random.Random(seed).shuffle(tasks)
|
||||
return tasks[:cap]
|
||||
tasks: List[Task] = []
|
||||
tasks.extend(load_generalthought(n=2000, seed=seed))
|
||||
tasks.extend(load_openthoughts(n_code=2000, n_math=2000, n_science=2000, seed=seed))
|
||||
rng = random.Random(seed)
|
||||
rng.shuffle(tasks)
|
||||
return tasks
|
||||
|
||||
|
||||
__all__ = [
|
||||
"GENERALTHOUGHT_ID",
|
||||
"OPENTHOUGHTS_ID",
|
||||
"Task",
|
||||
"load_generalthought",
|
||||
"load_openthoughts",
|
||||
"load_sft_tasks",
|
||||
]
|
||||
@@ -0,0 +1,58 @@
|
||||
"""Single source of truth for orchestrator dataset names.
|
||||
|
||||
raw/{name}-{stamp}/data.jsonl raw/qwen-july-7-2026-0553pm/data.jsonl
|
||||
sft/{name}-{split}-{stamp}.jsonl sft/qwen-train-july-7-2026-0553pm.jsonl
|
||||
|
||||
The raw dir and every file carved from it share the same ``stamp``, so a curated
|
||||
split always traces back to the generation run that produced it by eye.
|
||||
|
||||
The stamp is written out (``july-7-2026-0553pm``) rather than numeric because
|
||||
these names are read by humans far more often than they are parsed. The cost is
|
||||
that month names sort alphabetically, so ``ls`` is NOT chronological — use
|
||||
``ls -t``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
__all__ = ["run_stamp", "dataset_name", "raw_dir_name", "stamp_from"]
|
||||
|
||||
# {month}-{day}-{year}-{hhmm}{am|pm}, e.g. july-7-2026-0553pm
|
||||
_STAMP_RE = re.compile(r"[a-z]+-\d{1,2}-\d{4}-\d{4}(?:am|pm)", re.I)
|
||||
|
||||
|
||||
def stamp_from(filename: str) -> Optional[str]:
|
||||
"""Pull the stamp out of a raw dir / dataset name, or None if it carries none.
|
||||
|
||||
``qwen-clean-july-7-2026-0553pm.jsonl`` -> ``july-7-2026-0553pm``. Lets a split
|
||||
inherit its POOL's stamp instead of stamping itself with today's date — the
|
||||
split belongs to the run that generated the data, not to the day it was carved.
|
||||
"""
|
||||
m = _STAMP_RE.search(filename)
|
||||
return m.group(0).lower() if m else None
|
||||
|
||||
|
||||
def run_stamp(when: Optional[time.struct_time] = None) -> str:
|
||||
"""``july-7-2026-0553pm`` — month-in-words, day, year, 12h clock.
|
||||
|
||||
Local time, matching the wall-clock the run was launched at.
|
||||
"""
|
||||
t = when or time.localtime()
|
||||
month = time.strftime("%B", t).lower() # july
|
||||
day = t.tm_mday # no zero-pad: 7, not 07
|
||||
clock = time.strftime("%I%M%p", t).lower() # 0553pm (zero-padded hour)
|
||||
return f"{month}-{day}-{t.tm_year}-{clock}"
|
||||
|
||||
|
||||
def raw_dir_name(name: str, stamp: Optional[str] = None, tag: str = "") -> str:
|
||||
"""``qwen-july-7-2026-0553pm`` (+ ``-{tag}`` when disambiguating a variant)."""
|
||||
base = f"{name}-{stamp or run_stamp()}"
|
||||
return f"{base}-{tag}" if tag else base
|
||||
|
||||
|
||||
def dataset_name(name: str, split: str, stamp: Optional[str] = None) -> str:
|
||||
"""``qwen-train-july-7-2026-0553pm`` (no extension)."""
|
||||
return f"{name}-{split}-{stamp or run_stamp()}"
|
||||
@@ -0,0 +1,518 @@
|
||||
"""Rejection-sampling SFT-data generator (the ToolOrchestra cold-start).
|
||||
|
||||
For each reasoning task: roll out a teacher orchestrator N times, verify each
|
||||
trajectory, keep the passing ones (optionally just the cheapest), and serialize
|
||||
them into the unified-tool ``conversations`` JSONL the SFT trainer consumes.
|
||||
|
||||
The expensive/network parts are injected so the orchestration is pure and
|
||||
offline-testable:
|
||||
|
||||
* ``rollout_fn(task) -> UnifiedRollout`` — one teacher rollout (temperature>0).
|
||||
* ``verify_fn(task, rollout) -> bool`` — did the trajectory solve the task?
|
||||
(e.g. ``verify.make_verifier()``, an LLM judge over the final answer.)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Iterable, List, Optional
|
||||
|
||||
from openjarvis.agents.hybrid.expert_registry import ExpertTool
|
||||
from openjarvis.agents.hybrid.toolorchestra.rollout import UnifiedRollout
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import (
|
||||
_CONTROL_TOKEN_RE,
|
||||
_strip_control_tokens,
|
||||
trajectory_to_record,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# The sampler only touches ``task.task_id`` / ``task.instruction`` / ``task.domain``
|
||||
# (via ``trajectory_to_record``), so any task dataclass works — the reasoning
|
||||
# ``datasets.Task`` is what production uses, paired with its own ``verify_fn``
|
||||
# (e.g. ``verify.make_verifier()``).
|
||||
TaskLike = Any
|
||||
RolloutFn = Callable[[TaskLike], Optional[UnifiedRollout]]
|
||||
VerifyFn = Callable[[TaskLike, UnifiedRollout], bool]
|
||||
|
||||
# Human-readable blurb for each source dataset, stamped onto every record as
|
||||
# ``dataset_description`` so a row is self-describing (what corpus the question
|
||||
# came from) without a lookup. Keyed on the record's ``dataset`` value.
|
||||
DATASET_DESCRIPTIONS = {
|
||||
"GeneralThought": (
|
||||
"natolambert/GeneralThought-430K-filtered — filtered reasoning Q&A from "
|
||||
"gr.inc; each item has a question and a short reference (gold) answer "
|
||||
"distilled from DeepSeek-R1 traces."
|
||||
),
|
||||
"OpenThoughts3": (
|
||||
"open-thoughts/OpenThoughts3-1.2M — code/math/science reasoning tasks; "
|
||||
"the gold answer is extracted from the boxed final solution."
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def dataset_description(dataset: str) -> str:
|
||||
"""Blurb for a source dataset name (empty string if unknown)."""
|
||||
return DATASET_DESCRIPTIONS.get((dataset or "").strip(), "")
|
||||
|
||||
|
||||
# Markers of a broken expert/observation that must never become a training
|
||||
# target (empty turns, tracebacks, truncated tool errors, dead-provider strings).
|
||||
_ERR_MARKERS = (
|
||||
"Traceback (most recent call last)",
|
||||
"insufficient_quota",
|
||||
"rate limit",
|
||||
"RateLimitError",
|
||||
"[invalid tool",
|
||||
"Error code: 4",
|
||||
"Error code: 5",
|
||||
"InternalServerError",
|
||||
"ConnectionError",
|
||||
"timed out",
|
||||
# A tool call that reached dispatch with no/empty input, or a bridged-tool
|
||||
# crash — the call was malformed. The trajectory routed garbage; drop it.
|
||||
"no input provided",
|
||||
"[openjarvis-tool error",
|
||||
)
|
||||
|
||||
# A final answer cut off mid-expression ends on a dangling connective
|
||||
# (``a=1, b=``, ``result = ``, ``compute (``) — the decode hit the token cap
|
||||
# before finishing. A real distilled answer never ends this way.
|
||||
_TRUNCATED_TAIL_RE = re.compile(r"[=,(\[{+\-*/]\s*$")
|
||||
|
||||
# A line that only source code starts with. Used to spot an UNFENCED program so the
|
||||
# markdown/essay guards don't read its `#` comments as headers.
|
||||
_CODE_SIGNAL_RE = re.compile(
|
||||
r"(?m)^\s*(?:#!|def |class |import |from \s*\w+\s+import|function |const |let |"
|
||||
r"var |public |private |#include|package |using |fn |func )"
|
||||
)
|
||||
|
||||
|
||||
def clean_reason(roll: UnifiedRollout) -> Optional[str]:
|
||||
"""Why this trajectory is unfit to be an SFT *target* — ``None`` if it's fine.
|
||||
|
||||
Returns the FIRST failing check as a human-readable string so rejections are
|
||||
auditable (persisted as ``clean_reason`` on every record). Previously this
|
||||
returned a bare bool, which meant a rejected record gave no clue why — and
|
||||
the reason is genuinely not recoverable from the saved JSONL, because the
|
||||
gate runs on the RAW rollout while the serializer scrubs the stored copy.
|
||||
|
||||
|
||||
The judge decides semantic correctness; this catches the ~10-12% of records
|
||||
whose targets are empty / a Traceback / a truncated tool error / unrouted —
|
||||
garbage the model must never be trained to imitate. Requirements:
|
||||
* non-empty final answer, no error marker, balanced ``<think>`` (not
|
||||
truncated mid-thought);
|
||||
* at least one model-expert tool call (it actually *routed*);
|
||||
* no error-marker / empty observation in the kept trajectory.
|
||||
"""
|
||||
fa = (roll.final_answer or "").strip()
|
||||
if not fa:
|
||||
return "empty final answer"
|
||||
if any(m in fa for m in _ERR_MARKERS):
|
||||
return "error marker in final answer"
|
||||
# A thought that OPENS and never closes = the decode hit the cap mid-thought.
|
||||
# A stray CLOSING </think> with no opener is NOT truncation — Qwen3.x's chat
|
||||
# template pre-fills the opening `<think>` tag, so the model's completion
|
||||
# begins inside the reasoning block and only ever emits the closing tag. The
|
||||
# old check (`!=`) treated that template artifact as truncation and rejected
|
||||
# essentially every thinking rollout; the serializer strips the stray tag
|
||||
# anyway, so the stored target was fine. Only unclosed thoughts are bad.
|
||||
if fa.count("<think>") > fa.count("</think>"):
|
||||
return "unclosed <think> (truncated mid-thought)"
|
||||
if _TRUNCATED_TAIL_RE.search(fa): # truncated mid-expression (e.g. "a=1, b=")
|
||||
return "final truncated mid-expression"
|
||||
# Malformed final: a proper final answer is plain text, NOT a stray/broken
|
||||
# <tool_call> tag (the model's answer leaking inside a tool call that failed
|
||||
# to parse). (\boxed{} is NOT rejected here — the serializer de-boxes it, so
|
||||
# an otherwise-good boxed answer is salvaged rather than thrown away.)
|
||||
if "<tool_call>" in fa or "</tool_call>" in fa:
|
||||
return "tool_call tag leaked into final answer"
|
||||
# Control-token backstop (defense in depth). The serializer strips leaked
|
||||
# control/special tokens (<|im_end|>, <|tool_call>, <start_of_turn>, …) to
|
||||
# salvage good answers, so we mirror that strip and gate on the RESULT:
|
||||
# * strips to empty -> the answer WAS nothing but control tokens -> drop.
|
||||
# * a token survives the strip -> something malformed the stripper can't
|
||||
# safely delete mid-answer -> drop rather than train on it.
|
||||
# A clean answer (or one the stripper fully salvages) sails through.
|
||||
fa_stripped = _strip_control_tokens(fa)
|
||||
if not fa_stripped:
|
||||
return "final answer was only control tokens"
|
||||
if _CONTROL_TOKEN_RE.search(fa_stripped):
|
||||
return "control tokens in final answer"
|
||||
# Degenerate repetition: the same substantive line emitted many times (the
|
||||
# small-model decode loop). Reject rather than train on it.
|
||||
_lines = [ln.strip() for ln in fa.split("\n") if len(ln.strip()) > 30]
|
||||
if _lines and Counter(_lines).most_common(1)[0][1] >= 4:
|
||||
return "degenerate repetition"
|
||||
# Word-salad run-on: a giant unbroken line (no newline) is the other decode
|
||||
# collapse — the model spraying novel tokens instead of repeating. Reject.
|
||||
if any(len(ln) > 2000 for ln in fa.split("\n")):
|
||||
return "word-salad run-on line"
|
||||
# Essay-style final: the format wants a distilled answer, not a multi-section
|
||||
# writeup. Check the answer AFTER the FINAL_ANSWER marker (reasoning before it
|
||||
# is fine). NOTE: length/bold limits relaxed for the BALANCED/harder mix —
|
||||
# code & multi-step math answers are legitimately longer and use **bold** for
|
||||
# the key result, so the old 700-char + any-bold rejects were dropping ~half
|
||||
# the correct hard-task answers. Keep the STRUCTURAL essay signals (many
|
||||
# numbered sections, markdown headers, tables) which still catch real essays.
|
||||
_marks = list(re.finditer(r"(?im)FINAL[_\s]?ANSWER\s*:?", fa))
|
||||
_ans = fa[_marks[-1].end() :].strip() if _marks else fa
|
||||
# Every structural essay guard below runs on the PROSE only — a fenced code
|
||||
# block is a legitimate answer, not an essay, and its contents are not
|
||||
# markdown. Without this:
|
||||
# * a Python comment ("# Use download_as_text()") matches the markdown-header
|
||||
# regex `^\s*#{1,6}\s` and the whole answer is rejected as an "essay";
|
||||
# * a numbered list in a docstring trips the numbered-sections guard;
|
||||
# * an ASCII table in code output trips the markdown-table guard;
|
||||
# * a real program blows the 2000-char limit on its own.
|
||||
# Audit 2026-07-13: this was the difference between code correct=12/22 and code
|
||||
# USABLE=3/22 — we were binning ~3 of every 4 correct code answers.
|
||||
_prose = re.sub(r"```.*?```", "", _ans, flags=re.DOTALL)
|
||||
_prose = re.sub(r"(?m)^ {4,}\S.*$", "", _prose) # indented code blocks too
|
||||
# …and UNFENCED raw code. Haiku often answers a code task with a bare program
|
||||
# and no ``` fence at all, and then its shebang (`#!/usr/bin/env python3`) and
|
||||
# its comments (`# Read input`) trip the markdown-header regex exactly like the
|
||||
# fenced case did. If what's left after stripping fences still reads as code,
|
||||
# there is no prose to run essay guards against. (Audit 2026-07-13 — the fenced
|
||||
# fix caught only half of this bug.)
|
||||
if _CODE_SIGNAL_RE.search(_prose):
|
||||
_prose = ""
|
||||
# Markdown-table dump: the model wrote a formatted table instead of the short
|
||||
# exact value the format demands. A ``|---|`` separator row is the tell.
|
||||
if re.search(r"\|\s*:?-{2,}", _prose):
|
||||
return "markdown table dump"
|
||||
if len(_prose) > 2000: # prose-only: a long program is fine, an essay is not
|
||||
return "final answer >2000 chars (essay)"
|
||||
# Tool-status echo as the final answer: on code tasks the orchestrator
|
||||
# sometimes ends with a file_write/shell status line ("Successfully wrote to
|
||||
# /tmp/x.py") instead of the actual answer — a non-answer that slips the
|
||||
# length/format checks. Reject. (Audit: ~code/math trajectories where the
|
||||
# trace collapsed but the row was still scored correct.)
|
||||
if re.search(
|
||||
r"(?i)\b(successfully (wrote|created|saved|executed|ran)|written to /|"
|
||||
r"file (written|saved|created)|no further actions?)\b",
|
||||
_ans,
|
||||
):
|
||||
return "tool-status echo as the answer"
|
||||
if (
|
||||
len(re.findall(r"(?m)^\s*\d+\.\s", _prose)) >= 6
|
||||
): # many numbered sections = essay (was 4)
|
||||
return "many numbered sections (essay)"
|
||||
if re.search(r"(?m)^\s*#{1,6}\s", _prose): # markdown headers = essay
|
||||
return "markdown headers (essay)"
|
||||
# Garbled / shouty final: a multi-word answer that's mostly UPPERCASE, or has
|
||||
# an absurdly long merged all-letter token, is decode garble — the audit found
|
||||
# e.g. "RABES PEPTETANUS BOOSTERS" (for "rabies PEP, tetanus boosters") passing.
|
||||
if len(_ans) > 20:
|
||||
# Measure the caps ratio over PROSE tokens only. A chemistry answer is
|
||||
# legitimately uppercase-heavy — "HC≡CH → NaNH₂ → CH₃CH₂CH₂Br" is a
|
||||
# reaction scheme, not shouting — and the old whole-string ratio rejected
|
||||
# those outright (audit 2026-07-12), which is expensive because organic
|
||||
# chemistry is one of our highest-yield domains. So ignore any token that
|
||||
# carries a digit, a subscript/superscript, or a non-ASCII symbol (arrows,
|
||||
# bond glyphs): that's formula notation, not decode collapse. What's left
|
||||
# is real words, which is what the guard was built for ("RABES PEPTETANUS
|
||||
# BOOSTERS").
|
||||
_words = [
|
||||
w
|
||||
for w in _ans.split()
|
||||
if w.isalpha() and len(w) >= 3 and w.isascii()
|
||||
]
|
||||
_alpha = [c for w in _words for c in w]
|
||||
if (
|
||||
len(_words) >= 3
|
||||
and _alpha
|
||||
and sum(c.isupper() for c in _alpha) / len(_alpha) > 0.7
|
||||
):
|
||||
return "ALL-CAPS garble"
|
||||
if any(len(w) > 25 and w.isalpha() for w in _ans.split()):
|
||||
return "merged mega-token garble"
|
||||
# Reject runaway final answers — a distilled answer, not a multi-KB essay or
|
||||
# a verbatim dump of an expert observation. (Audit: 279 finals >4k chars, the
|
||||
# worst a 409k-char wordlist; ~363 were prefix-identical to a tool obs.)
|
||||
if len(fa) > 8000:
|
||||
return "runaway final (>8k chars)"
|
||||
fa_head = re.sub(r"\s+", " ", fa[:200]).strip()
|
||||
for t in roll.turns:
|
||||
if t.observation:
|
||||
obs_head = re.sub(r"\s+", " ", t.observation[:200]).strip()
|
||||
if (
|
||||
fa_head and fa_head == obs_head and len(fa) > 200
|
||||
): # long final = copied tool dump (short relays are legit)
|
||||
return "final answer copied verbatim from a tool result"
|
||||
# "routed" = delegated to a model EXPERT (not just a utility like web_search).
|
||||
# When anonymized, expert calls are the anon labels in anon_map; otherwise
|
||||
# fall back to "made any tool call".
|
||||
expert_names = set(roll.anon_map or {})
|
||||
called = {name for name, _ in roll.tool_calls()}
|
||||
if expert_names:
|
||||
if not (called & expert_names):
|
||||
return "NEVER routed to a model expert"
|
||||
elif roll.num_tool_calls < 1:
|
||||
return "no tool calls at all"
|
||||
# A broken EXPERT observation (rate limit / 5xx / dead provider) means the
|
||||
# trajectory routed into a hole and whatever follows is built on nothing —
|
||||
# reject it.
|
||||
#
|
||||
# A sandbox tool is different. `code_interpreter` returning a Python
|
||||
# "Traceback (most recent call last)" is the tool WORKING: the model ran code,
|
||||
# it raised, and the model reads the error and fixes it. That error-recovery
|
||||
# loop is precisely what we want to teach. Treating it as a broken tool threw
|
||||
# out 24% of rollouts (audit 2026-07-12: 21 of 23 flagged observations were
|
||||
# code_interpreter tracebacks the model then recovered from). Only an EMPTY
|
||||
# observation is fatal for a sandbox tool — that means the tool itself died.
|
||||
for t in roll.turns:
|
||||
if t.tool_name is None:
|
||||
continue
|
||||
obs = (t.observation or "").strip()
|
||||
if not obs:
|
||||
return "empty tool observation"
|
||||
is_expert = t.tool_name in expert_names if expert_names else False
|
||||
if is_expert and any(m in obs for m in _ERR_MARKERS):
|
||||
return "error observation from a model expert"
|
||||
# Reasoning-side decode collapse: an intermediate turn whose reasoning is a
|
||||
# giant unbroken line, or is dominated by exotic unicode / emoji spam, is
|
||||
# garbage even when the final answer looks fine (the answer-only guards above
|
||||
# miss it). The audit found a trajectory with a thousands-char symbol blob in
|
||||
# its reasoning that scored correct and slipped through.
|
||||
for t in roll.turns:
|
||||
r = t.reasoning or ""
|
||||
if any(len(ln) > 2000 for ln in r.split("\n")):
|
||||
return "reasoning decode collapse (giant line)"
|
||||
if len(r) > 200:
|
||||
weird = sum(1 for ch in r if ord(ch) > 0x2100 and not ch.isalnum())
|
||||
if weird / len(r) > 0.10: # >10% exotic-unicode/emoji = decode spam
|
||||
return "reasoning unicode/emoji spam"
|
||||
return None
|
||||
|
||||
|
||||
def _target_is_clean(roll: UnifiedRollout) -> bool:
|
||||
"""Back-compat bool wrapper around :func:`clean_reason`."""
|
||||
return clean_reason(roll) is None
|
||||
|
||||
|
||||
def _process_task(
|
||||
task: TaskLike,
|
||||
*,
|
||||
rollout_fn: RolloutFn,
|
||||
verify_fn: VerifyFn,
|
||||
samples_per_task: int,
|
||||
max_keep_per_task: int,
|
||||
stop_at_keep: bool = False,
|
||||
keep_all: bool = False,
|
||||
) -> List[tuple[UnifiedRollout, bool]]:
|
||||
"""One task's rejection-sampling unit of work: roll out ``samples_per_task``
|
||||
times and verify each. Returns ``(rollout, is_correct)`` for **every** sample
|
||||
(the caller decides what to write). Runs in a worker thread, so the only
|
||||
shared state it touches is the injected fns (each rollout builds its own
|
||||
OpenAI client; the tool executor is built under a lock).
|
||||
|
||||
``stop_at_keep``: short-circuit as soon as ``max_keep_per_task`` passing
|
||||
trajectories are found, instead of always exhausting ``samples_per_task``.
|
||||
Big throughput win (no wasted rollouts on already-solved tasks), but it
|
||||
trades away the cheapest-of-N cost optimisation (we keep the first passers,
|
||||
not the cheapest). Ignored when ``keep_all`` is set (we want every sample).
|
||||
Default off preserves the original semantics."""
|
||||
results: List[tuple[UnifiedRollout, bool]] = []
|
||||
n_correct = 0
|
||||
for _ in range(samples_per_task):
|
||||
roll = rollout_fn(task)
|
||||
if roll is None:
|
||||
continue
|
||||
ok = bool(verify_fn(task, roll))
|
||||
results.append((roll, ok))
|
||||
if ok:
|
||||
n_correct += 1
|
||||
if stop_at_keep and not keep_all and n_correct >= max_keep_per_task:
|
||||
break
|
||||
return results
|
||||
|
||||
|
||||
def generate_sft_dataset(
|
||||
out_path: str,
|
||||
*,
|
||||
tasks: Iterable[TaskLike],
|
||||
tools: List[ExpertTool],
|
||||
rollout_fn: RolloutFn,
|
||||
verify_fn: VerifyFn,
|
||||
samples_per_task: int = 4,
|
||||
max_keep_per_task: int = 1,
|
||||
reward_fn: Optional[Callable[[UnifiedRollout], float]] = None,
|
||||
concurrency: int = 1,
|
||||
stop_at_keep: bool = False,
|
||||
keep_all: bool = False,
|
||||
record_extra: Optional[dict] = None,
|
||||
) -> dict:
|
||||
"""Run rejection sampling over ``tasks`` and write the SFT JSONL.
|
||||
|
||||
``max_keep_per_task`` caps records kept per task; when >1 the cheapest
|
||||
passing trajectories are kept first. ``concurrency`` rolls out that many tasks
|
||||
in parallel (each task issues its samples sequentially, so ~``concurrency``
|
||||
requests hit the served model at once) — set 1 for the original sequential
|
||||
behaviour. Records are written as each task finishes, so peak memory is bounded
|
||||
by the in-flight tasks, not the whole dataset. Returns stats + writes a
|
||||
``.stats.json``.
|
||||
|
||||
``keep_all``: write **every** rolled-out trajectory (correct and incorrect)
|
||||
rather than dropping the failures. Each record gets ``correct`` (verifier
|
||||
verdict) and ``kept`` (True on the cheapest-correct sample — the one the
|
||||
rejection sampler would have selected). This preserves the full sample set so
|
||||
you can compute accuracy / inspect failures; the SFT trainer should filter on
|
||||
``correct`` (or ``kept``). Default off = original drop-the-failures behaviour.
|
||||
"""
|
||||
out = Path(out_path)
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
tasks = list(tasks)
|
||||
seen = len(tasks)
|
||||
written = 0
|
||||
correct_written = 0
|
||||
incorrect_written = 0
|
||||
dropped = 0 # tasks that produced no kept record
|
||||
tasks_solved = 0 # tasks with >=1 correct sample
|
||||
domain_counts: Counter[str] = Counter()
|
||||
|
||||
def _work(task: TaskLike) -> tuple[TaskLike, List[tuple[UnifiedRollout, bool]]]:
|
||||
return task, _process_task(
|
||||
task,
|
||||
rollout_fn=rollout_fn,
|
||||
verify_fn=verify_fn,
|
||||
samples_per_task=samples_per_task,
|
||||
max_keep_per_task=max_keep_per_task,
|
||||
stop_at_keep=stop_at_keep,
|
||||
keep_all=keep_all,
|
||||
)
|
||||
|
||||
def _emit(fh, task: TaskLike, roll: UnifiedRollout, ok: bool, kept: bool) -> None:
|
||||
nonlocal written, correct_written, incorrect_written
|
||||
reward = reward_fn(roll) if reward_fn else 0.0
|
||||
record = trajectory_to_record(
|
||||
task.task_id,
|
||||
task.instruction,
|
||||
tools,
|
||||
roll,
|
||||
reward=reward,
|
||||
domain=task.domain,
|
||||
)
|
||||
record["correct"] = ok
|
||||
record["kept"] = kept
|
||||
# Persist WHY a trajectory was rejected. The gate runs on the raw rollout
|
||||
# while the serializer scrubs the stored conversations, so the reason is
|
||||
# not recoverable from the saved record afterwards — capture it here.
|
||||
why = clean_reason(roll)
|
||||
record["clean"] = why is None
|
||||
record["clean_reason"] = why or ""
|
||||
# Task provenance so every record is self-describing: area (domain),
|
||||
# difficulty tier, source dataset, and fine subsector. Lets us slice
|
||||
# train/holdout/analysis by any of these later.
|
||||
record["area"] = task.domain
|
||||
record["difficulty"] = getattr(task, "difficulty", "") or ""
|
||||
record["dataset"] = getattr(task, "dataset", "") or ""
|
||||
record["dataset_description"] = dataset_description(record["dataset"])
|
||||
record["subsector"] = getattr(task, "subsector", "") or ""
|
||||
# Gold reference answer the verifier graded against. Persisted so every
|
||||
# record supports gold-vs-model inspection downstream (Braintrust, error
|
||||
# analysis) instead of only a bare `correct` boolean.
|
||||
record["gold_answer"] = getattr(task, "answer", "") or ""
|
||||
# Stamp provenance (which model/orchestrator generated this trajectory)
|
||||
# so records are self-identifying when pooled across model families.
|
||||
if record_extra:
|
||||
record.update(record_extra)
|
||||
fh.write(json.dumps(record) + "\n")
|
||||
fh.flush()
|
||||
written += 1
|
||||
correct_written += int(ok)
|
||||
incorrect_written += int(not ok)
|
||||
domain_counts[task.domain] += 1
|
||||
|
||||
def _write(fh, task: TaskLike, results: List[tuple[UnifiedRollout, bool]]) -> None:
|
||||
nonlocal dropped, tasks_solved
|
||||
# A trajectory is only eligible to be *kept* as a training target if the
|
||||
# judge passed it AND it's structurally clean (non-error, routed, not
|
||||
# truncated). Unclean-but-correct rollouts are still written in keep_all
|
||||
# mode (for analysis) but never marked kept.
|
||||
correct = [r for r in results if r[1] and _target_is_clean(r[0])]
|
||||
cheapest = min((r for r, _ in correct), key=lambda r: r.cost_usd, default=None)
|
||||
if correct:
|
||||
tasks_solved += 1
|
||||
if keep_all:
|
||||
# Write every sample; mark the cheapest-correct one as kept.
|
||||
if not results:
|
||||
dropped += 1
|
||||
return
|
||||
for roll, ok in results:
|
||||
_emit(fh, task, roll, ok, kept=(roll is cheapest))
|
||||
else:
|
||||
# Original behaviour: keep only cheapest-correct, capped.
|
||||
if not correct:
|
||||
dropped += 1
|
||||
return
|
||||
for roll in sorted((r for r, _ in correct), key=lambda r: r.cost_usd)[
|
||||
:max_keep_per_task
|
||||
]:
|
||||
_emit(fh, task, roll, True, kept=(roll is cheapest))
|
||||
|
||||
with out.open("w") as fh:
|
||||
if concurrency <= 1:
|
||||
for task in tasks:
|
||||
_, results = _work(task)
|
||||
_write(fh, task, results)
|
||||
else:
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
done = 0
|
||||
with ThreadPoolExecutor(max_workers=concurrency) as pool:
|
||||
futures = {pool.submit(_work, t): t for t in tasks}
|
||||
for fut in as_completed(futures):
|
||||
task = futures[fut]
|
||||
try:
|
||||
_, results = fut.result()
|
||||
except Exception as exc: # a task's worker died; count + skip
|
||||
logger.warning(
|
||||
"task %s failed: %s", getattr(task, "task_id", "?"), exc
|
||||
)
|
||||
results = []
|
||||
_write(fh, task, results)
|
||||
done += 1
|
||||
if done % 50 == 0 or done == seen:
|
||||
logger.info(
|
||||
"rejection-sampling: %d/%d tasks done "
|
||||
"(%d written, %d solved, %d dropped)",
|
||||
done,
|
||||
seen,
|
||||
written,
|
||||
tasks_solved,
|
||||
dropped,
|
||||
)
|
||||
|
||||
stats = {
|
||||
"out_path": str(out),
|
||||
"tasks_seen": seen,
|
||||
"records_written": written,
|
||||
"records_correct": correct_written,
|
||||
"records_incorrect": incorrect_written,
|
||||
"tasks_solved": tasks_solved,
|
||||
"tasks_dropped": dropped,
|
||||
"task_accuracy": round(tasks_solved / seen, 4) if seen else 0.0,
|
||||
"samples_per_task": samples_per_task,
|
||||
"keep_all": keep_all,
|
||||
"concurrency": concurrency,
|
||||
"domain_distribution": dict(domain_counts),
|
||||
}
|
||||
out.with_suffix(out.suffix + ".stats.json").write_text(json.dumps(stats, indent=2))
|
||||
logger.info(
|
||||
"Wrote %d SFT records to %s (%d correct, %d incorrect)",
|
||||
written,
|
||||
out,
|
||||
correct_written,
|
||||
incorrect_written,
|
||||
)
|
||||
return stats
|
||||
|
||||
|
||||
__all__ = ["generate_sft_dataset"]
|
||||
@@ -0,0 +1,279 @@
|
||||
"""Serialize a verified unified-tool rollout into an SFT ``conversations`` record.
|
||||
|
||||
Output matches what ``OrchestratorSFTDataset`` consumes, and trains the model to
|
||||
emit the ``<tool_call>{...}</tool_call>`` text form that
|
||||
``toolorchestra.parsing._parse_rl_tool_call`` already reads back. One record =
|
||||
one passing trajectory.
|
||||
|
||||
Roles: ``system`` (the unified tool catalog), ``user`` (the running ``Problem``
|
||||
prompt), ``assistant`` (reasoning + a ``<tool_call>`` tag, or the final answer),
|
||||
``tool`` (the executed observation).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from openjarvis.agents.hybrid.expert_registry import ExpertTool, build_tool_specs
|
||||
from openjarvis.agents.hybrid.toolorchestra.rollout import (
|
||||
_OBS_CAP as _ROLLOUT_OBS_CAP,
|
||||
UnifiedRollout,
|
||||
build_system_prompt,
|
||||
tool_call_tag,
|
||||
)
|
||||
|
||||
|
||||
def _system_prompt(tools: List[ExpertTool], rollout: UnifiedRollout) -> str:
|
||||
# Faithful ToolOrchestra system prompt (paper's prepare_sft_data.py format).
|
||||
# Prefer the EXACT specs the policy saw this rollout (anonymized labels when
|
||||
# anonymize=True) so the saved <tools> block matches the assistant's
|
||||
# <tool_call> tags. Falling back to the real registry here is the bug that
|
||||
# leaked real model names+pricing into anonymized records.
|
||||
specs = rollout.tool_specs if rollout.tool_specs else build_tool_specs(tools)
|
||||
return build_system_prompt(specs)
|
||||
|
||||
|
||||
def _normalize_think(text: str) -> str:
|
||||
"""Ensure a reasoning block has matched tags. Qwen rollouts routinely yield a
|
||||
dangling ``</think>`` with NO opening ``<think>`` (the chat template consumes
|
||||
the opener in the prompt, so only the closer survives in the completion). Add
|
||||
the opener back so the trained turn is well-formed."""
|
||||
t = (text or "").lstrip()
|
||||
if "</think>" in t and "<think>" not in t:
|
||||
t = "<think>\n" + t
|
||||
return t
|
||||
|
||||
|
||||
def _debox(text: str) -> str:
|
||||
"""Unwrap every ``\\boxed{X}`` -> ``X`` (balanced braces). The format kills
|
||||
\\boxed everywhere, but small models keep emitting it in otherwise-correct
|
||||
answers — de-box to salvage them rather than reject the whole trajectory."""
|
||||
out, i = [], 0
|
||||
marker = r"\boxed{"
|
||||
while i < len(text):
|
||||
j = text.find(marker, i)
|
||||
if j == -1:
|
||||
out.append(text[i:])
|
||||
break
|
||||
out.append(text[i:j])
|
||||
k = j + len(marker)
|
||||
depth, inner = 1, []
|
||||
while k < len(text) and depth:
|
||||
c = text[k]
|
||||
if c == "{":
|
||||
depth += 1
|
||||
elif c == "}":
|
||||
depth -= 1
|
||||
if depth == 0:
|
||||
break
|
||||
inner.append(c)
|
||||
k += 1
|
||||
out.append("".join(inner))
|
||||
i = k + 1
|
||||
return "".join(out)
|
||||
|
||||
|
||||
# Reasoning that PAROTS the training constraint back ("I must call a model per
|
||||
# my instructions", "this is testing whether I need to route") — the model
|
||||
# reasoning about its own harness. Never a good target: drop the whole line.
|
||||
_LEAK_RE = re.compile(
|
||||
r"(?im)^.*\b(?:"
|
||||
r"the user (?:is|was|wants|has|had|needs|'s|would|will|asked"
|
||||
r"|'ve asked|has asked)\b|"
|
||||
r"asked me to\b|" # echoes of an injected nudge turn
|
||||
r"make progress by delegat|"
|
||||
r"delegat\w+[^\n]*sub-?question|"
|
||||
r"(?:to |ask |send .{0,20}to )?one of the models\b|"
|
||||
r"instructions?|" # any self-reference to the rules
|
||||
r"requirement|"
|
||||
r"at least one (?:model|tool|expert)|" # "(call|invoke) at least one model"
|
||||
r"(?:call|invoke|route to|delegate to|use)\s+(?:a|one|at least one|the)\s+model\b|"
|
||||
r"route\s+(?:this|it|that)\b[^\n]{0,40}\bthrough\s+(?:a|one|the)\s+model\b|"
|
||||
r"without mentioning (?:any )?(?:tools?|reasoning|meta|steps)|"
|
||||
r"in the required format\b|"
|
||||
r"based on the (?:model|response|analysis|expert)|" # response-acknowledgment
|
||||
r"the model (?:provided|confirmed|gave|returned|correctly|analy\w+"
|
||||
r"|respon\w+|said|indicated|identified)|"
|
||||
r"i (?:got|received|have) (?:a |the )?"
|
||||
r"(?:comprehensive |clear |good |detailed )?answer\b|"
|
||||
r"now i can (?:confidently )?(?:give|provide|state|answer)|"
|
||||
r"perfect[!,]|great[!,]|excellent[!,]|" # narration exclamations
|
||||
r"as an orchestrator|"
|
||||
r"whose job is to (?:route|delegate|orchestrate)|"
|
||||
r"testing whether i (?:need|have|should|must)\b|"
|
||||
r"i(?:'m| am)? (?:required|instructed|supposed|expected) to\b"
|
||||
r")[^\n]*(?:\n|$)"
|
||||
)
|
||||
# A leading task-restatement opener ("The user is asking…") — the format
|
||||
# explicitly bans this meta-text; drop it when it opens the reasoning.
|
||||
_META_OPEN_RE = re.compile(
|
||||
r"(?is)^\s*(?:the user (?:is|was|wants|has|needs|would|'s)\b[^.\n]*[.\n]+\s*)+"
|
||||
)
|
||||
|
||||
|
||||
def _scrub_meta(text: str) -> str:
|
||||
"""Remove harness-leakage lines and a leading 'The user is asking…' meta
|
||||
opener from a reasoning block. Both are disallowed by the system prompt, so
|
||||
they must not survive into the supervised target. Preserves a leading
|
||||
``<think>`` marker so the block stays well-formed."""
|
||||
if not text:
|
||||
return text
|
||||
think = ""
|
||||
body = text
|
||||
m = re.match(r"(?is)^\s*(<think>)\s*", body)
|
||||
if m:
|
||||
think = "<think>\n"
|
||||
body = body[m.end() :]
|
||||
body = _LEAK_RE.sub("", body)
|
||||
body = _META_OPEN_RE.sub("", body).lstrip()
|
||||
return (think + body).strip() if think else body.strip()
|
||||
|
||||
|
||||
# Control / special tokens that must never survive into a training target's
|
||||
# final answer. Three shapes, in order of alternation:
|
||||
# 1. closed pipe token — ``<|im_end|>``, ``<|im_start|>``, ``<|eot_id|>``,
|
||||
# ``<|end_of_turn|>``, ``<|"|>`` (the stray-quote leak).
|
||||
# 2. UNCLOSED pipe token — ``<|tool_call>`` (opened with ``<|`` but closed on a
|
||||
# bare ``>`` because the decode was cut before the second pipe).
|
||||
# 3. named angle special tokens — gemma ``<start_of_turn>`` / ``<end_of_turn>``
|
||||
# and the sentinel set ``<eos> <bos> <pad> <unk> <s> </s>``.
|
||||
# Deliberately narrow: only the ``<|...|>`` pipe form and this known name list
|
||||
# match, so legitimate ``<``/``>`` in math/code (``x < 3 and y > 2``) is left
|
||||
# untouched.
|
||||
_CONTROL_TOKEN_RE = re.compile(
|
||||
r"<\|[^>]*?\|>" # closed pipe token
|
||||
r"|<\|[^|>]*>" # unclosed pipe token
|
||||
r"|</?(?:start_of_turn|end_of_turn|eos|bos|pad|unk|s)>" # named angle tokens
|
||||
)
|
||||
|
||||
|
||||
def _strip_control_tokens(text: str) -> str:
|
||||
"""Delete residual control/special tokens from an answer so a good answer
|
||||
with a stray token is salvaged rather than dropped. Loops (bounded) because
|
||||
removing one match can expose a nested/overlapping leftover (e.g.
|
||||
``<|<|im_end|>|>``). Re-strips whitespace at the end."""
|
||||
if not text:
|
||||
return text
|
||||
out = text
|
||||
for _ in range(4):
|
||||
stripped = _CONTROL_TOKEN_RE.sub("", out)
|
||||
if stripped == out:
|
||||
break
|
||||
out = stripped
|
||||
return out.strip()
|
||||
|
||||
|
||||
def _final_answer_block(text: str) -> str:
|
||||
"""Render the final-answer assistant message with a single clean
|
||||
``FINAL_ANSWER: <value>`` line. De-boxes the answer, strips leaked control
|
||||
tokens, and normalizes whatever spelling the model used (``FINALANSWER``,
|
||||
``FINAL ANSWER``, no colon, …) to the exact tag. Keeps the think block at
|
||||
most once."""
|
||||
text = _debox(_normalize_think((text or "").strip()))
|
||||
# Normalize any final-answer marker spelling to the canonical form; take the
|
||||
# LAST one as the real answer (idempotent — never emits two tags).
|
||||
marks = list(re.finditer(r"(?im)FINAL[_\s]?ANSWER\s*:?", text))
|
||||
if marks:
|
||||
last = marks[-1]
|
||||
answer = _strip_control_tokens(text[last.end() :].strip())
|
||||
# Final turn = the bare short answer only. Drop EVERYTHING before
|
||||
# FINAL_ANSWER — both the visible "Perfect! Based on the model's response…"
|
||||
# narration AND the final-turn <think>, which is just confirmatory
|
||||
# "both models agree… I've verified…" fluff. The routing reasoning already
|
||||
# lives in the earlier tool-call turns, so nothing of value is lost, and
|
||||
# this kills all final-turn narration deterministically (any wording).
|
||||
return f"FINAL_ANSWER: {answer}"
|
||||
if "</think>" in text:
|
||||
# No explicit marker: the answer is whatever follows the final </think>.
|
||||
# Drop the think block (same rationale as above) — keep only the answer.
|
||||
_, _, answer = text.rpartition("</think>")
|
||||
answer = _strip_control_tokens(answer.strip())
|
||||
return (
|
||||
f"FINAL_ANSWER: {answer}"
|
||||
if answer
|
||||
else f"FINAL_ANSWER: {_strip_control_tokens(_scrub_meta(text))}"
|
||||
)
|
||||
return f"FINAL_ANSWER: {_strip_control_tokens(text)}"
|
||||
|
||||
|
||||
def trajectory_to_record(
|
||||
task_id: str,
|
||||
question: str,
|
||||
tools: List[ExpertTool],
|
||||
rollout: UnifiedRollout,
|
||||
*,
|
||||
reward: float = 0.0,
|
||||
domain: str = "unknown",
|
||||
) -> Dict[str, Any]:
|
||||
"""Convert a passing :class:`UnifiedRollout` into one SFT JSONL record."""
|
||||
conversations: List[Dict[str, str]] = [
|
||||
{"role": "system", "content": _system_prompt(tools, rollout)},
|
||||
{"role": "user", "content": f"Problem: {question}"},
|
||||
]
|
||||
|
||||
for turn in rollout.turns:
|
||||
if turn.tool_name is None:
|
||||
# Final-answer turn. turn.reasoning is the model's actual output for
|
||||
# this turn (which already contains the answer); render it once.
|
||||
conversations.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": _final_answer_block(
|
||||
turn.reasoning or rollout.final_answer
|
||||
),
|
||||
}
|
||||
)
|
||||
continue
|
||||
tag = tool_call_tag(turn.tool_name, turn.arguments)
|
||||
reasoning = _scrub_meta(_normalize_think((turn.reasoning or "").rstrip()))
|
||||
conversations.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": (reasoning + "\n" + tag).strip(),
|
||||
}
|
||||
)
|
||||
# Store the observation EXACTLY as the orchestrator saw it — the rollout
|
||||
# caps tool output at _OBS_CAP before the model reads it (rollout.py), but
|
||||
# this serializer used to store the FULL text. That trained the student on
|
||||
# context the teacher never had, and produced monster rows (a 212k-char
|
||||
# raw-HTML http_request dump = ~50k tokens of <!DOCTYPE> boilerplate) that
|
||||
# then got BEHEADED by the trainer's max-seq. Train on what the policy saw.
|
||||
obs = turn.observation or ""
|
||||
if len(obs) > _ROLLOUT_OBS_CAP:
|
||||
obs = obs[:_ROLLOUT_OBS_CAP] + "\n…[truncated]"
|
||||
conversations.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"name": turn.tool_name,
|
||||
"content": obs,
|
||||
}
|
||||
)
|
||||
|
||||
# If the rollout terminated on max_turns (no None turn), append the answer.
|
||||
if not rollout.turns or rollout.turns[-1].tool_name is not None:
|
||||
conversations.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": _final_answer_block(rollout.final_answer),
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"conversations": conversations,
|
||||
"task_id": task_id,
|
||||
"domain": domain,
|
||||
"reward": reward,
|
||||
"metrics": {
|
||||
"cost_usd": rollout.cost_usd,
|
||||
"tokens": rollout.tokens,
|
||||
"num_tool_calls": rollout.num_tool_calls,
|
||||
"num_turns": len(rollout.turns),
|
||||
# Present only when experts were anonymized: opaque label -> real
|
||||
# tool name, so analysis can recover which model was actually picked.
|
||||
**({"anon_map": rollout.anon_map} if rollout.anon_map else {}),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
__all__ = ["trajectory_to_record"]
|
||||
@@ -0,0 +1,388 @@
|
||||
"""Correctness verifier for orchestrator reasoning tasks (:mod:`.datasets`).
|
||||
|
||||
The verifier dispatches on :attr:`Task.domain`:
|
||||
|
||||
- **math** — normalize and compare. Extract ``\\boxed{...}`` from both sides,
|
||||
try ``sympy`` for symbolic / numeric equality, and fall back to a normalized
|
||||
string / number match. No network.
|
||||
- **code** — best-effort: if a short expected output / answer exists, do a
|
||||
normalized substring match; otherwise defer to the LLM judge. (We don't run
|
||||
arbitrary code here.)
|
||||
- **science / medical / chat / misc / unknown** — Gemini LLM judge given the
|
||||
question + gold answer + candidate, asked for ``PASS`` / ``FAIL``. When no
|
||||
``GEMINI_API_KEY`` is set (e.g. offline tests) it falls back to a normalized
|
||||
string / token-F1 >= 0.6 match.
|
||||
|
||||
The OpenAI key is dead, so the LLM judge talks to Gemini through its
|
||||
OpenAI-compatible endpoint. Normalization helpers are copied (not imported)
|
||||
from ``hotpotqa.py`` to keep this module self-contained.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import string
|
||||
from typing import Any, Callable, Dict, Optional
|
||||
|
||||
from .datasets import Task
|
||||
|
||||
# LLM judge model. Switched from Gemini (free-tier rate limit stalled parallel gen)
|
||||
# to Anthropic Haiku: fast (~0.6s), direct PASS/FAIL, high rate limits.
|
||||
LOGGER = logging.getLogger(__name__)
|
||||
|
||||
JUDGE_MODEL = "claude-haiku-4-5-20251001"
|
||||
# The Anthropic SDK's own retry does exponential backoff and honours Retry-After.
|
||||
# A judge call that waits a few seconds beats a silently-wrong label.
|
||||
JUDGE_MAX_RETRIES = 6
|
||||
# Reused across calls — creating a fresh Anthropic client per judge call leaks an
|
||||
# httpx connection pool (CLOSE-WAIT pileup under parallel generation). Thread-safe.
|
||||
_JUDGE_CLIENT = None
|
||||
# Every swallowed judge error, so a run can never again silently mislabel while
|
||||
# reporting a clean bill of health.
|
||||
_JUDGE_FAILURES: list[str] = []
|
||||
|
||||
|
||||
# --- normalization (copied from hotpotqa.py — keep self-contained) -----------
|
||||
|
||||
|
||||
def _normalize_answer(s: str) -> str:
|
||||
"""Lowercase, strip punctuation / articles / extra whitespace (SQuAD/HotpotQA)."""
|
||||
s = s.lower()
|
||||
s = "".join(ch for ch in s if ch not in set(string.punctuation))
|
||||
s = re.sub(r"\b(a|an|the)\b", " ", s)
|
||||
return " ".join(s.split())
|
||||
|
||||
|
||||
def _f1(pred: str, gold: str) -> float:
|
||||
pt, gt = _normalize_answer(pred).split(), _normalize_answer(gold).split()
|
||||
if not pt or not gt:
|
||||
return float(pt == gt)
|
||||
common: Dict[str, int] = {}
|
||||
for w in pt:
|
||||
if w in gt:
|
||||
common[w] = common.get(w, 0) + 1
|
||||
num_same = sum(common.values())
|
||||
if num_same == 0:
|
||||
return 0.0
|
||||
precision = num_same / len(pt)
|
||||
recall = num_same / len(gt)
|
||||
return 2 * precision * recall / (precision + recall)
|
||||
|
||||
|
||||
def _string_or_f1(prediction: str, gold: str, *, f1_threshold: float = 0.6) -> bool:
|
||||
"""Normalized whole-word substring OR token-F1 >= threshold."""
|
||||
np_, ng = _normalize_answer(prediction), _normalize_answer(gold)
|
||||
if not ng:
|
||||
return False
|
||||
if f" {ng} " in f" {np_} ":
|
||||
return True
|
||||
return _f1(prediction, gold) >= f1_threshold
|
||||
|
||||
|
||||
# --- math --------------------------------------------------------------------
|
||||
|
||||
_BOXED_RE = re.compile(r"\\boxed\{")
|
||||
_NUM_RE = re.compile(r"-?\d+(?:\.\d+)?")
|
||||
|
||||
|
||||
def _extract_boxed(text: str) -> Optional[str]:
|
||||
"""Return the contents of the last ``\\boxed{...}`` (brace-balanced)."""
|
||||
last = None
|
||||
for m in _BOXED_RE.finditer(text):
|
||||
start = m.end()
|
||||
depth = 1
|
||||
i = start
|
||||
while i < len(text) and depth > 0:
|
||||
c = text[i]
|
||||
if c == "{":
|
||||
depth += 1
|
||||
elif c == "}":
|
||||
depth -= 1
|
||||
i += 1
|
||||
if depth == 0:
|
||||
last = text[start : i - 1].strip()
|
||||
return last
|
||||
|
||||
|
||||
# \frac{a}{b} / \dfrac{a}{b} / \tfrac{a}{b} -> (a)/(b)
|
||||
_FRAC_RE = re.compile(r"\\[dt]?frac\s*\{([^{}]+)\}\s*\{([^{}]+)\}")
|
||||
|
||||
|
||||
def _clean_math(s: str) -> str:
|
||||
s = s.strip()
|
||||
boxed = _extract_boxed(s)
|
||||
if boxed is not None:
|
||||
s = boxed
|
||||
# Common LaTeX wrappers / delimiters.
|
||||
s = s.replace("$", "").replace("\\!", "").replace("\\,", "").replace("\\;", "")
|
||||
s = s.replace("\\left", "").replace("\\right", "")
|
||||
# Normalize LaTeX fractions to plain division. Without this, `\frac{15}{64}`
|
||||
# never becomes `15/64`, the exact/numeric comparisons both fail, and we fall
|
||||
# through to a crude "grab the first number" match — which reads `15` out of
|
||||
# BOTH sides and returns True. That is a FALSE POSITIVE: it passed
|
||||
# gold \frac{15}{64} vs model 15/99 (different fraction)
|
||||
# gold \frac{1}{2} vs model 1/3 (different denominator)
|
||||
# gold \frac{15}{64} vs model 15 (just the numerator)
|
||||
# i.e. any wrong answer sharing a numerator with the gold was marked CORRECT
|
||||
# and fed into training. It also caused the mirror false-negative
|
||||
# (`-\dfrac{73}{143}` vs `-73/143`), because the minus sits outside the frac
|
||||
# and the number-grab loses it. (Audit 2026-07-13.)
|
||||
for _ in range(3): # nested fractions
|
||||
s, n = _FRAC_RE.subn(r"(\1)/(\2)", s)
|
||||
if not n:
|
||||
break
|
||||
s = s.strip().strip("$ ")
|
||||
return s
|
||||
|
||||
|
||||
# A pure rational expression: -73/143, (15)/(64), 3.5/2 — safe to evaluate.
|
||||
_RATIONAL_RE = re.compile(r"^-?\(?\s*-?\d+(?:\.\d+)?\s*\)?\s*/\s*\(?\s*-?\d+(?:\.\d+)?\s*\)?$")
|
||||
|
||||
|
||||
def _to_float(s: str) -> Optional[float]:
|
||||
s = s.replace(",", "").strip()
|
||||
try:
|
||||
return float(s)
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
# Evaluate a plain fraction rather than falling through to the number-grab
|
||||
# below, which would read only the NUMERATOR and happily equate 15/64 with
|
||||
# 15/99. Restricted by regex to digits and one '/', so nothing else is eval'd.
|
||||
if _RATIONAL_RE.match(s):
|
||||
try:
|
||||
num, den = s.replace("(", "").replace(")", "").split("/")
|
||||
d = float(den)
|
||||
if d != 0:
|
||||
return float(num) / d
|
||||
except (ValueError, ZeroDivisionError):
|
||||
pass
|
||||
m = _NUM_RE.search(s)
|
||||
if m:
|
||||
try:
|
||||
return float(m.group(0))
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _math_equal(prediction: str, gold: str) -> bool:
|
||||
pred = _clean_math(prediction)
|
||||
g = _clean_math(gold)
|
||||
if not g:
|
||||
return False
|
||||
|
||||
# Exact normalized-string shortcut.
|
||||
if pred.replace(" ", "") == g.replace(" ", "") and pred:
|
||||
return True
|
||||
|
||||
# Numeric comparison (handles "42" vs "42.0" vs trailing prose).
|
||||
pf, gf = _to_float(pred), _to_float(g)
|
||||
if pf is not None and gf is not None:
|
||||
if abs(pf - gf) <= 1e-6 * max(1.0, abs(gf)):
|
||||
return True
|
||||
|
||||
# NOTE: a sympy/parse_latex symbolic-equality block used to live here, but
|
||||
# sympy.simplify / antlr parse_latex can hang on pathological \boxed{} answers
|
||||
# while holding the GIL -> deadlocks the whole threaded rejection sampler
|
||||
# (all worker threads freeze in futex_wait, server goes idle, 0 records).
|
||||
# Symbolic equivalence is now deferred to the Gemini judge in verify_answer().
|
||||
|
||||
# Last resort: whole-word substring of the gold in the prediction.
|
||||
return _string_or_f1(prediction, g, f1_threshold=0.9)
|
||||
|
||||
|
||||
# --- LLM judge (Gemini, OpenAI-compatible) -----------------------------------
|
||||
|
||||
|
||||
def _judge_route() -> Optional[tuple]:
|
||||
"""(base_url, api_key, model) for the judge — OpenRouter if configured.
|
||||
|
||||
The judge must NOT share a rate-limit bucket with the orchestrator. They're
|
||||
the same model (Haiku), but the orchestrator makes 5-10 calls per rollout and
|
||||
the judge makes 1; when they contend, the judge 429s, returns None, and the
|
||||
caller falls back to string matching — silently marking correct answers WRONG.
|
||||
Set ``OJ_JUDGE_VIA_OPENROUTER=1`` to give the judge its own quota (~$2.50 per
|
||||
3000 rollouts).
|
||||
"""
|
||||
if os.environ.get("OJ_JUDGE_VIA_OPENROUTER") == "1":
|
||||
key = os.environ.get("OPENROUTER_API_KEY")
|
||||
if key:
|
||||
return ("https://openrouter.ai/api/v1", key, "anthropic/claude-haiku-4.5")
|
||||
key = os.environ.get("ANTHROPIC_API_KEY")
|
||||
return ("https://api.anthropic.com/v1/", key, JUDGE_MODEL) if key else None
|
||||
|
||||
|
||||
def _gemini_judge(task: Task, prediction: str) -> Optional[bool]:
|
||||
"""Ask the LLM judge PASS/FAIL. Returns None if unavailable (caller falls back)."""
|
||||
route = _judge_route()
|
||||
if route is None:
|
||||
return None
|
||||
base_url, api_key, model = route
|
||||
try:
|
||||
from openai import OpenAI # OpenAI-compatible: works for both routes
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
try:
|
||||
# Judge = Haiku (fast ~0.6s, direct PASS/FAIL), reached over whichever route
|
||||
# _judge_route() picked.
|
||||
#
|
||||
# This used to be fail-fast (max_retries=0, timeout=15) so a slow judge
|
||||
# couldn't stall the threaded sampler. That was actively destroying data:
|
||||
# under parallel generation the judge 429s, this returns None, the caller
|
||||
# falls back to string/f1, and a CORRECT answer gets marked WRONG. The
|
||||
# failure is invisible — the exception is swallowed and never logged, so
|
||||
# the run reports "0 rate limits" while silently mislabelling.
|
||||
# (Audit 2026-07-12: 3 consecutive judge calls at concurrency 100 -> 2x429.)
|
||||
#
|
||||
# Retries alone were NOT enough: with the judge on the same bucket as the
|
||||
# orchestrator, the retries just parked every worker thread in backoff and
|
||||
# stalled the run. The real fix is giving the judge its own quota
|
||||
# (OJ_JUDGE_VIA_OPENROUTER=1); retries are the belt to that suspenders.
|
||||
global _JUDGE_CLIENT
|
||||
if _JUDGE_CLIENT is None:
|
||||
_JUDGE_CLIENT = OpenAI(
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
max_retries=JUDGE_MAX_RETRIES,
|
||||
timeout=30,
|
||||
)
|
||||
client = _JUDGE_CLIENT
|
||||
# Grade on SEMANTIC equivalence, not surface form. The bare "correct and
|
||||
# consistent with the gold" instruction was far too strict: an audit
|
||||
# (2026-07-12) found ~half of all FAIL verdicts were correct answers the
|
||||
# judge rejected purely for phrasing —
|
||||
# gold "0" vs "Order w.r.t. A = 1; Order w.r.t. B = 0"
|
||||
# gold "use plt.figtext()" vs "Use fig.text()" (same matplotlib call)
|
||||
# gold "<terse fragment>" vs the same fact in a full sentence
|
||||
# Those FAILs silently threw away good training data. The guardrails at the
|
||||
# bottom keep it from swinging lenient: a different VALUE still fails.
|
||||
prompt = (
|
||||
"You are grading a candidate answer against a gold reference.\n\n"
|
||||
"The gold is often TERSE — a single word ('Yes'), a bare value ('0'), a "
|
||||
"compact expression ('4 > 3 > 1 > 2'). The candidate is typically a full "
|
||||
"sentence that states the same thing and adds correct supporting detail. "
|
||||
"That is a PASS.\n\n"
|
||||
"Your ONLY question is: does the candidate AGREE with the gold on the "
|
||||
"point the gold makes?\n\n"
|
||||
"PASS when:\n"
|
||||
" * the candidate says the same thing more verbosely, or explains it;\n"
|
||||
" * it uses different notation, exactly-converting units, or equivalent "
|
||||
"mathematical form (15/64 vs \\frac{15}{64});\n"
|
||||
" * the gold is a fragment and the candidate states the same fact in "
|
||||
"prose, or answers additional parts of the question as well;\n"
|
||||
" * for code: a different but equivalent implementation, or a different "
|
||||
"API call with the same effect.\n\n"
|
||||
"Do NOT fail the candidate merely for saying MORE than the gold, for "
|
||||
"being longer, for adding a derivation, or for answering other parts of "
|
||||
"the question too. Extra correct detail is never a reason to fail.\n\n"
|
||||
"FAIL only when the candidate gives a DIFFERENT VALUE, contradicts the "
|
||||
"gold, or does not actually answer the question.\n\n"
|
||||
f"Question:\n{task.question}\n\n"
|
||||
f"Gold answer:\n{task.answer}\n\n"
|
||||
f"Candidate answer:\n{prediction}\n\n"
|
||||
"Reply with exactly one word — PASS or FAIL.\nVerdict:"
|
||||
)
|
||||
# temperature=0: a grader must be DETERMINISTIC. Without it the SDK default
|
||||
# (1.0) sampled the PASS/FAIL token, so the same (question, gold, candidate)
|
||||
# could be graded differently on two runs — the label became a coin flip on
|
||||
# any borderline answer. (Audit 2026-07-13: gold "0" vs "order w.r.t. B is 0"
|
||||
# graded PASS in one run and FAIL in the next.)
|
||||
resp = client.chat.completions.create(
|
||||
model=model,
|
||||
max_tokens=8,
|
||||
temperature=0,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
)
|
||||
verdict = (resp.choices[0].message.content or "").strip().upper()
|
||||
if "PASS" in verdict:
|
||||
return True
|
||||
if "FAIL" in verdict:
|
||||
return False
|
||||
return None
|
||||
except Exception as exc: # noqa: BLE001
|
||||
# LOUDLY. Swallowing this silently is how a rate-limited judge quietly
|
||||
# mislabelled correct answers as wrong for a whole run while the monitor
|
||||
# cheerfully reported "0 rate limits" — the exception never reached a log.
|
||||
_JUDGE_FAILURES.append(type(exc).__name__)
|
||||
LOGGER.warning(
|
||||
"JUDGE CALL FAILED (%d so far this run) — falling back to string match, "
|
||||
"which will likely mark a correct answer WRONG: %s",
|
||||
len(_JUDGE_FAILURES),
|
||||
str(exc)[:160],
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
# --- public API --------------------------------------------------------------
|
||||
|
||||
|
||||
# The orchestrator is instructed to end with "FINAL_ANSWER: <answer>", so the raw
|
||||
# rollout text is "<reasoning…>\n\nFINAL_ANSWER: A". Everything before the marker
|
||||
# (and the marker itself) must come off before we compare, or we're matching the
|
||||
# model's prose against a one-token gold.
|
||||
_FINAL_ANSWER_RE = re.compile(r"(?im)FINAL[_\s]?ANSWER\s*:?")
|
||||
|
||||
|
||||
def _answer_only(prediction: str) -> str:
|
||||
"""The text AFTER the last FINAL_ANSWER marker (or the whole thing if absent).
|
||||
|
||||
Without this, ``_math_equal("FINAL_ANSWER: A", "A")`` is False while
|
||||
``_math_equal("A", "A")`` is True — the literal marker was never stripped, so
|
||||
every short non-numeric answer (multiple-choice letters, symbolic values) was
|
||||
auto-marked WRONG. Numeric answers happened to survive via float parsing,
|
||||
which is why this hid for so long. (audit 2026-07-12)
|
||||
"""
|
||||
marks = list(_FINAL_ANSWER_RE.finditer(prediction))
|
||||
return (prediction[marks[-1].end() :] if marks else prediction).strip()
|
||||
|
||||
|
||||
def verify_answer(task: Task, prediction: str) -> bool:
|
||||
"""Domain-dispatched correctness check for ``prediction`` vs ``task.answer``."""
|
||||
pred = _answer_only(prediction or "")
|
||||
if not pred or not (task.answer or "").strip():
|
||||
return False
|
||||
|
||||
domain = (task.domain or "").lower()
|
||||
|
||||
if domain == "math":
|
||||
# string + numeric + f1 only. sympy removed (GIL-deadlocked the threaded
|
||||
# sampler); Gemini judge NOT used here either (32 tasks x N samples of math
|
||||
# judge calls throttle Gemini and stall the whole run). Accept slightly lower
|
||||
# math yield for a fast, hang-free verifier.
|
||||
return _math_equal(pred, task.answer)
|
||||
|
||||
if domain == "code":
|
||||
# Best-effort: short gold -> normalized substring; otherwise LLM judge.
|
||||
gold = task.answer.strip()
|
||||
if len(gold) <= 200:
|
||||
if _string_or_f1(pred, gold, f1_threshold=0.8):
|
||||
return True
|
||||
judged = _gemini_judge(task, pred)
|
||||
if judged is not None:
|
||||
return judged
|
||||
return _string_or_f1(pred, gold, f1_threshold=0.6)
|
||||
|
||||
# science / medical / chat / misc / unknown -> LLM judge, then F1 fallback.
|
||||
judged = _gemini_judge(task, pred)
|
||||
if judged is not None:
|
||||
return judged
|
||||
return _string_or_f1(pred, task.answer, f1_threshold=0.6)
|
||||
|
||||
|
||||
def make_verifier() -> Callable[[Any, Any], bool]:
|
||||
"""Return ``verify_fn(task, rollout) -> bool`` for the rejection sampler."""
|
||||
|
||||
def verify(task: Any, rollout: Any) -> bool:
|
||||
ans = getattr(rollout, "final_answer", "") or ""
|
||||
return verify_answer(task, ans)
|
||||
|
||||
return verify
|
||||
|
||||
|
||||
__all__ = [
|
||||
"make_verifier",
|
||||
"verify_answer",
|
||||
]
|
||||
@@ -42,8 +42,22 @@ _MATH_FUNCS = {
|
||||
"sin": math.sin,
|
||||
"cos": math.cos,
|
||||
"tan": math.tan,
|
||||
"asin": math.asin,
|
||||
"acos": math.acos,
|
||||
"atan": math.atan,
|
||||
"atan2": math.atan2,
|
||||
"sinh": math.sinh,
|
||||
"cosh": math.cosh,
|
||||
"tanh": math.tanh,
|
||||
"radians": math.radians,
|
||||
"degrees": math.degrees,
|
||||
"exp": math.exp,
|
||||
"fabs": math.fabs,
|
||||
"hypot": math.hypot,
|
||||
"factorial": math.factorial,
|
||||
"pi": math.pi,
|
||||
"e": math.e,
|
||||
"tau": math.tau,
|
||||
"ceil": math.ceil,
|
||||
"floor": math.floor,
|
||||
}
|
||||
|
||||
@@ -116,11 +116,31 @@ class WebSearchTool(BaseTool):
|
||||
return text
|
||||
|
||||
def _duckduckgo_search(self, query: str, max_results: int) -> str:
|
||||
"""Search using DuckDuckGo as fallback."""
|
||||
"""Search using DuckDuckGo as fallback.
|
||||
|
||||
ddgs queries several engines (yandex/yahoo/brave/...) that frequently
|
||||
hang; with no timeout this blocks the whole rollout. Cap each engine
|
||||
request (DDGS timeout) AND wrap the call in a hard wall-clock deadline so
|
||||
a flaky search fast-fails instead of stalling data generation.
|
||||
"""
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
from ddgs import DDGS
|
||||
|
||||
ddgs = DDGS()
|
||||
raw_results = list(ddgs.text(query, max_results=max_results))
|
||||
def _run():
|
||||
ddgs = DDGS(timeout=5)
|
||||
return list(ddgs.text(query, max_results=max_results))
|
||||
|
||||
# Don't use `with` — its shutdown(wait=True) blocks on the hung thread,
|
||||
# defeating the deadline. submit, wait up to 8s, then abandon the thread
|
||||
# (it finishes in the background harmlessly) and fast-fail.
|
||||
_ex = ThreadPoolExecutor(max_workers=1)
|
||||
try:
|
||||
raw_results = _ex.submit(_run).result(timeout=8)
|
||||
except Exception:
|
||||
_ex.shutdown(wait=False, cancel_futures=True)
|
||||
return "[web_search: no results (search backend timed out)]"
|
||||
_ex.shutdown(wait=False)
|
||||
results = []
|
||||
for r in raw_results:
|
||||
title = r.get("title", "Untitled")
|
||||
|
||||
@@ -0,0 +1,154 @@
|
||||
"""Offline tests for the reasoning-SFT dataset loaders.
|
||||
|
||||
All loaders take a ``source=`` iterable of raw row dicts so we never touch the
|
||||
network: we hand them fake rows shaped like the real GeneralThought /
|
||||
OpenThoughts3 schemas and assert the normalized :class:`Task` fields.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import (
|
||||
Task,
|
||||
load_generalthought,
|
||||
load_openthoughts,
|
||||
)
|
||||
|
||||
# --- fake rows mirroring the real HF schemas ---------------------------------
|
||||
|
||||
GENERALTHOUGHT_ROWS = [
|
||||
{
|
||||
"question_id": 1,
|
||||
"question": "What is 2 + 2?",
|
||||
"reference_answer": "4",
|
||||
"model_answer": "The answer is 4.",
|
||||
"task": "High School Math",
|
||||
"question_source": "Numina/NuminaMath",
|
||||
},
|
||||
{
|
||||
"question_id": 2,
|
||||
"question": "How do the neural respiratory centers operate?",
|
||||
"reference_answer": "Via the medulla oblongata.",
|
||||
"model_answer": "long answer ...",
|
||||
"task": "NHS QA",
|
||||
"question_source": "CogStack/NHSQA",
|
||||
},
|
||||
{
|
||||
"question_id": 3,
|
||||
"question": "Sort a list.",
|
||||
"reference_answer": "",
|
||||
"model_answer": "use sorted()",
|
||||
"task": "Sorting",
|
||||
"question_source": "BAAI/TACO",
|
||||
},
|
||||
# Should be skipped: empty question and answer.
|
||||
{"question_id": 4, "question": "", "reference_answer": "", "model_answer": ""},
|
||||
]
|
||||
|
||||
OPENTHOUGHTS_ROWS = [
|
||||
{
|
||||
"domain": "code",
|
||||
"source": "nvidia/OpenCodeReasoning",
|
||||
"difficulty": 7,
|
||||
"conversations": [
|
||||
{"from": "human", "value": "Write a program that prints hi."},
|
||||
{"from": "gpt", "value": "<think>reason</think> Final: print('hi')"},
|
||||
],
|
||||
},
|
||||
{
|
||||
"domain": "math",
|
||||
"source": "Numina/NuminaMath",
|
||||
"difficulty": 5,
|
||||
"conversations": [
|
||||
{"from": "human", "value": "What is 6 times 7?"},
|
||||
{
|
||||
"from": "gpt",
|
||||
"value": "<think>6*7=42</think> The answer is \\boxed{42}.",
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"domain": "science",
|
||||
"source": "sci",
|
||||
"difficulty": 6,
|
||||
"conversations": [
|
||||
{"from": "human", "value": "What gas do plants release?"},
|
||||
{"from": "gpt", "value": "Oxygen."},
|
||||
],
|
||||
},
|
||||
# Should be skipped: no human/gpt turns.
|
||||
{"domain": "math", "conversations": []},
|
||||
# Should be skipped: codegolf is an UNUSABLE_SOURCE (adversarial, unscorable).
|
||||
{
|
||||
"domain": "code",
|
||||
"source": "stackexchange_codegolf",
|
||||
"conversations": [
|
||||
{"from": "human", "value": "Golf: print hi."},
|
||||
{"from": "gpt", "value": "print('hi')"},
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def test_load_generalthought_fields_and_domains():
|
||||
# exclude_domains=() so the lookup-shaped medical row survives for the
|
||||
# classification asserts below — by DEFAULT it is dropped (see end of test).
|
||||
tasks = list(
|
||||
load_generalthought(n=10, source=GENERALTHOUGHT_ROWS, exclude_domains=())
|
||||
)
|
||||
assert all(isinstance(t, Task) for t in tasks)
|
||||
# Row 4 (empty) is dropped; remaining 3 yielded.
|
||||
assert len(tasks) == 3
|
||||
by_q = {t.question: t for t in tasks}
|
||||
assert by_q["What is 2 + 2?"].answer == "4"
|
||||
assert by_q["What is 2 + 2?"].domain == "math"
|
||||
assert by_q["What is 2 + 2?"].instruction == "What is 2 + 2?"
|
||||
assert by_q["How do the neural respiratory centers operate?"].domain == "medical"
|
||||
# No reference_answer -> falls back to model_answer; TACO source -> code.
|
||||
assert by_q["Sort a list."].answer == "use sorted()"
|
||||
assert by_q["Sort a list."].domain == "code"
|
||||
# By DEFAULT the lookup-shaped medical row is dropped (LOOKUP_DOMAINS), so the
|
||||
# same source yields only the math + code reasoning tasks.
|
||||
default = list(load_generalthought(n=10, source=GENERALTHOUGHT_ROWS))
|
||||
assert {t.domain for t in default} == {"math", "code"}
|
||||
|
||||
|
||||
def test_load_generalthought_respects_n():
|
||||
tasks = list(load_generalthought(n=2, source=GENERALTHOUGHT_ROWS))
|
||||
assert len(tasks) == 2
|
||||
|
||||
|
||||
def test_load_openthoughts_fields_and_domains():
|
||||
tasks = list(
|
||||
load_openthoughts(n_code=5, n_math=5, n_science=5, source=OPENTHOUGHTS_ROWS)
|
||||
)
|
||||
assert all(isinstance(t, Task) for t in tasks)
|
||||
assert len(tasks) == 3
|
||||
by_dom = {t.domain: t for t in tasks}
|
||||
assert set(by_dom) == {"code", "math", "science"}
|
||||
# <think> stripped; boxed answer extracted.
|
||||
assert by_dom["math"].answer == "42"
|
||||
assert by_dom["math"].question == "What is 6 times 7?"
|
||||
# Code: think stripped, visible solution kept.
|
||||
assert "print('hi')" in by_dom["code"].answer
|
||||
assert by_dom["science"].answer == "Oxygen."
|
||||
|
||||
|
||||
def test_load_openthoughts_quota_caps_per_domain():
|
||||
rows = [
|
||||
{
|
||||
"domain": "code",
|
||||
"conversations": [
|
||||
{"from": "human", "value": f"q{i}"},
|
||||
{"from": "gpt", "value": f"a{i}"},
|
||||
],
|
||||
}
|
||||
for i in range(10)
|
||||
]
|
||||
tasks = list(load_openthoughts(n_code=3, n_math=0, n_science=0, source=rows))
|
||||
assert len(tasks) == 3
|
||||
assert all(t.domain == "code" for t in tasks)
|
||||
|
||||
|
||||
def test_task_dataclass_instruction_property():
|
||||
t = Task(task_id="x", question="hi?", answer="yes", domain="misc")
|
||||
assert t.instruction == "hi?"
|
||||
@@ -0,0 +1,150 @@
|
||||
"""Offline tests for the rejection-sampling SFT pipeline (unified tools)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog, tools_by_name
|
||||
from openjarvis.agents.hybrid.toolorchestra.rollout import (
|
||||
UnifiedRollout,
|
||||
UnifiedTurn,
|
||||
run_unified_rollout,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import Task
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
|
||||
generate_sft_dataset,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import (
|
||||
trajectory_to_record,
|
||||
)
|
||||
|
||||
|
||||
def test_run_unified_rollout_terminates_on_no_tool_call():
|
||||
tools = orchestrator_catalog()
|
||||
by = tools_by_name(tools)
|
||||
name = tools[0].name # any real expert tool from the live catalog
|
||||
assert name in by
|
||||
|
||||
scripted = [
|
||||
("reason 1\n", [(name, {"input": "do step 1"})], 5, 5),
|
||||
("here is the answer", [], 3, 3), # no tool call -> terminate
|
||||
]
|
||||
calls = iter(scripted)
|
||||
|
||||
def call_orch(messages, specs):
|
||||
# Signature is now (messages, specs): the rollout drives a running
|
||||
# system/user/assistant/tool conversation, not a flattened (system, user).
|
||||
return next(calls)
|
||||
|
||||
def dispatch(tool, args):
|
||||
return (f"OBS for {tool.name}", 0.01, 10, False)
|
||||
|
||||
roll = run_unified_rollout(
|
||||
"What is X?",
|
||||
tools,
|
||||
call_orchestrator=call_orch,
|
||||
dispatch=dispatch,
|
||||
max_turns=5,
|
||||
)
|
||||
assert roll.final_answer == "here is the answer"
|
||||
assert roll.num_tool_calls == 1
|
||||
assert roll.tool_calls() == [(name, {"input": "do step 1"})]
|
||||
assert abs(roll.cost_usd - 0.01) < 1e-9
|
||||
|
||||
|
||||
def test_serialize_record_shape_and_tool_call_tags():
|
||||
tools = orchestrator_catalog()
|
||||
roll = UnifiedRollout(
|
||||
turns=[
|
||||
UnifiedTurn(
|
||||
reasoning="think",
|
||||
tool_name="qwen3_32b",
|
||||
arguments={"input": "q"},
|
||||
observation="obs",
|
||||
),
|
||||
# The final turn's reasoning IS the model's real final output (it
|
||||
# already contains the answer). The serializer now renders this turn
|
||||
# via _final_answer_block(turn.reasoning), not rollout.final_answer.
|
||||
UnifiedTurn(
|
||||
reasoning="The result is 42.\nFINAL_ANSWER: 42", tool_name=None
|
||||
),
|
||||
],
|
||||
final_answer="42",
|
||||
cost_usd=0.02,
|
||||
tokens=30,
|
||||
num_tool_calls=1,
|
||||
)
|
||||
rec = trajectory_to_record("t1", "Q?", tools, roll, reward=0.5, domain="math")
|
||||
roles = [m["role"] for m in rec["conversations"]]
|
||||
assert roles[0] == "system" and roles[1] == "user"
|
||||
assert "tool" in roles and roles[-1] == "assistant"
|
||||
# Tool call is emitted as a <tool_call> tag (what the parser reads back).
|
||||
assert any(
|
||||
"<tool_call>" in m["content"] and "qwen3_32b" in m["content"]
|
||||
for m in rec["conversations"]
|
||||
if m["role"] == "assistant"
|
||||
)
|
||||
assert "FINAL_ANSWER: 42" in rec["conversations"][-1]["content"]
|
||||
assert rec["reward"] == 0.5 and rec["domain"] == "math"
|
||||
|
||||
|
||||
def test_generate_sft_dataset_end_to_end(tmp_path):
|
||||
tools = orchestrator_catalog()
|
||||
tasks = [
|
||||
Task(
|
||||
task_id="movie-001",
|
||||
question="Cancel ticket A03 and refund the user.",
|
||||
answer="refunded $20.90",
|
||||
domain="entertainment",
|
||||
),
|
||||
Task(
|
||||
task_id="unsolvable",
|
||||
question="Cancel ticket A03 and refund the user.",
|
||||
answer="refunded $20.90",
|
||||
domain="entertainment",
|
||||
),
|
||||
]
|
||||
|
||||
def rollout_fn(task):
|
||||
# Solve the first task; the second answers wrongly and never routes.
|
||||
if task.task_id == "movie-001":
|
||||
return UnifiedRollout(
|
||||
turns=[
|
||||
UnifiedTurn("", "cancel", {"booking": "A03"}, "ok"),
|
||||
UnifiedTurn("", "refund", {"user": "8612"}, "ok"),
|
||||
UnifiedTurn("done", None),
|
||||
],
|
||||
# num_tool_calls must reflect the two routed calls: the structural
|
||||
# _target_is_clean gate drops a trajectory with num_tool_calls < 1.
|
||||
final_answer="refunded $20.90",
|
||||
cost_usd=0.03,
|
||||
num_tool_calls=2,
|
||||
)
|
||||
return UnifiedRollout(
|
||||
turns=[UnifiedTurn("x", "cancel", {}, "ok")],
|
||||
final_answer="nope",
|
||||
cost_usd=0.05,
|
||||
)
|
||||
|
||||
# Simple, meaningful verifier: the final answer must mention the refund.
|
||||
def verify_fn(task, rollout):
|
||||
return "refund" in (rollout.final_answer or "").lower()
|
||||
|
||||
out = tmp_path / "sft.jsonl"
|
||||
stats = generate_sft_dataset(
|
||||
str(out),
|
||||
tasks=tasks,
|
||||
tools=tools,
|
||||
rollout_fn=rollout_fn,
|
||||
verify_fn=verify_fn,
|
||||
samples_per_task=2,
|
||||
)
|
||||
assert stats["tasks_seen"] == 2
|
||||
assert stats["records_written"] == 1 # only the solvable task
|
||||
assert stats["tasks_dropped"] == 1
|
||||
lines = out.read_text().strip().splitlines()
|
||||
assert len(lines) == 1
|
||||
rec = json.loads(lines[0])
|
||||
assert rec["task_id"] == "movie-001"
|
||||
assert rec["domain"] == "entertainment"
|
||||
assert (tmp_path / "sft.jsonl.stats.json").exists()
|
||||
@@ -0,0 +1,111 @@
|
||||
"""Offline tests for the answer verifier.
|
||||
|
||||
Math and code paths run with no network. The Gemini judge is only reachable for
|
||||
non-math/code domains; we monkeypatch ``_gemini_judge`` (or rely on its no-key
|
||||
fallback) so nothing hits the network.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data import verify as V
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import Task
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.verify import (
|
||||
make_verifier,
|
||||
verify_answer,
|
||||
)
|
||||
|
||||
|
||||
def _math(answer: str) -> Task:
|
||||
return Task(task_id="m", question="q", answer=answer, domain="math")
|
||||
|
||||
|
||||
def test_math_boxed_vs_plain_true():
|
||||
assert verify_answer(_math("42"), "\\boxed{42}") is True
|
||||
|
||||
|
||||
def test_math_boxed_both_sides():
|
||||
assert verify_answer(_math("\\boxed{42}"), "The answer is \\boxed{42}") is True
|
||||
|
||||
|
||||
def test_math_wrong_number_false():
|
||||
assert verify_answer(_math("42"), "41") is False
|
||||
|
||||
|
||||
def test_math_numeric_float_match():
|
||||
assert verify_answer(_math("42"), "42.0") is True
|
||||
|
||||
|
||||
def test_math_symbolic_equivalence_no_longer_verified_locally():
|
||||
# The sympy symbolic-equality block was removed from _math_equal: sympy.simplify
|
||||
# / parse_latex could hang on pathological \boxed{} answers while holding the GIL
|
||||
# and deadlock the threaded rejection sampler. The math domain now uses
|
||||
# string+numeric+f1 only (no judge call either), so a purely SYMBOLIC
|
||||
# equivalence that isn't string/numeric-equal is intentionally NOT recognized.
|
||||
t = _math("(x+1)^2")
|
||||
assert verify_answer(t, "x^2 + 2*x + 1") is False
|
||||
|
||||
|
||||
def test_empty_prediction_false():
|
||||
assert verify_answer(_math("42"), "") is False
|
||||
|
||||
|
||||
def test_code_short_gold_substring():
|
||||
t = Task(task_id="c", question="print hi", answer="hello", domain="code")
|
||||
assert verify_answer(t, "the output is hello") is True
|
||||
|
||||
|
||||
def test_code_falls_back_to_judge(monkeypatch):
|
||||
# Long gold + no substring match -> hits judge; monkeypatch to PASS.
|
||||
monkeypatch.setattr(V, "_gemini_judge", lambda task, pred: True)
|
||||
t = Task(
|
||||
task_id="c",
|
||||
question="solve",
|
||||
answer="x" * 300, # long -> skips substring shortcut
|
||||
domain="code",
|
||||
)
|
||||
assert verify_answer(t, "totally different") is True
|
||||
|
||||
|
||||
def test_judge_domain_uses_gemini(monkeypatch):
|
||||
calls = {}
|
||||
|
||||
def fake_judge(task, pred):
|
||||
calls["hit"] = (task.domain, pred)
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(V, "_gemini_judge", fake_judge)
|
||||
t = Task(task_id="s", question="q", answer="Oxygen", domain="science")
|
||||
assert verify_answer(t, "plants release oxygen") is True
|
||||
assert calls["hit"][0] == "science"
|
||||
|
||||
|
||||
def test_judge_domain_fallback_f1_when_no_key(monkeypatch):
|
||||
# Judge unavailable (returns None) -> falls back to string / F1.
|
||||
monkeypatch.setattr(V, "_gemini_judge", lambda task, pred: None)
|
||||
t = Task(task_id="s", question="q", answer="the capital is paris", domain="misc")
|
||||
assert verify_answer(t, "I think the capital is paris indeed") is True
|
||||
assert verify_answer(t, "completely unrelated text here") is False
|
||||
|
||||
|
||||
def test_make_verifier_reads_final_answer(monkeypatch):
|
||||
monkeypatch.setattr(V, "_gemini_judge", lambda task, pred: None)
|
||||
verify = make_verifier()
|
||||
|
||||
class Rollout:
|
||||
final_answer = "\\boxed{42}"
|
||||
|
||||
assert verify(_math("42"), Rollout()) is True
|
||||
|
||||
class BadRollout:
|
||||
final_answer = "999"
|
||||
|
||||
assert verify(_math("42"), BadRollout()) is False
|
||||
|
||||
|
||||
def test_make_verifier_missing_final_answer():
|
||||
verify = make_verifier()
|
||||
|
||||
class Empty:
|
||||
pass
|
||||
|
||||
assert verify(_math("42"), Empty()) is False
|
||||
@@ -0,0 +1,156 @@
|
||||
"""Offline smoke test for the v1 orchestrator-SFT driver (base self-sampling).
|
||||
|
||||
No network / no GPU: we feed fake ``datasets.Task`` objects, a canned
|
||||
``UnifiedRollout`` (via a fake ``rollout_fn``), and an accept-all verifier, then
|
||||
assert ``generate_sft_dataset`` writes a JSONL record in the expected
|
||||
``conversations`` shape that the SFT trainer consumes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog
|
||||
from openjarvis.agents.hybrid.toolorchestra.rollout import UnifiedRollout, UnifiedTurn
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import Task
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
|
||||
generate_sft_dataset,
|
||||
)
|
||||
|
||||
|
||||
def _fake_tasks(**_kwargs) -> list[Task]:
|
||||
# Accepts **kwargs so it can stand in for ``load_sft_tasks(cap=..., balanced=...)``.
|
||||
return [
|
||||
Task(task_id="t-math-1", question="What is 6 * 7?", answer="42", domain="math"),
|
||||
Task(task_id="t-code-1", question="Reverse 'ab'.", answer="ba", domain="code"),
|
||||
]
|
||||
|
||||
|
||||
def _canned_rollout(task: Task) -> UnifiedRollout:
|
||||
# One tool turn + a final-answer turn that echoes the gold answer.
|
||||
return UnifiedRollout(
|
||||
turns=[
|
||||
UnifiedTurn(
|
||||
reasoning="let me compute",
|
||||
tool_name="code_interpreter",
|
||||
arguments={"code": "print(6*7)"},
|
||||
observation="42",
|
||||
),
|
||||
UnifiedTurn(
|
||||
reasoning=(
|
||||
f"The tool returned {task.answer}. FINAL_ANSWER: {task.answer}"
|
||||
),
|
||||
tool_name=None,
|
||||
),
|
||||
],
|
||||
final_answer=task.answer,
|
||||
cost_usd=0.01,
|
||||
tokens=20,
|
||||
num_tool_calls=1,
|
||||
)
|
||||
|
||||
|
||||
def test_build_v1_writes_expected_conversations_shape(tmp_path, monkeypatch):
|
||||
tools = orchestrator_catalog() # specialists unwired (math/coder endpoints None)
|
||||
|
||||
# Accept-all verifier (the real make_verifier may hit Gemini / math checkers).
|
||||
monkeypatch.setattr(
|
||||
"openjarvis.learning.intelligence.orchestrator.sft_data.verify.make_verifier",
|
||||
lambda: lambda task, rollout: True,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.verify import (
|
||||
make_verifier,
|
||||
)
|
||||
|
||||
out = tmp_path / "orchestrator_sft_v1.jsonl"
|
||||
stats = generate_sft_dataset(
|
||||
str(out),
|
||||
tasks=_fake_tasks(),
|
||||
tools=tools,
|
||||
rollout_fn=_canned_rollout,
|
||||
verify_fn=make_verifier(),
|
||||
samples_per_task=2,
|
||||
max_keep_per_task=1,
|
||||
reward_fn=lambda r: -r.cost_usd,
|
||||
)
|
||||
|
||||
assert stats["tasks_seen"] == 2
|
||||
assert stats["records_written"] == 2
|
||||
assert stats["tasks_dropped"] == 0
|
||||
|
||||
lines = out.read_text().strip().splitlines()
|
||||
assert len(lines) == 2
|
||||
|
||||
rec = json.loads(lines[0])
|
||||
assert rec["task_id"] == "t-math-1"
|
||||
assert rec["domain"] == "math"
|
||||
|
||||
roles = [m["role"] for m in rec["conversations"]]
|
||||
assert roles[0] == "system" and roles[1] == "user"
|
||||
assert "tool" in roles
|
||||
assert roles[-1] == "assistant"
|
||||
# Tool call emitted as a <tool_call> tag the parser reads back.
|
||||
assert any(
|
||||
"<tool_call>" in m["content"] and "code_interpreter" in m["content"]
|
||||
for m in rec["conversations"]
|
||||
if m["role"] == "assistant"
|
||||
)
|
||||
assert "FINAL_ANSWER: 42" in rec["conversations"][-1]["content"]
|
||||
assert (tmp_path / "orchestrator_sft_v1.jsonl.stats.json").exists()
|
||||
|
||||
|
||||
def _load_driver():
|
||||
"""Load the (non-package) CLI driver module by path."""
|
||||
import importlib.util
|
||||
from pathlib import Path
|
||||
|
||||
root = Path(__file__).resolve().parents[2]
|
||||
path = root / "scripts" / "orchestrator" / "build_orchestrator_sft.py"
|
||||
spec = importlib.util.spec_from_file_location("build_orchestrator_sft", path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def test_driver_runs_end_to_end_with_fakes(tmp_path, monkeypatch):
|
||||
"""Drive ``main`` with no network: fake task loader, rollout, and verifier."""
|
||||
drv = _load_driver()
|
||||
|
||||
monkeypatch.setattr(drv, "load_sft_tasks", _fake_tasks)
|
||||
monkeypatch.setattr(
|
||||
drv,
|
||||
"run_unified_rollout",
|
||||
lambda question, tools, **kw: _canned_rollout(_fake_tasks()[0]),
|
||||
)
|
||||
monkeypatch.setattr(drv, "make_verifier", lambda: lambda task, rollout: True)
|
||||
# make_call_orchestrator builds an OpenAI client lazily, so it never connects
|
||||
# here (run_unified_rollout is stubbed out).
|
||||
|
||||
# OJ_DATA_ROOT is an ABSOLUTE path in the real .env (the data tree lives
|
||||
# outside the repo). If it leaks into the test env, chdir(tmp_path) below no
|
||||
# longer sandboxes anything and the run writes into the real dataset dir.
|
||||
# Drop it so the driver falls back to its repo-relative default.
|
||||
monkeypatch.delenv("OJ_DATA_ROOT", raising=False)
|
||||
monkeypatch.chdir(tmp_path)
|
||||
rc = drv.main(
|
||||
[
|
||||
"--out",
|
||||
"v1",
|
||||
"--samples-per-task",
|
||||
"1",
|
||||
"--max-tasks",
|
||||
"2",
|
||||
]
|
||||
)
|
||||
assert rc == 0
|
||||
# raw dir is {label}-{month}-{day}-{year}-{hhmm}{am|pm}[-{tag}] — e.g.
|
||||
# orch-july-7-2026-0553pm-v1 (see sft_data.naming.raw_dir_name).
|
||||
produced = list(
|
||||
(tmp_path / "data" / "orchestrator" / "raw").glob("*-v1/data.jsonl")
|
||||
)
|
||||
assert len(produced) == 1
|
||||
lines = produced[0].read_text().strip().splitlines()
|
||||
assert len(lines) == 2
|
||||
rec = json.loads(lines[0])
|
||||
assert rec["conversations"][0]["role"] == "system"
|
||||
assert "FINAL_ANSWER" in rec["conversations"][-1]["content"]
|
||||
Reference in New Issue
Block a user