From c7d788a089a47cc7729c3ddcdd2db981008fe220 Mon Sep 17 00:00:00 2001 From: pradipta-lyzr Date: Mon, 3 Aug 2026 23:13:08 +0530 Subject: [PATCH 01/14] traces: recover tool schemas; more_plus: route on every phrasing in a unit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three fixes the synthesizer needs underneath it, each standing on its own. traces._span_call hardcoded tools=None, so a Trajectory rebuilt from spans never carried the tool definitions the spans were already reporting — and to_dataset therefore never emitted a "tools" key. Read them from gen_ai.request.tools, or OpenInference's indexed llm.tools.{i}.tool.json_schema, wrapping bare function schemas. more_plus.split_units built a unit's BM25 surrogate from grp[0] alone, so when a unit holds several phrasings of one fact, rows 2..k contributed nothing to routing — an expert reachable only by the wording that happened to come first. Join the whole group. Neither backend passes a dataset's tool schemas to the training chat template, while chat() does pass them at inference. Whether trl can carry a tools column is version-dependent and unverified here, so say it out loud rather than let the skew pass silently — the convention this repo already holds for config a backend can't honor. Co-Authored-By: Claude Opus 5 (1M context) --- shadowlm/backends/base.py | 15 +++++++++++++++ shadowlm/backends/mlx.py | 4 +++- shadowlm/backends/torch.py | 4 +++- shadowlm/more_plus.py | 8 ++++++-- shadowlm/traces.py | 28 +++++++++++++++++++++++++++- tests/test_more_plus_config.py | 4 +++- tests/test_more_plus_router.py | 14 ++++++++++++++ tests/test_traces.py | 34 +++++++++++++++++++++++++++++++--- 8 files changed, 102 insertions(+), 9 deletions(-) diff --git a/shadowlm/backends/base.py b/shadowlm/backends/base.py index 0c362df..8d5a756 100644 --- a/shadowlm/backends/base.py +++ b/shadowlm/backends/base.py @@ -50,6 +50,21 @@ def stopped(self) -> bool: return bool(self._should_stop and self._should_stop()) +def warn_dropped_tools(dataset: Dataset, callbacks: Callbacks) -> None: + """Say so when a dataset's tool schemas won't reach the training template. + + Rows may carry `tools`, but neither backend passes them to the chat template + while training — inference (`chat(tools=...)`) does. Never let that skew pass + silently. + """ + if any(r.get("tools") for r in getattr(dataset, "rows", ())): + callbacks.log( + "[shadowlm] dataset carries tool schemas: training renders the tool " + "*calls* but not the tool *definitions*, which inference does pass — " + "the train and inference prompts differ on that block." + ) + + class FinetuneResult: """What a backend returns from `finetune`: where the adapter/model landed.""" diff --git a/shadowlm/backends/mlx.py b/shadowlm/backends/mlx.py index f56dcd7..2d3a4a3 100644 --- a/shadowlm/backends/mlx.py +++ b/shadowlm/backends/mlx.py @@ -16,7 +16,7 @@ from .._quiet import quiet_backend from ..data import CHAT, INSTRUCTION, SHAREGPT, TEXT, Dataset from ..training import ATTENTION_MODULES, MLP_MODULES, Metric, TrainConfig, resolve_total_steps -from .base import Backend, Callbacks, FinetuneResult +from .base import Backend, Callbacks, FinetuneResult, warn_dropped_tools DEFAULT_LORA_LAYERS = 16 # how many transformer blocks get LoRA adapters @@ -215,6 +215,8 @@ def finetune(self, dataset: Dataset, config: TrainConfig, callbacks: Callbacks, if self.accelerator != "none": callbacks.log(shadow.note) + warn_dropped_tools(dataset, callbacks) + # The method spec drives everything below: base requirements, trainable # surface, and data rendering. Backends never branch on the method name. spec = methods.get(config.method) diff --git a/shadowlm/backends/torch.py b/shadowlm/backends/torch.py index d8f9f42..9142ecd 100644 --- a/shadowlm/backends/torch.py +++ b/shadowlm/backends/torch.py @@ -20,7 +20,7 @@ from .._quiet import quiet_backend from ..data import Dataset from ..training import Metric, TrainConfig, resolve_total_steps -from .base import Backend, Callbacks, FinetuneResult +from .base import Backend, Callbacks, FinetuneResult, warn_dropped_tools _REQUIRED = ("torch", "transformers", "trl", "peft", "datasets") @@ -188,6 +188,8 @@ def finetune(self, dataset: Dataset, config: TrainConfig, callbacks: Callbacks, if shadow.fused_kernels: self._apply_liger(callbacks) + warn_dropped_tools(dataset, callbacks) + # The method spec drives base requirements and the trainable surface — # no branching on method names here. spec = methods.get(config.method) diff --git a/shadowlm/more_plus.py b/shadowlm/more_plus.py index 04f3ec2..3afafdf 100644 --- a/shadowlm/more_plus.py +++ b/shadowlm/more_plus.py @@ -236,13 +236,17 @@ def _surrogate(row: dict) -> str: def split_units(dataset, group_size: int = 1) -> list[tuple[str, list[dict]]]: - """Partition rows into (surrogate_text, rows) units — one expert per unit.""" + """Partition rows into (surrogate_text, rows) units — one expert per unit. + + The surrogate joins every row's query side, so a unit holding several + phrasings of one fact is routable by all of them, not just the first. + """ rows = list(getattr(dataset, "rows", dataset)) g = max(1, int(group_size)) units = [] for i in range(0, len(rows), g): grp = rows[i:i + g] - units.append((_surrogate(grp[0]), grp)) + units.append((" ".join(s for s in map(_surrogate, grp) if s), grp)) return units diff --git a/shadowlm/traces.py b/shadowlm/traces.py index 0e5ec38..b72e161 100644 --- a/shadowlm/traces.py +++ b/shadowlm/traces.py @@ -46,6 +46,7 @@ _CONVERSATION = "gen_ai.conversation.id" _REQ_MODEL = "gen_ai.request.model" _RESP_MODEL = "gen_ai.response.model" +_TOOLS = "gen_ai.request.tools" # Indexed form (older OpenLLMetry/Traceloop convention) — kept as a fallback. _PROMPT = "gen_ai.prompt" # gen_ai.prompt.{i}.{role,content,tool_calls...} _COMPLETION = "gen_ai.completion" # gen_ai.completion.{i}.{role,content,tool_calls...} @@ -136,6 +137,28 @@ def _tool_calls(attrs: dict, base: str) -> list[dict]: return calls +def _span_tools(attrs: dict) -> list[dict] | None: + """Tool definitions declared on the call, in OpenAI schema shape. + + From `gen_ai.request.tools` (a JSON list — what a spec-compliant instrumentor + and our own synth emitter write) or OpenInference's indexed + `llm.tools.{i}.tool.json_schema`. Bare function schemas are wrapped. + """ + raw = _as_list(attrs.get(_TOOLS)) or [ + attrs.get(f"llm.tools.{i}.tool.json_schema") for i in _indices(attrs, "llm.tools") + ] + tools = [] + for t in raw: + if isinstance(t, str): + try: + t = json.loads(t) + except json.JSONDecodeError: + continue + if isinstance(t, dict): + tools.append(t if "function" in t else {"type": "function", "function": t}) + return tools or None + + def _message(attrs: dict, p: str) -> dict | None: """One message from the attributes under index prefix `p`.""" role = attrs.get(f"{p}.role") @@ -361,7 +384,8 @@ def _span_call(span: dict) -> _Call | None: ts = float(ts) except (TypeError, ValueError): ts = 0.0 - return _Call(trace, ts, prompt, completion[-1] if completion else None, None, model) + return _Call(trace, ts, prompt, completion[-1] if completion else None, + _span_tools(attrs), model) def _system_text(val: Any) -> list[dict]: @@ -433,11 +457,13 @@ def from_spans( out: list[Trajectory] = [] for trace, calls in by_trace.items(): model = next((c.model for c in calls if c.model), None) + tools = next((c.tools for c in calls if c.tools), None) for convo in _reconstruct(calls, per_request=(builder == "per_request")): if not convo: continue out.append(Trajectory( messages=convo, + tools=tools, reward=rewards.get(trace, 0.0), metadata={"trace_id": trace, "model": model, "source": "otel"}, )) diff --git a/tests/test_more_plus_config.py b/tests/test_more_plus_config.py index e74f05d..b30885d 100644 --- a/tests/test_more_plus_config.py +++ b/tests/test_more_plus_config.py @@ -32,7 +32,9 @@ def test_split_units_cardinality_and_surrogate(): assert len(mp.split_units(ds, 1)) == 6 assert len(mp.split_units(ds, 3)) == 2 surrogate, rows = mp.split_units(ds, 3)[0] - assert surrogate == "q0" and len(rows) == 3 + # every row in the group joins the surrogate — a unit holding several + # phrasings of one fact must be routable by all of them, not just the first + assert surrogate == "q0 q1 q2" and len(rows) == 3 def test_split_units_chat_surrogate(): diff --git a/tests/test_more_plus_router.py b/tests/test_more_plus_router.py index 39cacd2..249d961 100644 --- a/tests/test_more_plus_router.py +++ b/tests/test_more_plus_router.py @@ -13,6 +13,20 @@ def _router(): ]) +def test_grouped_unit_routes_on_every_phrasing(): + """A unit of paraphrases must be reachable by any of them — the surrogate + joins the whole group, so phrasing #3 routes as well as phrasing #1.""" + rows = [{"question": q, "answer": "$0.08"} for q in ( + "What does Lyzr Cloud cost per agent run?", + "How much am I billed each time an agent executes?", + "Pricing for a single agent invocation?", + )] + [{"question": "Where is Lyzr headquartered?", "answer": "Boston"}] + surrogates = [s for s, _ in mp.split_units(rows, group_size=3)] + r = mp.BM25Router.build(surrogates) + assert r.rank("billed each time an agent executes", 1)[0][0] == 0 + assert r.rank("pricing single invocation", 1)[0][0] == 0 + + def test_tokenize_lowercases_and_splits(): assert mp._tokenize("Lyzr's $0.08 / agent-run!") == ["lyzr", "s", "0", "08", "agent", "run"] diff --git a/tests/test_traces.py b/tests/test_traces.py index 3dc5a03..7dc9c28 100644 --- a/tests/test_traces.py +++ b/tests/test_traces.py @@ -6,13 +6,21 @@ from shadowlm import traces -def _llm_span(trace_id, ts, prompt, completion, *, model="gpt-4o", reward=None): +_WEATHER_TOOL = {"type": "function", "function": { + "name": "get_weather", "description": "Current weather for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}}} + + +def _llm_span(trace_id, ts, prompt, completion, *, model="gpt-4o", reward=None, + tools=None): """A chat span in OTel GenAI indexed-attribute form.""" attrs = {"gen_ai.response.model": model} for i, m in enumerate(prompt): _put_message(attrs, f"gen_ai.prompt.{i}", m) for i, m in enumerate(completion): _put_message(attrs, f"gen_ai.completion.{i}", m) + if tools is not None: + attrs["gen_ai.request.tools"] = json.dumps(tools) if reward is not None: attrs["eval.score"] = reward return {"trace_id": trace_id, "start_time": ts, "name": "chat", "attributes": attrs} @@ -41,8 +49,9 @@ def _agent_run(trace_id): tool_res = {"role": "tool", "tool_call_id": "call_0", "content": "18C, sunny"} final = {"role": "assistant", "content": "It's 18°C and sunny in Paris."} return [ - _llm_span(trace_id, 1.0, [sys, user], [tool_call]), - _llm_span(trace_id, 2.0, [sys, user, tool_call, tool_res], [final], reward=1.0), + _llm_span(trace_id, 1.0, [sys, user], [tool_call], tools=[_WEATHER_TOOL]), + _llm_span(trace_id, 2.0, [sys, user, tool_call, tool_res], [final], + tools=[_WEATHER_TOOL], reward=1.0), ] @@ -83,6 +92,25 @@ def test_to_dataset_is_chat_format(): assert "messages" in ds.rows[0] +def test_tool_schemas_are_recovered_from_spans(): + traj = traces.from_spans(_agent_run("t1"))[0] + assert traj.tools == [_WEATHER_TOOL] + # and they ride through to the training rows, where the model needs them + assert traces.to_dataset([traj]).rows[0]["tools"] == [_WEATHER_TOOL] + + +def test_openinference_tool_schemas_and_bare_schemas_are_wrapped(): + bare = _WEATHER_TOOL["function"] + span = {"trace_id": "oi", "start_time": 1.0, "name": "chat", "attributes": { + "llm.input_messages.0.message.role": "user", + "llm.input_messages.0.message.content": "Weather in Paris?", + "llm.output_messages.0.message.role": "assistant", + "llm.output_messages.0.message.content": "18C.", + "llm.tools.0.tool.json_schema": json.dumps(bare), + }} + assert traces.from_spans([span])[0].tools == [_WEATHER_TOOL] + + def test_min_reward_filters(): spans = _agent_run("t1") + _agent_run("t2") # only the reward-bearing completion spans carry eval.score=1.0 → both pass; From 52af8457d38ed7e8dd0407f5266569b13aa8714f Mon Sep 17 00:00:00 2001 From: pradipta-lyzr Date: Mon, 3 Aug 2026 23:13:26 +0530 Subject: [PATCH 02/14] =?UTF-8?q?synth:=20the=20fourth=20inlet=20=E2=80=94?= =?UTF-8?q?=20training=20data=20for=20traffic=20that=20doesn't=20exist=20y?= =?UTF-8?q?et?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit capture() records a live agent, traces reads one that already ran, and Dataset.from_* loads what you already have. All three need the traffic to exist. This makes it: describe the task, point at a document, or amplify a few real episodes, and a teacher writes the rest. Two orthogonal axes again, mirroring backends × methods. Seeds decide where scenarios come from (task / document / episodes); modes decide what gets written per scenario (conversation / preference / paraphrases). Any seed composes with any mode. Everything converges on Trajectory — the same type capture and traces produce — and emit.py renders that into the shape the consumer takes, chosen from the method's *spec* and never its name. Generation is taxonomy first, instances second. A teacher asked for variety in one breath rewrites one example n times; a teacher asked to fill a named slot in a scenario tree doesn't. That structure, not prompt wording, is the anti-mode-collapse mechanism. Two consequences of the architecture worth naming: - trajectory-GRPO learns from the spread between good and bad attempts, so the judge scores but does not filter for that format — gating it would throw away exactly the signal it needs. - an outcome is accepted or rejected whole. MoRE+ groups rows by fixed size, so dropping one row of a paraphrase unit would misalign every unit after it; since a conversation outcome is one row, that is the same rule everywhere. Nothing is dropped silently: SynthReport reconciles exactly (generated == kept + invalid + duplicate + low-scoring + surplus) and report.balanced asserts it. Teachers are duck-typed .chat(), so a frontier model, a local slm.load() model and — later, for SDPO — the student itself are interchangeable. Co-Authored-By: Claude Opus 5 (1M context) --- shadowlm/__init__.py | 5 + shadowlm/synth/__init__.py | 284 ++++++++++++++++++++ shadowlm/synth/emit.py | 169 ++++++++++++ shadowlm/synth/generate.py | 417 +++++++++++++++++++++++++++++ shadowlm/synth/quality.py | 140 ++++++++++ shadowlm/synth/run.py | 106 ++++++++ shadowlm/synth/seeds.py | 117 ++++++++ shadowlm/synth/teacher.py | 131 +++++++++ tests/test_synth_emit.py | 90 +++++++ tests/test_synth_otlp_roundtrip.py | 99 +++++++ tests/test_synth_pipeline.py | 191 +++++++++++++ tests/test_synth_quality.py | 114 ++++++++ tests/test_synth_teacher.py | 135 ++++++++++ 13 files changed, 1998 insertions(+) create mode 100644 shadowlm/synth/__init__.py create mode 100644 shadowlm/synth/emit.py create mode 100644 shadowlm/synth/generate.py create mode 100644 shadowlm/synth/quality.py create mode 100644 shadowlm/synth/run.py create mode 100644 shadowlm/synth/seeds.py create mode 100644 shadowlm/synth/teacher.py create mode 100644 tests/test_synth_emit.py create mode 100644 tests/test_synth_otlp_roundtrip.py create mode 100644 tests/test_synth_pipeline.py create mode 100644 tests/test_synth_quality.py create mode 100644 tests/test_synth_teacher.py diff --git a/shadowlm/__init__.py b/shadowlm/__init__.py index eefb40a..0a7af6a 100644 --- a/shadowlm/__init__.py +++ b/shadowlm/__init__.py @@ -24,6 +24,8 @@ from .models import Model, Reply, load from .rl import Trajectory, TrajectoryGroup, judge_group from .training import Metric, TrainConfig, TrainingRun +from . import synth # noqa: E402 — imports models/apo, so it lands after them +from .synth import SynthRun, synthesize __version__ = "0.4.12" @@ -43,6 +45,9 @@ "load", "methods", "runs", + "synth", + "synthesize", + "SynthRun", "traces", "Metric", "TrainConfig", diff --git a/shadowlm/synth/__init__.py b/shadowlm/synth/__init__.py new file mode 100644 index 0000000..199a230 --- /dev/null +++ b/shadowlm/synth/__init__.py @@ -0,0 +1,284 @@ +"""Data synthesis — the fourth inlet. + +`capture()` records a live agent, `traces` reads one that already ran, and +`Dataset.from_*` loads what you already have. All three need the traffic to +exist. This makes training data that doesn't: describe the task in plain +English, point at a document, or hand over a few real episodes, and a teacher +model writes the rest — emitted in the exact shape the method you name accepts. + + import shadowlm as slm + + run = slm.synthesize( + task="Triage billing emails: classify urgency, draft a reply, and " + "escalate refunds over $200.", + teacher=slm.synth.frontier("gpt-4o"), + n=200, method="lora") + print(run.report.summary()) + model.finetune(run.dataset, method="lora") + +The teacher is anything that answers `.chat()` — a frontier model, a local +`slm.load(...)` model, or the student itself. Naming a `method=` picks the +output shape from that method's spec, so `method="dpo"` yields preference pairs +and `method="more_plus"` yields query-diverse paraphrase units, with no +format bookkeeping on your side. +""" + +from __future__ import annotations + +import math +import random +import time + +from .. import methods +from ..apo import _judge_one +from . import emit +from .generate import (FLAWS, STYLES, conversation, paraphrases, plan_leaves, + preference) +from .quality import Dedup +from .run import SynthReport, SynthRun +from .seeds import Seed, chunk_text, resolve_seed +from .teacher import OpenAIChatTeacher, as_teacher, CountingTeacher, frontier + +__all__ = [ + "synthesize", "SynthRun", "SynthReport", "Seed", "FORMATS", "resolve_output", + "frontier", "as_teacher", "OpenAIChatTeacher", "chunk_text", "emit", +] + +FORMATS = ("chat", "text", "preference", "grpo", "groups", "otlp") +_MORE_ADAPTERS = (methods.ADAPTER_MORE, methods.ADAPTER_MORE_PLUS) +_MAX_ROUNDS = 3 # a round that keeps nothing ends the run; this caps the rest + + +def synthesize( + *, + teacher, + task: str | None = None, + document=None, + episodes=None, + n: int = 100, + method: str | None = None, + format: str | None = None, + tools: list[dict] | None = None, + student=None, + judge=None, + min_score: float | None = 0.6, + dedup_threshold: float = 0.7, + per_scenario: int = 4, + seed: int = 3407, + verbose: bool = True, + on_progress=None, +) -> SynthRun: + """Generate `n` training rows and return a `SynthRun`. + + teacher: who writes the data — a loaded Model, or `frontier("gpt-4o")`. + task / document / episodes: the seed. Exactly one is required (`task` may + also accompany the other two as extra context). + method: the training method the data is for — its spec picks the output + shape. Override with `format=` ("chat", "text", "preference", "grpo", + "groups", "otlp"); "otlp" emits OpenTelemetry GenAI spans. + tools: OpenAI-style tool schemas, to synthesize tool-calling episodes. + student: a Model whose answers become the `rejected` side of preference + pairs — DPO then targets exactly the teacher/student gap. + judge / min_score: quality gate; the judge defaults to the teacher, and for + document seeds it scores answers against the source passage, so the same + call doubles as a grounding check. `min_score=None` disables it. + dedup_threshold: reject a row whose question overlaps an accepted one by at + least this much (token Jaccard). + per_scenario: rows generated per scenario — attempts per group for "groups", + and paraphrases per fact for the MoRE methods. + """ + if per_scenario < 1: + raise ValueError("per_scenario must be at least 1") + started = time.time() + fmt, mode = resolve_output(method, format) + source = resolve_seed(task=task, document=document, episodes=episodes) + teacher = CountingTeacher(as_teacher(teacher)) + scorer = CountingTeacher(as_teacher(judge)) if judge is not None else teacher + student = as_teacher(student) if student is not None else None + rng = random.Random(seed) + report = SynthReport(requested=n) + dedup = Dedup(dedup_threshold) + if source.exemplars: + # real episodes are for evaluating, not for handing back as "synthetic" + dedup.seed(t.first_user_content() for t in source.exemplars) + + # Trajectory-GRPO learns from the spread between good and bad attempts, so + # filtering the weak ones out would throw away exactly the signal it needs — + # score every row, gate on nothing. + gate = None if fmt == "groups" else min_score + kept: list = [] + rejected: list = [] + covered: list[str] = [] + if verbose: + print(f"[synth] {source.kind} seed · target {n} rows · format {fmt} · " + f"teacher {teacher.name}", flush=True) + + for round_no in range(1, _MAX_ROUNDS + 1): + if len(kept) >= n: + break + leaves = plan_leaves(source, teacher, rng=rng, avoid=covered, + count=math.ceil((n - len(kept)) / per_scenario)) + if not leaves: + break + report.scenarios += len(leaves) + covered.extend(leaf.scenario for leaf in leaves) + outcomes = _generate(leaves, source, teacher, mode=mode, tools=tools, + student=student, per_scenario=per_scenario) + _score(outcomes, scorer, enabled=min_score is not None) + before = len(kept) + for outcome in outcomes: + _absorb(outcome, report=report, dedup=dedup, min_score=gate, + kept=kept, rejected=rejected) + if verbose: + print(f"[synth] round {round_no} · {len(kept)}/{n} rows kept", + flush=True) + if on_progress: + on_progress(len(kept), n) + if len(kept) == before: + break # a whole round survived nothing — stop spending teacher calls + + if mode != "paraphrases" and len(kept) > n: + # a parallel round finishes every job it started, so it can overshoot; + # paraphrase units are left whole because MoRE+ groups by fixed size + report.surplus = len(kept) - n + kept = kept[:n] + if not kept: + gated = (f", {report.rejected_judge} below min_score={min_score}" + if gate is not None else "") + raise RuntimeError( + f"synthesis produced nothing usable — {report.rejected_validation} " + f"invalid, {report.rejected_dedup} duplicate{gated}. Loosen the gate " + "(min_score=, dedup_threshold=) or check what the teacher is emitting.") + + report.kept = len(kept) + scored = [t.reward for t in kept if t.reward] + report.mean_score = sum(scored) / len(scored) if scored else None + report.teacher_calls = teacher.calls + (0 if scorer is teacher else scorer.calls) + report.duration_s = time.time() - started + if mode == "paraphrases": + report.note = f"train with more_plus_group_size={per_scenario}" + + run = _emit(fmt, kept, report, rejected, seed=seed) + if verbose: + print(report.summary(), flush=True) + return run + + +def resolve_output(method: str | None, fmt: str | None) -> tuple[str, str]: + """(output format, generation mode) for the method you plan to train with. + + Dispatches on the method's *spec*, never its name — so a method registered + tomorrow that reuses an existing trainer gets the right data shape for free. + """ + if fmt is not None and fmt not in FORMATS: + raise ValueError( + f"unknown format {fmt!r} (expected one of {', '.join(FORMATS)})") + mode = "conversation" + if method is not None: + spec = methods.get(method) # raises, listing the registered methods + if spec.trainer == "dpo": + fmt, mode = fmt or "preference", "preference" + elif spec.adapter in _MORE_ADAPTERS: + fmt, mode = fmt or "chat", "paraphrases" + elif spec.trainer == "grpo": + fmt = fmt or "groups" + elif spec.raw_text: + fmt = fmt or "text" + if fmt == "preference": + mode = "preference" + return fmt or "chat", mode + + +def _generate(leaves, source, teacher, *, mode, tools, student, per_scenario): + """One job per row wanted, run at the teacher's parallelism.""" + jobs = [] + for i, leaf in enumerate(leaves): + if mode == "paraphrases": + jobs.append(lambda leaf=leaf, i=i: paraphrases( + leaf, source, teacher, k=per_scenario, + style=STYLES[i % len(STYLES)])) + continue + for j in range(per_scenario): + index = i * per_scenario + j + style = STYLES[index % len(STYLES)] + if mode == "preference": + jobs.append(lambda leaf=leaf, style=style, index=index: preference( + leaf, source, teacher, student=student, style=style, + flaw=FLAWS[index % len(FLAWS)])) + else: + jobs.append(lambda leaf=leaf, style=style: conversation( + leaf, source, teacher, tools=tools, style=style)) + return _run_jobs(jobs, workers=teacher.parallelism) + + +def _run_jobs(jobs, *, workers: int) -> list: + """Run jobs, preserving submission order — order groups MoRE+ units.""" + if workers <= 1 or len(jobs) <= 1: + return [job() for job in jobs] + from concurrent.futures import ThreadPoolExecutor # noqa: PLC0415 + + with ThreadPoolExecutor(max_workers=workers) as pool: + return list(pool.map(lambda job: job(), jobs)) + + +def _score(outcomes, judge, *, enabled: bool) -> None: + """Judge one row per outcome; a unit's rows share an answer, so one score + speaks for all of them (and a conversation outcome is a single row anyway).""" + if not enabled: + return + heads = [o.trajectories[0] for o in outcomes if o.trajectories] + _run_jobs([lambda t=t: _judge(t, judge) for t in heads], + workers=judge.parallelism) + for outcome in outcomes: + for traj in outcome.trajectories[1:]: + traj.reward = outcome.trajectories[0].reward + + +def _judge(traj, judge) -> None: + """Score an episode 0–1. When the row is grounded in a source passage that + passage is the reference, so this doubles as the hallucination check.""" + traj.reward = _judge_one(judge, traj.first_user_content(), + traj.final_content(), + traj.metadata.get("grounding") or "") + traj.metrics["judge_score"] = traj.reward + + +def _absorb(outcome, *, report, dedup, min_score, kept, rejected) -> None: + """Take or drop one outcome whole, counting it either way.""" + report.generated += outcome.attempted + report.rejected_validation += outcome.invalid + report.repaired += outcome.repaired + rows = outcome.trajectories + if not rows: + return + if not dedup.accept(_unit_text(rows), key=outcome.key): + report.rejected_dedup += len(rows) + _mark(rows, "duplicate", rejected) + elif min_score is not None and rows[0].reward < min_score: + report.rejected_judge += len(rows) + _mark(rows, f"judge score {rows[0].reward:.2f} < {min_score}", rejected) + else: + kept.extend(rows) + + +def _unit_text(rows) -> str: + return " ".join(m.get("content") or "" for t in rows for m in t.messages) + + +def _mark(rows, reason: str, rejected: list) -> None: + for traj in rows: + traj.metadata["reject_reason"] = reason + rejected.extend(rows) + + +def _emit(fmt: str, kept: list, report, rejected: list, *, seed: int) -> SynthRun: + run = SynthRun(format=fmt, report=report, trajectories=kept, rejected=rejected) + if fmt == "groups": + run.groups = emit.to_groups(kept) + elif fmt == "otlp": + run.spans = emit.to_otlp(kept, seed=seed) + else: + run.dataset = {"chat": emit.to_chat, "text": emit.to_text, + "preference": emit.to_preference, + "grpo": emit.to_grpo_prompts}[fmt](kept) + return run diff --git a/shadowlm/synth/emit.py b/shadowlm/synth/emit.py new file mode 100644 index 0000000..a8b4391 --- /dev/null +++ b/shadowlm/synth/emit.py @@ -0,0 +1,169 @@ +"""Emitters — one synthesized episode, rendered into whatever the consumer needs. + +The generator produces `Trajectory` objects and nothing else; these turn them +into the exact shape a training method (or somebody else's OTel pipeline) +accepts. "Inject ready" is not a claim we make about the JSON — it means the +code that eats the output accepts it, which is what the emit tests assert by +feeding every shape back into its real reader. +""" + +from __future__ import annotations + +import json +import random +from pathlib import Path + +from .. import traces +from ..data import Dataset +from ..rl import Trajectory, TrajectoryGroup + +SCORE_KEY = "shadowlm.synth.score" +_NAME = "synth" +# Fixed epoch: span order is all the reader uses, and a wall-clock read would +# make the same seed produce a different file every run. +_EPOCH_NS = 1_700_000_000_000_000_000 +_ONE_SECOND_NS = 1_000_000_000 + + +def to_chat(trajectories: list[Trajectory]) -> Dataset: + """Chat rows for the supervised methods (lora, qlora, dora, full, …).""" + return traces.to_dataset(trajectories, name=_NAME) + + +def to_text(trajectories: list[Trajectory]) -> Dataset: + """Raw domain text for continued pretraining. + + The assistant prose only: CPT trains on the material itself, not on a + transcript of somebody being asked about it. + """ + rows = [] + for traj in trajectories: + text = "\n\n".join(m["content"] for m in traj.messages + if m.get("role") == "assistant" and m.get("content")) + if text: + rows.append({"text": text}) + if not rows: + raise ValueError("no assistant prose to train on") + return Dataset.from_list(rows, name=_NAME, format="text") + + +def to_preference(trajectories: list[Trajectory]) -> Dataset: + """`{prompt, chosen, rejected}` rows — the shape trl's DPOTrainer demands. + + The rejected side rides on the trajectory's metadata, so a preference sample + is still a valid chat episode (its chosen answer) and can be emitted as one. + """ + rows = [] + for traj in trajectories: + prompt, chosen = traj.first_user_content(), traj.final_content() + rejected = traj.metadata.get("rejected") + if prompt and chosen and rejected and chosen != rejected: + rows.append({"prompt": prompt, "chosen": chosen, "rejected": rejected}) + if not rows: + raise ValueError( + "no usable preference pairs — every candidate was missing a side or " + "scored its rejected answer as good as the chosen one") + return Dataset.from_list(rows, name=_NAME, format="preference") + + +def to_grpo_prompts(trajectories: list[Trajectory]) -> Dataset: + """`{prompt, answer}` rows for reward-function GRPO (`reward_fns=[...]`).""" + rows = [{"prompt": t.first_user_content(), "answer": t.final_content()} + for t in trajectories if t.first_user_content()] + if not rows: + raise ValueError("no rows carried an opening user turn to prompt with") + return Dataset.from_list(rows, name=_NAME, format="instruction") + + +def to_groups(trajectories: list[Trajectory]) -> list[TrajectoryGroup]: + """Attempts at the same scenario, bucketed for trajectory-GRPO. + + A group needs two members and some spread in reward or it carries no signal; + `rl.weighted_rows` would skip those silently, so drop them here where the + report can say how many went. + """ + buckets: dict[str, list[Trajectory]] = {} + for traj in trajectories: + buckets.setdefault(traj.metadata.get("taxonomy_path", ""), []).append(traj) + groups = [TrajectoryGroup(ts) for ts in buckets.values() + if len(ts) >= 2 and len({t.reward for t in ts}) > 1] + if not groups: + raise ValueError( + "no scored groups — trajectory-GRPO needs several attempts per " + "scenario with differing judge scores (raise attempts=, or leave " + "min_score enabled so the judge actually runs)") + return groups + + +def to_otlp(trajectories: list[Trajectory], *, path: str | Path | None = None, + model: str = "shadowlm-synth", seed: int = 0) -> dict: + """Trajectories → an OTLP/JSON export of OpenTelemetry GenAI spans. + + One span per assistant turn, each span's input extending the previous span's + input plus output — the shape a real agent loop emits. That is what lets + `traces.from_otlp` fold them back into exactly these episodes, which is the + round-trip the OTLP test asserts. + """ + rng = random.Random(seed) + spans = [] + for traj in trajectories: + trace_id = f"{rng.getrandbits(128):032x}" + timestamp = _EPOCH_NS + for i, message in enumerate(traj.messages): + if message.get("role") != "assistant": + continue + attributes = { + "gen_ai.operation.name": "chat", + "gen_ai.request.model": model, + "gen_ai.conversation.id": trace_id, + "gen_ai.input.messages": json.dumps( + [_to_parts(m) for m in traj.messages[:i]]), + "gen_ai.output.messages": json.dumps([_to_parts(message)]), + } + if traj.tools: + attributes["gen_ai.request.tools"] = json.dumps(traj.tools) + spans.append({ + "traceId": trace_id, + "spanId": f"{rng.getrandbits(64):016x}", + "name": "chat", + "startTimeUnixNano": str(timestamp), + "attributes": [_kv(k, v) for k, v in attributes.items()] + + [_kv(SCORE_KEY, float(traj.reward))], + }) + timestamp += _ONE_SECOND_NS + payload = {"resourceSpans": [{ + "resource": {"attributes": [_kv("service.name", "shadowlm-synth")]}, + "scopeSpans": [{"scope": {"name": "shadowlm.synth"}, "spans": spans}], + }]} + if path is not None: + Path(path).write_text(json.dumps(payload, indent=2)) + return payload + + +def _to_parts(message: dict) -> dict: + """An OpenAI-wire message → the OTel GenAI `{role, parts}` form. + + The exact inverse of `traces._spec_message` — read the two together before + changing either, because their agreement *is* the round-trip guarantee. + """ + if message.get("role") == "tool": + return {"role": "tool", "parts": [{ + "type": "tool_call_response", + "id": message.get("tool_call_id"), + "response": message.get("content") or "", + }]} + parts = [] + if message.get("content"): + parts.append({"type": "text", "content": message["content"]}) + for call in message.get("tool_calls") or []: + fn = call.get("function") or {} + parts.append({"type": "tool_call", "id": call.get("id"), + "name": fn.get("name", ""), "arguments": fn.get("arguments")}) + return {"role": message.get("role", "assistant"), "parts": parts} + + +def _kv(key: str, value) -> dict: + """One OTLP attribute in `AnyValue` encoding.""" + encoded = ({"doubleValue": value} if isinstance(value, float) + else {"stringValue": str(value)}) + return {"key": key, "value": encoded} diff --git a/shadowlm/synth/generate.py b/shadowlm/synth/generate.py new file mode 100644 index 0000000..4b8477a --- /dev/null +++ b/shadowlm/synth/generate.py @@ -0,0 +1,417 @@ +"""Generation — scenarios first, then instances. + +The anti-mode-collapse mechanism here is structural, not a clever prompt: we +never ask a teacher for "500 examples". We ask it to lay out a *taxonomy* of +distinct scenarios, then generate a few instances per scenario with rotated user +styles. A teacher asked for variety in one breath will rewrite one example five +hundred times; a teacher asked to fill a named slot won't. + +Two axes compose here. Where the scenarios come from is the **seed** (a task +description, a document's facts, or real episodes to vary), and what gets +written per scenario is the **mode** (a conversation, a preference pair, or a +set of paraphrases for retrieval routing). Any seed works with any mode. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field + +from ..apo import _norm +from ..models import _first_json_object +from ..more_plus import _tokenize +from ..rl import Trajectory +from .quality import first_json_array, jaccard, validate + +# Asking for more than this in one array is where teachers start repeating +# themselves and truncating JSON. +_MAX_LEAVES_PER_CALL = 25 + +# Rotated so consecutive rows never read like the same person wrote them. +STYLES = ( + "terse, a single line", + "verbose, with background detail the model must look past", + "non-native English, slightly broken grammar", + "lowercase and typo-prone", + "frustrated — this is their second attempt at getting help", + "formal and businesslike", + "asks two things at once", + "vague, so the model must ask or state a reasonable assumption", + "pastes in a log line or error message", + "polite and chatty before getting to the point", + "keyword-style, barely a sentence", + "confident but wrong about one premise", +) + +# Rotated failure modes for the rejected side of a preference pair. +FLAWS = ( + "a confident factual error", + "it ignores part of the instruction", + "the right idea in the wrong format", + "vague and padded with filler", + "it answers a subtly different question", +) + + +@dataclass +class Leaf: + """One scenario to write training data for.""" + + scenario: str + difficulty: str = "medium" + angle: str = "" + grounding: str | None = None # source passage, for document seeds + exemplars: list[Trajectory] = field(default_factory=list) + + +@dataclass +class Outcome: + """What one generation attempt produced, and what it cost in rejections. + + An outcome is the unit of acceptance: its rows are kept or dropped together. + That matters for paraphrases, where MoRE+ groups rows by fixed size and + losing one row of a unit would misalign every unit after it. + """ + + trajectories: list[Trajectory] = field(default_factory=list) + key: str = "" # what "already have this" means for these rows + attempted: int = 1 # candidate rows asked for + invalid: int = 0 # candidates still broken after the corrective retry + repaired: int = 0 # candidates the retry rescued + + +# ---- scenarios --------------------------------------------------------------- +_TAXONOMY = """You are designing a training curriculum for a small model that must learn this task: + +{task} + +Produce a JSON array of exactly {n} training scenarios. Each element: +{{"scenario": "", + "difficulty": "easy" | "medium" | "hard", + "angle": ""}} + +Cover the full breadth of the task. No two scenarios may share an angle.{avoid} +Reply with JSON only.""" + +_FACTS = """List the distinct factual claims stated in this passage. + +PASSAGE: +\"\"\" +{chunk} +\"\"\" + +Reply with a JSON array of strings — one self-contained claim per element, and +only claims the passage actually makes. JSON only.""" + +_PATTERNS = """Below are real conversations from an agent. + +{examples} + +TASK CONTEXT: {task} + +Describe {n} DIFFERENT situations the same agent would plausibly face, in the +same domain and register. Reply with a JSON array of objects: +{{"scenario": "...", "difficulty": "easy" | "medium" | "hard", "angle": "..."}}{avoid} +JSON only.""" + + +def plan_leaves(seed, teacher, *, count: int, rng, avoid=()) -> list[Leaf]: + """Expand a seed into `count` distinct scenarios to generate against.""" + if seed.kind == "document": + return _document_leaves(seed, teacher, count=count) + if seed.kind == "episodes": + return _episode_leaves(seed, teacher, count=count, rng=rng, avoid=avoid) + return _batched_leaves( + teacher, lambda n, block: _TAXONOMY.format( + task=seed.context(), n=n, avoid=block), + count=count, avoid=avoid) + + +def _batched_leaves(teacher, prompt_for, *, count: int, avoid) -> list[Leaf]: + """Ask for scenarios in batches, steering each round clear of the last.""" + leaves: list[Leaf] = [] + seen = list(avoid) + while len(leaves) < count: + batch = min(_MAX_LEAVES_PER_CALL, count - len(leaves)) + fresh = _parse_leaves(teacher.chat( + [{"role": "user", "content": prompt_for(batch, _avoid_block(seen))}], + temperature=0.9, max_new_tokens=160 * batch + 300)) + if not fresh: + break # the teacher stopped producing usable JSON — report what we have + for leaf in fresh[:batch]: + leaves.append(leaf) + seen.append(leaf.scenario) + return leaves + + +def _document_leaves(seed, teacher, *, count: int) -> list[Leaf]: + """One leaf per fact stated in the document, carrying its source passage.""" + leaves: list[Leaf] = [] + for chunk in seed.chunks: + if len(leaves) >= count: + break + raw = teacher.chat([{"role": "user", "content": _FACTS.format(chunk=chunk)}], + temperature=0.2, max_new_tokens=1200) + for fact in first_json_array(str(raw)) or []: + if isinstance(fact, str) and fact.strip(): + leaves.append(Leaf(scenario=fact.strip(), grounding=chunk)) + return leaves[:count] + + +def _episode_leaves(seed, teacher, *, count: int, rng, avoid) -> list[Leaf]: + examples = _render_exemplars(seed.exemplars[:3]) + leaves = _batched_leaves( + teacher, lambda n, block: _PATTERNS.format( + examples=examples, task=seed.context(), n=n, avoid=block), + count=count, avoid=avoid) + for leaf in leaves: + leaf.exemplars = [rng.choice(seed.exemplars)] + return leaves + + +def _parse_leaves(raw) -> list[Leaf]: + leaves = [] + for item in first_json_array(str(raw)) or []: + if isinstance(item, str) and item.strip(): + leaves.append(Leaf(scenario=item.strip())) + elif isinstance(item, dict) and item.get("scenario"): + leaves.append(Leaf(scenario=str(item["scenario"]), + difficulty=str(item.get("difficulty", "medium")), + angle=str(item.get("angle", "")))) + return leaves + + +def _avoid_block(scenarios) -> str: + if not scenarios: + return "" + recent = "\n".join(f"- {s[:120]}" for s in scenarios[-30:]) + return f"\n\nAlready covered — do not repeat or rephrase these:\n{recent}\n" + + +# ---- instances --------------------------------------------------------------- +_TOOL_RULES = """- A tool call is {"role":"assistant","content":null,"tool_calls":[{"id":"call_1","type":"function","function":{"name":"...","arguments":""}}]}. +- Every tool call must be followed by {"role":"tool","tool_call_id":"call_1","content":""}. +- Call only the tools listed above.""" + +_CONVERSATION = """Write ONE realistic training conversation for this task. + +TASK: {task} +SCENARIO: {scenario} +DIFFICULTY: {difficulty}{angle} +USER STYLE: {style} +{tools}{grounding}{exemplars} +Output strict JSON only: +{{"messages": [{{"role": "system" | "user" | "assistant" | "tool", "content": "..."}}]}} + +Rules: +- The conversation MUST end with an assistant turn. +- Write specific, realistic content — never placeholders like [NAME] or [DATE]. +- The assistant answers the way the task demands, not as a generic chatbot. +{rules}""" + +_QUESTION = """TASK: {task} +SCENARIO: {scenario} +USER STYLE: {style} + +Write the single user message this scenario would produce. The message only — +no preamble, no quotes.""" + +_ANSWER = """{task} +{grounding} +Answer this as well as you possibly can. + +QUESTION: {question} + +Reply with the answer only.""" + +_FLAWED = """{task} + +Write a plausible but FLAWED answer to the question below. The flaw: {flaw}. +Never mention or label the flaw — it must read as a sincere attempt. + +QUESTION: {question} + +Reply with the answer only.""" + +_PARAPHRASES = """A user could ask about this fact in many different ways. Write {k} +questions that should all retrieve it, varying vocabulary, question form and +specificity — include one keyword-style query with no sentence structure. + +FACT: {fact} +{grounding} +Then write the single canonical answer, stated plainly and completely. + +Reply with JSON only: +{{"questions": ["...", "..."], "answer": "..."}}""" + + +def conversation(leaf: Leaf, seed, teacher, *, tools=None, style: str) -> Outcome: + """A full multi-turn episode for this scenario.""" + prompt = _CONVERSATION.format( + task=seed.context(), scenario=leaf.scenario, difficulty=leaf.difficulty, + angle=f"\nANGLE: {leaf.angle}" if leaf.angle else "", style=style, + tools=_tools_block(tools), grounding=_grounding_block(leaf), + exemplars=_exemplar_block(leaf), rules=_TOOL_RULES if tools else "") + outcome = Outcome() + messages = _ask_messages(teacher, prompt, tools=tools, outcome=outcome) + if messages: + traj = _trajectory(messages, leaf, teacher, tools=tools, style=style) + # instances of one scenario are meant to differ, so identity is the + # question asked — not the scenario they came from + outcome.key = traj.first_user_content() + outcome.trajectories.append(traj) + return outcome + + +def preference(leaf: Leaf, seed, teacher, *, student=None, style: str, + flaw: str) -> Outcome: + """A `chosen` / `rejected` pair for the same question. + + With a `student`, the rejected side is the student's own answer — so DPO + targets exactly the teacher/student gap, which is the whole point of + shadowing. Without one, the teacher writes a deliberately flawed answer. + """ + outcome = Outcome() + question = str(teacher.chat( + [{"role": "user", "content": _QUESTION.format( + task=seed.context(), scenario=leaf.scenario, style=style)}], + temperature=0.9, max_new_tokens=300)).strip() + if not question: + outcome.invalid += 1 + return outcome + chosen = _ask_answer(teacher, _ANSWER.format( + task=seed.context(), grounding=_grounding_block(leaf), question=question), + temperature=0.3) + rejected = ( + _ask_answer(student, _ANSWER.format( + task=seed.context(), grounding="", question=question), temperature=0.9) + if student is not None else + _ask_answer(teacher, _FLAWED.format( + task=seed.context(), flaw=flaw, question=question), temperature=0.9)) + if not chosen or not rejected or _norm(chosen) == _norm(rejected): + outcome.invalid += 1 + return outcome + traj = _trajectory( + [{"role": "user", "content": question}, + {"role": "assistant", "content": chosen}], + leaf, teacher, tools=None, style=style) + traj.metadata["rejected"] = rejected + traj.metadata["rejected_from"] = "student" if student is not None else "flaw" + outcome.key = question + outcome.trajectories.append(traj) + return outcome + + +def paraphrases(leaf: Leaf, seed, teacher, *, k: int, style: str) -> Outcome: + """`k` differently-worded questions sharing one canonical answer. + + MoRE and MoRE+ index the *user* turn — it is the retrieval key, and for + MoRE+ the BM25 routing surrogate. One fact phrased one way yields an expert + nobody can route to, so phrasing diversity is the objective here, not a + nicety. + """ + outcome = Outcome() + raw = teacher.chat([{"role": "user", "content": _PARAPHRASES.format( + k=k, fact=leaf.scenario, grounding=_grounding_block(leaf))}], + temperature=0.9, max_new_tokens=200 * k + 400) + parsed = _first_json_object(str(raw)) or {} + questions = [str(q).strip() for q in parsed.get("questions") or [] + if str(q).strip()] + answer = str(parsed.get("answer") or "").strip() + picked = _diverse(questions, k) if answer else [] + if len(picked) < k: + # a unit must be exactly k rows — see Outcome — so a fact that can't be + # phrased k genuinely different ways is dropped rather than shipped short + outcome.invalid += 1 + return outcome + for question in picked: + outcome.trajectories.append(_trajectory( + [{"role": "user", "content": question}, + {"role": "assistant", "content": answer}], + leaf, teacher, tools=None, style=style)) + outcome.key = leaf.scenario # the fact is the identity; phrasings are its rows + outcome.attempted = len(outcome.trajectories) + return outcome + + +def _diverse(questions: list[str], k: int, *, threshold: float = 0.6) -> list[str]: + """Drop phrasings that echo one already kept — a duplicate paraphrase is + routing dead weight, not extra coverage.""" + picked: list[str] = [] + seen: list[set[str]] = [] + for question in questions: + tokens = set(_tokenize(question)) + if any(jaccard(tokens, other) > threshold for other in seen): + continue + picked.append(question) + seen.append(tokens) + if len(picked) == k: + break + return picked + + +def _ask_messages(teacher, prompt: str, *, tools, outcome: Outcome) -> list[dict] | None: + """Ask for a conversation, validate it, allow exactly one corrective retry. + + One retry, not a loop: a teacher that can't emit valid JSON twice won't on + the fifth attempt, and the scenario is cheap to abandon. + """ + for attempt in (0, 1): + raw = str(teacher.chat([{"role": "user", "content": prompt}], + temperature=0.9, max_new_tokens=1600)) + parsed = _first_json_object(raw) + messages = parsed.get("messages") if isinstance(parsed, dict) else None + problems = validate(messages or [], tools=tools) + if not problems: + outcome.repaired += attempt + return messages + if attempt: + break + prompt = (f"{prompt}\n\nYour previous attempt was rejected: " + f"{'; '.join(problems)}.\nEmit corrected JSON only.") + outcome.invalid += 1 + return None + + +def _ask_answer(teacher, prompt: str, *, temperature: float) -> str: + return str(teacher.chat([{"role": "user", "content": prompt}], + temperature=temperature, max_new_tokens=800)).strip() + + +def _trajectory(messages, leaf: Leaf, teacher, *, tools, style: str) -> Trajectory: + """Wrap generated messages with the provenance every synthetic row carries.""" + return Trajectory(messages=messages, tools=tools, metadata={ + "source": "synth", "teacher": teacher.name, "style": style, + "taxonomy_path": leaf.scenario, "difficulty": leaf.difficulty, + "grounding": leaf.grounding, + }) + + +def _tools_block(tools) -> str: + if not tools: + return "" + return f"\nTOOLS THE ASSISTANT MAY CALL:\n{json.dumps(tools, indent=2)}\n" + + +def _grounding_block(leaf: Leaf) -> str: + if not leaf.grounding: + return "" + return ("\nSOURCE MATERIAL — every claim in the answer must be supported by " + f"this passage:\n\"\"\"\n{leaf.grounding}\n\"\"\"\n") + + +def _exemplar_block(leaf: Leaf) -> str: + if not leaf.exemplars: + return "" + return ("\nREAL EXAMPLES from this agent — match their voice and format, but " + f"do not reuse their content:\n{_render_exemplars(leaf.exemplars)}\n") + + +def _render_exemplars(trajectories) -> str: + blocks = [] + for i, traj in enumerate(trajectories, 1): + turns = "\n".join( + f"{m.get('role')}: {m.get('content') or json.dumps(m.get('tool_calls', ''))}" + for m in traj.messages) + blocks.append(f"### Example {i}\n{turns}") + return "\n\n".join(blocks) diff --git a/shadowlm/synth/quality.py b/shadowlm/synth/quality.py new file mode 100644 index 0000000..053bc5d --- /dev/null +++ b/shadowlm/synth/quality.py @@ -0,0 +1,140 @@ +"""Quality control — validate, deduplicate, judge. + +Synthetic data is cheap to make and easy to make badly, so nothing reaches a +dataset without passing through here. Every rejection is counted rather than +quietly dropped: the run report has to be able to tell you the truth about what +survived and why the rest didn't. +""" + +from __future__ import annotations + +import json +import re + +from ..apo import _norm +from ..more_plus import _tokenize + +_ROLES = frozenset({"system", "user", "assistant", "tool"}) + +# Two tells that a teacher phoned it in: an un-filled template slot, or the +# refusal boilerplate that would teach the student to refuse. +_JUNK = re.compile( + r"\[(?:NAME|COMPANY|DATE|TODO|INSERT|PRODUCT|X)\]|as an AI language model", + re.IGNORECASE) + + +def first_json_array(text: str) -> list | None: + """The first valid JSON array in `text` — teachers wrap JSON in prose.""" + decoder = json.JSONDecoder() + for i, ch in enumerate(text): + if ch == "[": + try: + obj, _ = decoder.raw_decode(text[i:]) + except json.JSONDecodeError: + continue + if isinstance(obj, list): + return obj + return None + + +def validate(messages: list[dict], *, tools: list[dict] | None = None) -> list[str]: + """Everything wrong with a generated conversation — empty list means usable. + + The last rule is the load-bearing one: the torch backend only takes the + prompt-masking path when every row ends on an assistant turn, so a row that + doesn't is worse than useless — it silently degrades the whole batch. + """ + if not messages: + return ["no messages"] + problems = [] + roles = [m.get("role") for m in messages] + unknown = set(roles) - _ROLES + if unknown: + problems.append(f"unknown roles {sorted(map(str, unknown))}") + if roles.count("system") > 1 or ("system" in roles[1:]): + problems.append("a system turn may appear only once, first") + last = messages[-1] + if last.get("role") != "assistant": + problems.append("the conversation must end with an assistant turn") + elif not (last.get("content") or last.get("tool_calls")): + problems.append("the final assistant turn is empty") + problems += _tool_problems(messages, tools) + if _JUNK.search(" ".join(str(m.get("content") or "") for m in messages)): + problems.append("contains a placeholder or assistant boilerplate") + return problems + + +def _tool_problems(messages: list[dict], tools: list[dict] | None) -> list[str]: + """Tool calls that don't parse, aren't declared, or are never answered.""" + declared = {(t.get("function") or {}).get("name") for t in (tools or [])} + problems: list[str] = [] + unanswered: list[str] = [] + for msg in messages: + for call in msg.get("tool_calls") or []: + fn = call.get("function") or {} + name = fn.get("name") + args = fn.get("arguments") + if isinstance(args, str): + try: + json.loads(args) + except json.JSONDecodeError: + problems.append(f"tool call {name!r} has unparseable arguments") + elif not isinstance(args, dict): + problems.append(f"tool call {name!r} has no arguments") + if declared and name not in declared: + problems.append(f"calls undeclared tool {name!r}") + unanswered.append(call.get("id")) + if msg.get("role") == "tool": + call_id = msg.get("tool_call_id") + if call_id in unanswered: + unanswered.remove(call_id) + else: + problems.append(f"tool result for unknown call {call_id!r}") + if unanswered: + problems.append(f"{len(unanswered)} tool call(s) never answered") + return problems + + +def jaccard(a: set, b: set) -> float: + """Token-set overlap, 0–1. Used for both dedup and paraphrase diversity.""" + if not a or not b: + return 0.0 + return len(a & b) / len(a | b) + + +class Dedup: + """Rejects repeats: exact on normalized text, near on query-side overlap. + + Mode collapse is *the* failure of LLM synthesis — a teacher asked for + variety will happily rewrite one example five hundred times — so near-dup + rejection is on by default. One pass over the accepted pool per candidate is + fine at the scale a run produces; swap in minhash if that ever changes. + """ + + def __init__(self, threshold: float = 0.7) -> None: + self.threshold = threshold + self._exact: set[str] = set() + self._tokens: list[set[str]] = [] + + def seed(self, texts) -> None: + """Pre-load texts to reject against — real episodes we must not clone.""" + for text in texts: + self._exact.add(_norm(text)) + self._tokens.append(set(_tokenize(text))) + + def accept(self, text: str, *, key: str | None = None) -> bool: + """True when `text` is new (and now remembered), False when it repeats. + + `key` is the query-side text near-duplicates are judged on; two rows with + different answers to the same question are still a duplicate question. + """ + exact = _norm(text) + if exact in self._exact: + return False + tokens = set(_tokenize(key if key is not None else text)) + if tokens and any(jaccard(tokens, seen) >= self.threshold + for seen in self._tokens): + return False + self._exact.add(exact) + self._tokens.append(tokens) + return True diff --git a/shadowlm/synth/run.py b/shadowlm/synth/run.py new file mode 100644 index 0000000..44185d8 --- /dev/null +++ b/shadowlm/synth/run.py @@ -0,0 +1,106 @@ +"""The handle `synthesize()` returns — what came out, and what didn't. + +Synthetic data is only trustworthy if you can see how much of it was thrown +away, so the report is not a summary line bolted on at the end: every rejection +is counted as it happens, and the counts reconcile exactly — + + generated == kept + rejected_validation + rejected_dedup + rejected_judge +""" + +from __future__ import annotations + +import json +from dataclasses import asdict, dataclass, field +from pathlib import Path + +from ..data import Dataset +from ..rl import Trajectory, TrajectoryGroup + + +@dataclass +class SynthReport: + """Honest accounting for one synthesis run.""" + + requested: int = 0 + kept: int = 0 + scenarios: int = 0 # distinct taxonomy leaves explored + generated: int = 0 # candidate rows the teacher was asked for + rejected_validation: int = 0 # malformed, even after the corrective retry + rejected_dedup: int = 0 # a repeat of something already kept + rejected_judge: int = 0 # scored below min_score + surplus: int = 0 # good rows past `requested`, trimmed + repaired: int = 0 # rescued by the corrective retry + mean_score: float | None = None + teacher_calls: int = 0 + duration_s: float = 0.0 + note: str | None = None # e.g. the more_plus_group_size to train with + + @property + def balanced(self) -> bool: + """The funnel reconciles — nothing vanished unaccounted for.""" + return self.generated == (self.kept + self.rejected_validation + + self.rejected_dedup + self.rejected_judge + + self.surplus) + + def summary(self) -> str: + score = f" · mean score {self.mean_score:.2f}" if self.mean_score else "" + surplus = f" · {self.surplus} surplus" if self.surplus else "" + note = f"\n note: {self.note}" if self.note else "" + return ( + f" ♥ {self.kept} rows from {self.generated} generated " + f"({self.scenarios} scenarios)\n" + f" rejected: {self.rejected_validation} invalid · " + f"{self.rejected_dedup} duplicate · {self.rejected_judge} low-scoring" + f"{surplus} · {self.repaired} repaired\n" + f" {self.teacher_calls} teacher calls · {self.duration_s:.1f}s{score}{note}" + ) + + def to_dict(self) -> dict: + return asdict(self) + + def __repr__(self) -> str: + return (f"SynthReport(kept={self.kept}/{self.requested}, " + f"generated={self.generated}, duration={self.duration_s:.1f}s)") + + +@dataclass +class SynthRun: + """The result of `synthesize()`: the data, plus how it came to be.""" + + format: str + report: SynthReport + trajectories: list[Trajectory] = field(default_factory=list) + rejected: list[Trajectory] = field(default_factory=list) + dataset: Dataset | None = None # chat | text | preference | grpo + groups: list[TrajectoryGroup] | None = None # trajectory-GRPO + spans: dict | None = None # OTLP payload + + def rows(self) -> list[dict]: + """The training rows this run produced, in the shape the method takes.""" + if self.dataset is not None: + return self.dataset.rows + if self.groups is not None: + from ..rl import weighted_rows # noqa: PLC0415 + + return weighted_rows(self.groups) + return [] + + def save(self, path: str | Path) -> str: + """Write the run's artifact — JSONL rows, or the OTLP payload as JSON.""" + path = Path(path) + if self.spans is not None: + path.write_text(json.dumps(self.spans, indent=2)) + else: + path.write_text("\n".join(json.dumps(r) for r in self.rows())) + return str(path) + + def to_otlp(self, path: str | Path | None = None, **kwargs) -> dict: + """Re-emit the episodes as OpenTelemetry GenAI spans.""" + from .emit import to_otlp # noqa: PLC0415 + + return to_otlp(self.trajectories, path=path, **kwargs) + + def __repr__(self) -> str: + what = (f"{len(self.groups)} groups" if self.groups is not None + else f"{len(self.rows())} rows") + return f"SynthRun(format={self.format!r}, {what}, {self.report!r})" diff --git a/shadowlm/synth/seeds.py b/shadowlm/synth/seeds.py new file mode 100644 index 0000000..90d34a0 --- /dev/null +++ b/shadowlm/synth/seeds.py @@ -0,0 +1,117 @@ +"""Seeds — what the user brings to the synthesizer. + +Three ways in, one shape out. Describe the task in plain English, point at a +document to ground on, or hand over a few real episodes to amplify; `resolve_seed` +normalizes all three into a `Seed` the generator reads. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from pathlib import Path + +from ..data import Dataset +from ..rl import Trajectory + +# Documents we can read as-is. Anything else is the user's to convert — silently +# extracting text from a PDF would be a whole dependency and a lossy guess. +_TEXT_SUFFIXES = (".txt", ".md", ".markdown", ".rst") +_MAX_PATHLIKE = 4096 # beyond this a string is content, not a filename + + +@dataclass +class Seed: + """Normalized synthesis input.""" + + kind: str # "task" | "document" | "episodes" + task: str | None = None + chunks: list[str] = field(default_factory=list) # document passages + exemplars: list[Trajectory] = field(default_factory=list) + + def context(self) -> str: + """The task description every generation prompt is written against.""" + if self.task: + return self.task + if self.kind == "document": + return ("Answer questions about the source material accurately and " + "concisely, using only what it states.") + return "Perform the task demonstrated by the example conversations." + + +def chunk_text(text: str, *, target_chars: int = 2000, + overlap_chars: int = 200) -> list[str]: + """Split prose into ~`target_chars` passages on paragraph boundaries. + + Each chunk carries the tail of the previous one so a fact spanning the seam + is still stated whole somewhere. + """ + paragraphs = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()] + chunks: list[str] = [] + current = "" + for para in paragraphs: + if current and len(current) + len(para) + 2 > target_chars: + chunks.append(current) + tail = current[-overlap_chars:] if overlap_chars else "" + current = f"{tail}\n\n{para}" if tail else para + else: + current = f"{current}\n\n{para}" if current else para + if current: + chunks.append(current) + return chunks + + +def resolve_seed(*, task=None, document=None, episodes=None) -> Seed: + """Normalize the caller's starting point into a `Seed`.""" + if document is not None: + chunks = chunk_text(_document_text(document)) + if not chunks: + raise ValueError("document= is empty — nothing to ground on") + return Seed(kind="document", task=task, chunks=chunks) + if episodes is not None: + exemplars = _as_trajectories(episodes) + if not exemplars: + raise ValueError( + "episodes= held no usable conversations (each needs a " + "'messages' list)") + return Seed(kind="episodes", task=task, exemplars=exemplars) + if task: + return Seed(kind="task", task=task) + raise ValueError( + "synthesize needs a seed: task='what the model should learn', " + "document='notes.md', or episodes=[...]") + + +def _document_text(document) -> str: + """The document's text — read from a path, or taken as the content itself.""" + if isinstance(document, (str, Path)) and len(str(document)) < _MAX_PATHLIKE: + path = Path(document) + if path.exists(): + if path.suffix.lower() not in _TEXT_SUFFIXES: + raise ValueError( + f"unsupported document type {path.suffix!r} — convert it to " + f".txt or .md first (supported: {', '.join(_TEXT_SUFFIXES)})") + return path.read_text() + return str(document) + + +def _as_trajectories(episodes) -> list[Trajectory]: + """Episodes from anywhere — trajectories, a Dataset, rows, or a file.""" + if isinstance(episodes, (str, Path)): + from .. import traces # noqa: PLC0415 — avoids an import cycle + + path = Path(episodes) + if path.suffix.lower() in (".jsonl", ".ndjson", ".csv", ".parquet"): + return _rows_to_trajectories(Dataset.load(path).as_chat().rows) + return traces.from_otlp(path) + if isinstance(episodes, Dataset): + return _rows_to_trajectories(episodes.as_chat().rows) + items = list(episodes) + if items and isinstance(items[0], Trajectory): + return items + return _rows_to_trajectories(items) + + +def _rows_to_trajectories(rows) -> list[Trajectory]: + return [Trajectory(messages=r["messages"], tools=r.get("tools")) + for r in rows if r.get("messages")] diff --git a/shadowlm/synth/teacher.py b/shadowlm/synth/teacher.py new file mode 100644 index 0000000..fa89cf4 --- /dev/null +++ b/shadowlm/synth/teacher.py @@ -0,0 +1,131 @@ +"""Teachers — whoever writes the synthetic data. + +A teacher is anything that answers `chat(messages) -> str`: a frontier model over +an OpenAI-compatible API (`frontier("gpt-4o")`), a model you loaded yourself with +`slm.load(...)`, or — for self-distillation — the student. The synthesizer only +ever calls `.chat()`, so the three are interchangeable and nothing downstream +knows which one produced a row. + +Pure stdlib, like the rest of the transport (see `remote.py`). +""" + +from __future__ import annotations + +import json +import os +import threading +import time +import urllib.error +import urllib.request + +_RETRIES = 3 +_RETRY_AFTER_S = 1.0 # doubled ×4 per attempt: 1s, 4s +_RETRY_CODES = frozenset({408, 429, 500, 502, 503, 504}) + + +class OpenAIChatTeacher: + """A frontier (or any OpenAI-compatible) model, over plain HTTP. + + Works against OpenAI, vLLM, Ollama, a ShadowLM capture proxy — anything + serving `/chat/completions`. Retries the transient failures that make long + synthesis runs fall over; everything else raises with the server's own words. + """ + + def __init__(self, model: str, *, base_url: str | None = None, + api_key: str | None = None, parallelism: int = 4, + timeout: float = 120.0) -> None: + self.name = self.model = model + self.base_url = (base_url or os.environ.get("OPENAI_BASE_URL") + or "https://api.openai.com/v1").rstrip("/") + self.api_key = api_key or os.environ.get("OPENAI_API_KEY") or "" + self.parallelism = max(1, parallelism) + self.timeout = timeout + + def chat(self, messages: list[dict], *, temperature: float = 0.7, + max_new_tokens: int = 1024, **_) -> str: + body = json.dumps({"model": self.model, "messages": messages, + "temperature": temperature, + "max_tokens": max_new_tokens}).encode() + headers = {"Content-Type": "application/json"} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + req = urllib.request.Request(f"{self.base_url}/chat/completions", + data=body, headers=headers, method="POST") + delay = _RETRY_AFTER_S + for attempt in range(_RETRIES): + final = attempt == _RETRIES - 1 + try: + with urllib.request.urlopen(req, timeout=self.timeout) as resp: + payload = json.loads(resp.read()) + return payload["choices"][0]["message"].get("content") or "" + except urllib.error.HTTPError as e: + if final or e.code not in _RETRY_CODES: + detail = e.read()[:200].decode("utf-8", "replace") + raise RuntimeError( + f"teacher {self.name!r} returned HTTP {e.code}: {detail}" + ) from None + except (urllib.error.URLError, TimeoutError) as e: + if final: + raise RuntimeError( + f"teacher {self.name!r} unreachable at {self.base_url}: {e}" + ) from None + time.sleep(delay) + delay *= 4 + + +class _ModelTeacher: + """A loaded shadowlm `Model` as a teacher. + + Serialized (`parallelism = 1`): neither backend's generate is thread-safe — + the same reason the capture proxy holds a generation lock. + """ + + parallelism = 1 + + def __init__(self, model) -> None: + self._model = model + self.name = getattr(model, "name", "local") + + def chat(self, messages: list[dict], *, temperature: float = 0.7, + max_new_tokens: int = 1024, **_) -> str: + return str(self._model.chat(messages, temperature=temperature, + max_new_tokens=max_new_tokens)) + + +class CountingTeacher: + """Wraps a teacher to count calls for the run report.""" + + def __init__(self, inner) -> None: + self._inner = inner + self.name, self.parallelism = inner.name, inner.parallelism + self.calls = 0 + self._lock = threading.Lock() + + def chat(self, messages, **kwargs) -> str: + with self._lock: + self.calls += 1 + return self._inner.chat(messages, **kwargs) + + +def frontier(model: str, **kwargs) -> OpenAIChatTeacher: + """A frontier teacher by name: `frontier("gpt-4o")`. + + Reads `OPENAI_API_KEY` / `OPENAI_BASE_URL` unless you pass `api_key=` / + `base_url=`. + """ + return OpenAIChatTeacher(model, **kwargs) + + +def as_teacher(obj): + """Coerce a `Model`, or anything already teacher-shaped, into a teacher.""" + if obj is None: + raise ValueError( + "synthesize needs teacher= — a loaded Model, or " + "slm.synth.frontier('gpt-4o') for an OpenAI-compatible endpoint") + if hasattr(obj, "chat") and hasattr(obj, "parallelism"): + return obj + if hasattr(obj, "chat"): + return _ModelTeacher(obj) + raise TypeError( + f"{type(obj).__name__} is not a teacher — expected a loaded Model or an " + "object with .chat(messages)") diff --git a/tests/test_synth_emit.py b/tests/test_synth_emit.py new file mode 100644 index 0000000..777478e --- /dev/null +++ b/tests/test_synth_emit.py @@ -0,0 +1,90 @@ +"""Emitter contracts — asserted by the code that actually consumes each shape. + +"Inject ready" is not a property of the JSON; it is the consuming reader +accepting it. So every test here hands the output to the real consumer: +`Dataset` format detection, `rl.weighted_rows`, `more_plus.split_units` + the +BM25 router. Nothing is checked by eyeballing keys alone. +""" + +import pytest + +from shadowlm import more_plus as mp +from shadowlm.rl import Trajectory, weighted_rows +from shadowlm.synth import emit + + +def _traj(question, answer, *, scenario="s", reward=0.0, rejected=None): + traj = Trajectory( + messages=[{"role": "user", "content": question}, + {"role": "assistant", "content": answer}], + reward=reward, metadata={"taxonomy_path": scenario}) + if rejected is not None: + traj.metadata["rejected"] = rejected + return traj + + +def test_chat_rows_end_on_an_assistant_turn(): + """The torch backend only takes the prompt-masking path when they all do.""" + ds = emit.to_chat([_traj("q1", "a1"), _traj("q2", "a2")]) + assert ds.format == "chat" + assert all(r["messages"][-1]["role"] == "assistant" for r in ds.rows) + + +def test_text_rows_are_the_prose_not_a_transcript(): + ds = emit.to_text([_traj("what is X?", "X is a thing.")]) + assert ds.format == "text" + assert ds.rows == [{"text": "X is a thing."}] + + +def test_preference_rows_carry_all_three_keys(): + ds = emit.to_preference([_traj("q", "good", rejected="bad")]) + assert ds.format == "preference" # what Dataset detection calls it + assert set(ds.rows[0]) == {"prompt", "chosen", "rejected"} + # the exact key set trl's DPOTrainer validates on + assert ds.rows[0]["chosen"] != ds.rows[0]["rejected"] + + +def test_preference_drops_pairs_with_no_contrast(): + with pytest.raises(ValueError, match="no usable preference pairs"): + emit.to_preference([_traj("q", "same", rejected="same")]) + + +def test_grpo_prompt_rows_have_the_column_the_backend_looks_for(): + ds = emit.to_grpo_prompts([_traj("q", "a")]) + assert "prompt" in ds.rows[0] and "answer" in ds.rows[0] + + +def test_groups_are_accepted_by_weighted_rows(): + """The real inject-ready assertion: the GRPO row builder takes them.""" + trajs = [_traj(f"q{i}", f"a{i}", scenario="shared", reward=r) + for i, r in enumerate((0.9, 0.5, 0.1))] + groups = emit.to_groups(trajs) + assert len(groups) == 1 and len(groups[0]) == 3 + rows = weighted_rows(groups) + assert all(set(r) == {"messages", "weight"} for r in rows) + + +def test_groups_without_reward_spread_are_dropped_loudly(): + flat = [_traj(f"q{i}", "a", scenario="shared", reward=0.5) for i in range(3)] + with pytest.raises(ValueError, match="no scored groups"): + emit.to_groups(flat) + + +def test_paraphrase_units_route_through_the_real_bm25_router(): + """MoRE+'s contract: k consecutive rows per fact, reachable by any phrasing.""" + fact_a = ["what does cloud cost per agent run", + "how am i billed each execution", + "pricing for one invocation"] + fact_b = ["where is the company headquartered", + "which city is the office in", + "what is the head office location"] + rows = [_traj(q, "$0.08").messages for q in fact_a] + rows += [_traj(q, "Boston").messages for q in fact_b] + ds = emit.to_chat([Trajectory(messages=m) for m in rows]) + + surrogates = [s for s, _ in mp.split_units(ds, group_size=3)] + assert len(surrogates) == 2 + router = mp.BM25Router.build(surrogates) + # a query phrased like the *third* row of unit 0 still routes to unit 0 + assert router.rank("pricing one invocation", 1)[0][0] == 0 + assert router.rank("which city is the office", 1)[0][0] == 1 diff --git a/tests/test_synth_otlp_roundtrip.py b/tests/test_synth_otlp_roundtrip.py new file mode 100644 index 0000000..8fffea0 --- /dev/null +++ b/tests/test_synth_otlp_roundtrip.py @@ -0,0 +1,99 @@ +"""The OTLP round trip: what we write, our own reader reads back unchanged. + +This is the load-bearing test of the whole synthesizer. `emit.to_otlp` writes +OpenTelemetry GenAI spans and `traces.from_otlp` reconstructs episodes from +them; if those two ever disagree, "inject ready" is a claim rather than a fact. +Tool calls and multi-span agent loops are included because that is exactly where +a naive emitter loses information. +""" + +import json + +from shadowlm import traces +from shadowlm.rl import Trajectory +from shadowlm.synth import emit + +_TOOL = {"type": "function", "function": { + "name": "get_weather", "description": "Weather for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}}} + + +def _agent_episode(reward=1.0): + """ask → tool call → tool result → answer: two assistant turns, so two spans.""" + return Trajectory( + messages=[ + {"role": "system", "content": "You are a weather agent."}, + {"role": "user", "content": "Weather in Paris?"}, + {"role": "assistant", "content": None, "tool_calls": [ + {"id": "call_0", "type": "function", "function": { + "name": "get_weather", "arguments": '{"city": "Paris"}'}}]}, + {"role": "tool", "tool_call_id": "call_0", "content": "18C, sunny"}, + {"role": "assistant", "content": "It's 18°C and sunny in Paris."}, + ], + tools=[_TOOL], reward=reward) + + +def _roundtrip(trajectories, **kwargs): + payload = emit.to_otlp(trajectories, **kwargs) + return traces.from_otlp(payload, reward_key=emit.SCORE_KEY) + + +def test_plain_episode_survives_intact(): + original = Trajectory(messages=[ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi there"}], reward=0.8) + back = _roundtrip([original]) + assert len(back) == 1 + assert back[0].messages == original.messages + assert back[0].reward == 0.8 + + +def test_agent_loop_folds_back_into_one_episode(): + original = _agent_episode() + back = _roundtrip([original]) + assert len(back) == 1, "the two spans should prefix-merge into one episode" + assert back[0].messages == original.messages + + +def test_tool_calls_and_schemas_survive(): + back = _roundtrip([_agent_episode()])[0] + call = back.messages[2]["tool_calls"][0] + assert call["type"] == "function" + assert call["function"]["name"] == "get_weather" + assert json.loads(call["function"]["arguments"]) == {"city": "Paris"} + assert back.messages[3]["tool_call_id"] == "call_0" + assert back.tools == [_TOOL] + + +def test_separate_episodes_stay_separate(): + back = _roundtrip([_agent_episode(), _agent_episode()]) + assert len(back) == 2 + + +def test_rewards_ride_along(): + back = _roundtrip([_agent_episode(reward=0.25)]) + assert back[0].reward == 0.25 + + +def test_written_file_reloads(tmp_path): + path = tmp_path / "spans.json" + emit.to_otlp([_agent_episode()], path=path) + back = traces.from_otlp(path, reward_key=emit.SCORE_KEY) + assert back[0].messages[-1]["content"] == "It's 18°C and sunny in Paris." + + +def test_output_is_deterministic_for_a_seed(): + once = emit.to_otlp([_agent_episode()], seed=7) + twice = emit.to_otlp([_agent_episode()], seed=7) + assert once == twice # ids and timestamps included — no wall clock anywhere + assert emit.to_otlp([_agent_episode()], seed=8) != once + + +def test_payload_is_a_real_otlp_envelope(): + payload = emit.to_otlp([_agent_episode()]) + spans = payload["resourceSpans"][0]["scopeSpans"][0]["spans"] + assert len(spans) == 2 # one per assistant turn + attrs = {a["key"]: a["value"] for a in spans[0]["attributes"]} + assert "stringValue" in attrs["gen_ai.input.messages"] + assert attrs["gen_ai.operation.name"]["stringValue"] == "chat" + assert len(spans[0]["traceId"]) == 32 and len(spans[0]["spanId"]) == 16 diff --git a/tests/test_synth_pipeline.py b/tests/test_synth_pipeline.py new file mode 100644 index 0000000..e15e846 --- /dev/null +++ b/tests/test_synth_pipeline.py @@ -0,0 +1,191 @@ +"""The synthesis pipeline end to end, against a scripted teacher. + +No network, no model: `FakeTeacher` recognises which prompt it was handed and +answers in kind, so every branch — taxonomy, instances, repair, dedup, the judge +gate — is exercised deterministically. +""" + +import json + +import pytest + +import shadowlm as slm +from shadowlm.synth import synthesize + + +class FakeTeacher: + """Answers by recognising the prompt it got. + + `junk_replies` makes the first N conversation calls return unparseable text: + 1 exercises the corrective retry, 2 exhausts it and forces a rejection. + """ + + name = "fake" + parallelism = 1 + + def __init__(self, *, scenarios=8, scores=("0.9",), junk_replies=0): + self.scenarios = scenarios + self.scores = list(scores) + self.junk_replies = junk_replies + self.prompts = [] + self._conversations = 0 + self._scored = 0 + + def chat(self, messages, **_): + prompt = messages[-1]["content"] + self.prompts.append(prompt) + if "0.0 to 1.0" in prompt: + score = self.scores[self._scored % len(self.scores)] + self._scored += 1 + return score + if '"scenario"' in prompt: # taxonomy, or episode-pattern variations + return json.dumps([{"scenario": f"scenario {i}", "difficulty": "easy", + "angle": f"angle {i}"} for i in range(self.scenarios)]) + if "factual claims" in prompt: + return json.dumps(["the sky is blue", "water boils at 100C"]) + if "questions that should all retrieve" in prompt: + fact = prompt.split("FACT: ")[1].split("\n")[0] + return json.dumps({ + "questions": [f"alpha bravo {fact}", f"charlie delta {fact}", + f"echo foxtrot {fact}", f"golf hotel {fact}"], + "answer": f"the answer to {fact}"}) + if "user message this scenario would produce" in prompt: + self._conversations += 1 + return f"user question {self._conversations}" + if "FLAWED" in prompt: + return "a flawed answer" + if "Answer this as well as you possibly can" in prompt: + return "the good answer" + if "training conversation" in prompt: + self._conversations += 1 + if self._conversations <= self.junk_replies: + return "sorry, prose instead of JSON" + return json.dumps({"messages": [ + {"role": "user", "content": f"question {self._conversations}"}, + {"role": "assistant", "content": f"answer {self._conversations}"}]}) + raise AssertionError(f"FakeTeacher got an unexpected prompt:\n{prompt[:300]}") + + +def test_task_seed_produces_chat_rows(): + run = synthesize(task="triage billing email", teacher=FakeTeacher(), + n=8, method="lora", verbose=False) + assert run.format == "chat" + assert run.dataset.format == "chat" + assert len(run.dataset.rows) == 8 + assert all(r["messages"][-1]["role"] == "assistant" for r in run.dataset.rows) + + +def test_the_funnel_reconciles(): + """Every generated row is accounted for — nothing silently disappears.""" + run = synthesize(task="t", teacher=FakeTeacher(scores=("0.9", "0.1")), + n=6, method="lora", verbose=False) + r = run.report + assert r.balanced, r.summary() + assert r.generated == (r.kept + r.rejected_validation + r.rejected_dedup + + r.rejected_judge + r.surplus) + assert r.rejected_judge > 0 # the 0.1 scores were gated out + assert r.teacher_calls > 0 + + +def test_rows_carry_provenance(): + run = synthesize(task="t", teacher=FakeTeacher(), n=4, verbose=False) + meta = run.trajectories[0].metadata + assert meta["source"] == "synth" + assert meta["teacher"] == "fake" + assert meta["taxonomy_path"].startswith("scenario") + assert meta["style"] and "judge_score" in run.trajectories[0].metrics + + +def test_same_seed_gives_the_same_rows(): + kwargs = dict(task="t", n=6, method="lora", verbose=False) + first = synthesize(teacher=FakeTeacher(), **kwargs) + second = synthesize(teacher=FakeTeacher(), **kwargs) + assert first.dataset.rows == second.dataset.rows + + +def test_one_bad_reply_is_repaired_two_is_rejected(): + repaired = synthesize(task="t", teacher=FakeTeacher(junk_replies=1), n=4, + verbose=False) + assert repaired.report.repaired == 1 + assert repaired.report.rejected_validation == 0 + + rejected = synthesize(task="t", teacher=FakeTeacher(junk_replies=2), n=4, + verbose=False) + assert rejected.report.rejected_validation == 1 + assert rejected.report.balanced + + +def test_duplicates_are_rejected_not_shipped(): + """A teacher that repeats itself gets caught rather than padding the set.""" + class Repeater(FakeTeacher): + def chat(self, messages, **_): + if "training conversation" in messages[-1]["content"]: + return json.dumps({"messages": [ + {"role": "user", "content": "the very same question"}, + {"role": "assistant", "content": "the very same answer"}]}) + return super().chat(messages) + + run = synthesize(task="t", teacher=Repeater(), n=8, verbose=False) + assert len(run.dataset.rows) == 1 + assert run.report.rejected_dedup > 0 + assert run.report.balanced + + +def test_document_seed_grounds_on_the_passage(tmp_path): + doc = tmp_path / "notes.md" + doc.write_text("The sky is blue.\n\nWater boils at 100C at sea level.") + teacher = FakeTeacher() + run = synthesize(document=doc, teacher=teacher, n=2, verbose=False) + # the passage reaches both the generator (as source material to stay inside) + # and the judge (as the reference its answer is scored against) + assert "SOURCE MATERIAL" in "".join(teacher.prompts) + assert "REFERENCE: The sky is blue" in "".join(teacher.prompts) + assert run.trajectories[0].metadata["grounding"].startswith("The sky is blue") + + +def test_document_rejects_formats_it_cannot_read(tmp_path): + pdf = tmp_path / "report.pdf" + pdf.write_bytes(b"%PDF-1.4") + with pytest.raises(ValueError, match="convert it to"): + synthesize(document=pdf, teacher=FakeTeacher(), n=2, verbose=False) + + +def test_episodes_seed_never_clones_the_real_data(): + real = slm.Trajectory(messages=[ + {"role": "user", "content": "question 1"}, + {"role": "assistant", "content": "real answer"}]) + + class Cloner(FakeTeacher): + def chat(self, messages, **_): + if "training conversation" in messages[-1]["content"]: + return json.dumps({"messages": [ + {"role": "user", "content": "question 1"}, + {"role": "assistant", "content": "real answer"}]}) + return super().chat(messages) + + # the dedup pool is pre-seeded with the real episodes, so a teacher that + # regurgitates them produces nothing at all rather than laundering the + # user's own data back as "synthetic" + with pytest.raises(RuntimeError, match="nothing usable"): + synthesize(episodes=[real, real], teacher=Cloner(), n=4, verbose=False, + min_score=None) + + +def test_no_seed_and_no_teacher_fail_loudly(): + with pytest.raises(ValueError, match="needs a seed"): + synthesize(teacher=FakeTeacher(), n=2, verbose=False) + with pytest.raises(ValueError, match="needs teacher"): + synthesize(task="t", teacher=None, n=2, verbose=False) + + +def test_everything_rejected_raises_with_the_counts(): + with pytest.raises(RuntimeError, match="nothing usable"): + synthesize(task="t", teacher=FakeTeacher(scores=("0.0",)), n=4, + min_score=0.9, verbose=False) + + +def test_unknown_format_and_method_are_caught_early(): + with pytest.raises(ValueError, match="unknown format"): + synthesize(task="t", teacher=FakeTeacher(), n=2, format="parquet") + with pytest.raises(ValueError, match="unknown method"): + synthesize(task="t", teacher=FakeTeacher(), n=2, method="telepathy") diff --git a/tests/test_synth_quality.py b/tests/test_synth_quality.py new file mode 100644 index 0000000..1f28eb7 --- /dev/null +++ b/tests/test_synth_quality.py @@ -0,0 +1,114 @@ +"""Validators and dedup — the gate everything synthesized passes through.""" + +from shadowlm.synth.quality import Dedup, first_json_array, validate + +_TOOLS = [{"type": "function", "function": {"name": "search", "parameters": {}}}] + + +def _call(name="search", args='{"q": "x"}', cid="call_1"): + return {"id": cid, "type": "function", + "function": {"name": name, "arguments": args}} + + +def test_a_clean_conversation_passes(): + assert validate([{"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}]) == [] + + +def test_must_end_on_an_assistant_turn(): + problems = validate([{"role": "assistant", "content": "hi"}, + {"role": "user", "content": "and you?"}]) + assert any("end with an assistant" in p for p in problems) + + +def test_empty_final_turn_is_caught(): + assert validate([{"role": "user", "content": "hi"}, + {"role": "assistant", "content": ""}]) + + +def test_system_turn_must_come_first_and_once(): + assert validate([{"role": "user", "content": "hi"}, + {"role": "system", "content": "be nice"}, + {"role": "assistant", "content": "ok"}]) + + +def test_unknown_role_is_caught(): + assert validate([{"role": "narrator", "content": "meanwhile"}, + {"role": "assistant", "content": "ok"}]) + + +def test_placeholders_and_boilerplate_are_caught(): + assert validate([{"role": "user", "content": "hi"}, + {"role": "assistant", "content": "Dear [NAME], hello"}]) + assert validate([{"role": "user", "content": "hi"}, + {"role": "assistant", "content": "As an AI language model I"}]) + + +def test_wellformed_tool_exchange_passes(): + assert validate([ + {"role": "user", "content": "search x"}, + {"role": "assistant", "content": None, "tool_calls": [_call()]}, + {"role": "tool", "tool_call_id": "call_1", "content": "found"}, + {"role": "assistant", "content": "here it is"}], tools=_TOOLS) == [] + + +def test_unparseable_tool_arguments_are_caught(): + problems = validate([ + {"role": "user", "content": "x"}, + {"role": "assistant", "content": None, "tool_calls": [_call(args="{oops")]}, + {"role": "tool", "tool_call_id": "call_1", "content": "r"}, + {"role": "assistant", "content": "done"}], tools=_TOOLS) + assert any("unparseable arguments" in p for p in problems) + + +def test_undeclared_tool_is_caught(): + problems = validate([ + {"role": "user", "content": "x"}, + {"role": "assistant", "content": None, "tool_calls": [_call(name="rm_rf")]}, + {"role": "tool", "tool_call_id": "call_1", "content": "r"}, + {"role": "assistant", "content": "done"}], tools=_TOOLS) + assert any("undeclared tool" in p for p in problems) + + +def test_unanswered_tool_call_is_caught(): + problems = validate([ + {"role": "user", "content": "x"}, + {"role": "assistant", "content": None, "tool_calls": [_call()]}, + {"role": "assistant", "content": "done"}], tools=_TOOLS) + assert any("never answered" in p for p in problems) + + +def test_orphan_tool_result_is_caught(): + problems = validate([ + {"role": "user", "content": "x"}, + {"role": "tool", "tool_call_id": "ghost", "content": "r"}, + {"role": "assistant", "content": "done"}]) + assert any("unknown call" in p for p in problems) + + +def test_dedup_rejects_exact_and_near_repeats(): + dedup = Dedup(threshold=0.7) + assert dedup.accept("how much does the plan cost") + assert not dedup.accept("how much does the plan cost") # exact + assert not dedup.accept("how much does the plan cost really") # near + assert dedup.accept("where is the office located") # different + + +def test_dedup_judges_near_repeats_on_the_query_side(): + """Same question, different answer, is still a duplicate question.""" + dedup = Dedup(threshold=0.7) + assert dedup.accept("q: what is the refund window", key="what is the refund window") + assert not dedup.accept("q: what is the refund window — other answer", + key="what is the refund window") + + +def test_dedup_can_be_preseeded_with_real_data(): + dedup = Dedup() + dedup.seed(["the exact question a real user asked"]) + assert not dedup.accept("the exact question a real user asked") + + +def test_first_json_array_digs_it_out_of_prose(): + assert first_json_array('Sure! ["a", "b"] hope that helps') == ["a", "b"] + assert first_json_array("no array here") is None + assert first_json_array('broken [1, 2 then [3, 4]') == [3, 4] diff --git a/tests/test_synth_teacher.py b/tests/test_synth_teacher.py new file mode 100644 index 0000000..ef12d0b --- /dev/null +++ b/tests/test_synth_teacher.py @@ -0,0 +1,135 @@ +"""Teachers: coercion, the HTTP client, and its retry behaviour. + +The OpenAI-compatible teacher is exercised against a real stdlib server on a +loopback port — no mocking of urllib, so the wire format is actually tested. +""" + +import json +import threading +from http.server import BaseHTTPRequestHandler, HTTPServer + +import pytest + +from shadowlm.synth import teacher as tch + + +class _Stub: + """A tiny OpenAI-compatible endpoint. `fail_first` 429s that many times.""" + + def __init__(self, *, reply="hello", fail_first=0, status=None): + self.reply, self.remaining_failures, self.status = reply, fail_first, status + self.requests: list[dict] = [] + outer = self + + class Handler(BaseHTTPRequestHandler): + def log_message(self, *_): + pass + + def do_POST(self): # noqa: N802 + body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + outer.requests.append({"body": body, "auth": self.headers.get("Authorization")}) + if outer.status is not None: + self.send_error(outer.status, "nope") + return + if outer.remaining_failures > 0: + outer.remaining_failures -= 1 + self.send_error(429, "slow down") + return + payload = json.dumps({"choices": [ + {"message": {"role": "assistant", "content": outer.reply}}]}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + self._server = HTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=self._server.serve_forever, daemon=True).start() + + @property + def base_url(self): + return f"http://127.0.0.1:{self._server.server_port}/v1" + + def close(self): + self._server.shutdown() + self._server.server_close() + + +def _teacher(stub, **kwargs): + return tch.OpenAIChatTeacher("test-model", base_url=stub.base_url, + api_key="sk-test", **kwargs) + + +def test_chat_speaks_the_openai_wire_format(): + stub = _Stub(reply="the answer") + try: + out = _teacher(stub).chat([{"role": "user", "content": "hi"}], + temperature=0.3, max_new_tokens=64) + assert out == "the answer" + sent = stub.requests[0] + assert sent["body"]["model"] == "test-model" + assert sent["body"]["messages"] == [{"role": "user", "content": "hi"}] + assert sent["body"]["temperature"] == 0.3 + assert sent["body"]["max_tokens"] == 64 + assert sent["auth"] == "Bearer sk-test" + finally: + stub.close() + + +def test_transient_failures_are_retried(): + stub = _Stub(reply="eventually", fail_first=1) + try: + assert _teacher(stub).chat([{"role": "user", "content": "hi"}]) == "eventually" + assert len(stub.requests) == 2 + finally: + stub.close() + + +def test_a_permanent_error_raises_with_the_servers_words(): + stub = _Stub(status=401) + try: + with pytest.raises(RuntimeError, match="HTTP 401"): + _teacher(stub).chat([{"role": "user", "content": "hi"}]) + assert len(stub.requests) == 1 # 401 is not retried + finally: + stub.close() + + +def test_an_unreachable_endpoint_says_where_it_tried(): + teacher = tch.OpenAIChatTeacher("m", base_url="http://127.0.0.1:1", api_key="k") + with pytest.raises(RuntimeError, match="unreachable at http://127.0.0.1:1"): + teacher.chat([{"role": "user", "content": "hi"}]) + + +def test_a_loaded_model_becomes_a_serialized_teacher(): + class FakeModel: + name = "qwen-tiny" + + def chat(self, messages, **kwargs): + return f"reply to {messages[-1]['content']}" + + teacher = tch.as_teacher(FakeModel()) + assert teacher.name == "qwen-tiny" + assert teacher.parallelism == 1 # backends are not thread-safe + assert teacher.chat([{"role": "user", "content": "x"}]) == "reply to x" + + +def test_an_existing_teacher_passes_through_and_junk_is_rejected(): + original = tch.frontier("gpt-4o", api_key="k") + assert tch.as_teacher(original) is original + with pytest.raises(ValueError, match="needs teacher"): + tch.as_teacher(None) + with pytest.raises(TypeError, match="not a teacher"): + tch.as_teacher(42) + + +def test_counting_teacher_tallies_calls(): + stub = _Stub() + try: + counted = tch.CountingTeacher(_teacher(stub)) + counted.chat([{"role": "user", "content": "a"}]) + counted.chat([{"role": "user", "content": "b"}]) + assert counted.calls == 2 + assert counted.name == "test-model" + finally: + stub.close() From 209bc7303c035933ce51126553dc762953b5c6a3 Mon Sep 17 00:00:00 2001 From: pradipta-lyzr Date: Mon, 3 Aug 2026 23:13:37 +0530 Subject: [PATCH 03/14] synth: reach it from the shell, the server, and the studio MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `shadowlm synth` with --dry-run, so you can see the resolved output shape before spending a teacher call on it. POST/GET /v1/synth follows the model-download pattern — a background thread polled for status — rather than the training queue, since synthesis is teacher calls and would otherwise hold the one training slot for the whole run. The finished rows land in DatasetStore like any other dataset. Teacher API keys are used for the run and never written to disk, unlike the HF token, which is deliberately persisted. Datasets gains a Synthesize tab beside Upload and Hugging Face: task, optional grounding document, target method, and a teacher that is either an OpenAI-compatible endpoint or a model already on this machine. Co-Authored-By: Claude Opus 5 (1M context) --- frontend/src/api.ts | 25 ++++++ frontend/src/pages/Datasets.tsx | 148 ++++++++++++++++++++++++++++++-- shadowlm/cli.py | 83 ++++++++++++++++++ shadowlm/serve.py | 79 +++++++++++++++++ tests/test_serve_synth.py | 94 ++++++++++++++++++++ tests/test_synth_cli.py | 65 ++++++++++++++ 6 files changed, 486 insertions(+), 8 deletions(-) create mode 100644 tests/test_serve_synth.py create mode 100644 tests/test_synth_cli.py diff --git a/frontend/src/api.ts b/frontend/src/api.ts index f271648..1ace553 100644 --- a/frontend/src/api.ts +++ b/frontend/src/api.ts @@ -151,6 +151,31 @@ export const addHFDataset = ( eval_split: evalSplit || null }) }); export const deleteDataset = (id: string) => api<{ ok: boolean }>(`/v1/datasets/${id}`, { method: "DELETE" }); + +// ---- synthesis: make a dataset instead of bringing one ---------------------- +export interface SynthStatus { + synth_id: string; + name: string; + status: "running" | "succeeded" | "failed"; + kept: number; + requested: number; + dataset_id?: string; + error?: string; + logs?: string[]; +} +export interface SynthRequest { + name: string; + n: number; + method?: string; + min_score?: number; + task?: string; + document?: string; + dataset_id?: string; + teacher: { kind: "openai" | "local"; model: string; base_url?: string; api_key?: string }; +} +export const startSynth = (body: SynthRequest) => + api<{ synth_id: string }>("/v1/synth", { method: "POST", body: JSON.stringify(body) }); +export const getSynthRun = (id: string) => api(`/v1/synth/${id}`); export const getModels = () => api<{ catalog: CatalogModel[]; recent: string[]; server_backend: string }>("/v1/models"); export const getDownloads = () => diff --git a/frontend/src/pages/Datasets.tsx b/frontend/src/pages/Datasets.tsx index ae78f35..a3a8d4a 100644 --- a/frontend/src/pages/Datasets.tsx +++ b/frontend/src/pages/Datasets.tsx @@ -1,12 +1,14 @@ -// Dataset library — upload JSONL, or reference a HuggingFace dataset (with a -// streamed preview before you add it). Both become trainable by reference. +// Dataset library — upload JSONL, reference a HuggingFace dataset (with a +// streamed preview before you add it), or synthesize one from a task +// description. All three become trainable by reference. import { useEffect, useRef, useState } from "react"; -import { Database, Search, Upload } from "lucide-react"; +import { Database, Search, Sparkles, Upload } from "lucide-react"; import { - addHFDataset, createDataset, deleteDataset, getDataset, getDatasets, hfInfo, previewHF, + addHFDataset, createDataset, deleteDataset, getDataset, getDatasets, getMethods, + getSynthRun, hfInfo, previewHF, startSynth, } from "../api"; -import type { DatasetMeta, HFPreview } from "../api"; -import { Modal, ModalHeader, PageHeader, btnGhost, btnPrimary } from "../ui"; +import type { DatasetMeta, HFPreview, MethodInfo, SynthStatus } from "../api"; +import { Field, Modal, ModalHeader, PageHeader, btnGhost, btnPrimary } from "../ui"; const FORMAT_COLORS: Record = { chat: "bg-primary/10 text-primary border-primary/30", @@ -33,7 +35,7 @@ export default function Datasets() { const [list, setList] = useState([]); const [search, setSearch] = useState(""); const [view, setView] = useState<"mine" | "explore">("mine"); - const [tab, setTab] = useState<"none" | "upload" | "hf">("none"); + const [tab, setTab] = useState<"none" | "upload" | "hf" | "synth">("none"); const [rowPreview, setRowPreview] = useState(null); const [previewing, setPreviewing] = useState(null); @@ -74,9 +76,13 @@ export default function Datasets() { + + )} + + ); + } + + return ( +
+ setName(e.target.value)} /> + +