diff --git a/scripts/orchestrator/build_orchestrator_sft.py b/scripts/orchestrator/build_orchestrator_sft.py new file mode 100644 index 00000000..22b0595e --- /dev/null +++ b/scripts/orchestrator/build_orchestrator_sft.py @@ -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 + # /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 /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 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()) diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/__init__.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/__init__.py new file mode 100644 index 00000000..4f1357ad --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/__init__.py @@ -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 +```` ``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", +] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/datasets.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/datasets.py new file mode 100644 index 00000000..8e2e6c7c --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/datasets.py @@ -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 ``…`` 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 … trace; keep the post-reasoning solution. + visible = re.sub(r"(?s).*?", "", full).strip() + # …and an UNCLOSED (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", " Okay, so...")` is always + # False and the task becomes unwinnable no matter what the model answers. + if "" in visible: + visible = visible.split("", 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. + + * ```` 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 "" 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", +] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/naming.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/naming.py new file mode 100644 index 00000000..88c7c749 --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/naming.py @@ -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()}" diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/reject_sample.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/reject_sample.py new file mode 100644 index 00000000..c52f9cbf --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/reject_sample.py @@ -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 ```` (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 with no opener is NOT truncation — Qwen3.x's chat + # template pre-fills the opening `` 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("") > fa.count(""): + return "unclosed (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 + # 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 "" in fa or "" 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>, , …) 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"] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/unified_serialize.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/unified_serialize.py new file mode 100644 index 00000000..163ac7bc --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/unified_serialize.py @@ -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 ``{...}`` 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 ```` 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 block matches the assistant's + # 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 ```` with NO opening ```` (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 "" in t and "" not in t: + t = "\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 + ```` marker so the block stays well-formed.""" + if not text: + return text + think = "" + body = text + m = re.match(r"(?is)^\s*()\s*", body) + if m: + 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 ```` / ```` +# and the sentinel set `` ``. +# 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"|" # 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: `` 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 , 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 "" in text: + # No explicit marker: the answer is whatever follows the final . + # Drop the think block (same rationale as above) — keep only the answer. + _, _, answer = text.rpartition("") + 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 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"] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/verify.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/verify.py new file mode 100644 index 00000000..2db4f09f --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/verify.py @@ -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 "" 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: ", so the raw +# rollout text is "\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", +] diff --git a/src/openjarvis/tools/calculator.py b/src/openjarvis/tools/calculator.py index 47ce773f..3e391706 100644 --- a/src/openjarvis/tools/calculator.py +++ b/src/openjarvis/tools/calculator.py @@ -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, } diff --git a/src/openjarvis/tools/web_search.py b/src/openjarvis/tools/web_search.py index 3d71f5de..950cc132 100644 --- a/src/openjarvis/tools/web_search.py +++ b/src/openjarvis/tools/web_search.py @@ -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") diff --git a/tests/test_orchestrator_learning/sft_data/__init__.py b/tests/test_orchestrator_learning/sft_data/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/test_orchestrator_learning/sft_data/test_datasets.py b/tests/test_orchestrator_learning/sft_data/test_datasets.py new file mode 100644 index 00000000..c07cf73e --- /dev/null +++ b/tests/test_orchestrator_learning/sft_data/test_datasets.py @@ -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": "reason Final: print('hi')"}, + ], + }, + { + "domain": "math", + "source": "Numina/NuminaMath", + "difficulty": 5, + "conversations": [ + {"from": "human", "value": "What is 6 times 7?"}, + { + "from": "gpt", + "value": "6*7=42 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"} + # 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?" diff --git a/tests/test_orchestrator_learning/sft_data/test_rejection_pipeline.py b/tests/test_orchestrator_learning/sft_data/test_rejection_pipeline.py new file mode 100644 index 00000000..a90bc2f1 --- /dev/null +++ b/tests/test_orchestrator_learning/sft_data/test_rejection_pipeline.py @@ -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 tag (what the parser reads back). + assert any( + "" 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() diff --git a/tests/test_orchestrator_learning/sft_data/test_verify.py b/tests/test_orchestrator_learning/sft_data/test_verify.py new file mode 100644 index 00000000..af2907ab --- /dev/null +++ b/tests/test_orchestrator_learning/sft_data/test_verify.py @@ -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 diff --git a/tests/test_orchestrator_learning/test_build_orchestrator_sft.py b/tests/test_orchestrator_learning/test_build_orchestrator_sft.py new file mode 100644 index 00000000..a8fdeebe --- /dev/null +++ b/tests/test_orchestrator_learning/test_build_orchestrator_sft.py @@ -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 tag the parser reads back. + assert any( + "" 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"]