From d25daf68c1273bcf044b8ccdf003394b98d2d4ed Mon Sep 17 00:00:00 2001 From: mgoin Date: Wed, 22 Jul 2026 15:19:01 +0000 Subject: [PATCH] Add Verifiers agentic regeneration demo --- docs/.nav.yml | 1 + .../tutorials/agentic_regeneration.md | 107 ++++++++ docs/user_guide/tutorials/index.md | 6 + scripts/agentic_regeneration/.gitignore | 2 + scripts/agentic_regeneration/.python-version | 1 + .../configs/livecodebench.toml | 32 +++ .../agentic_regeneration/configs/prolog.toml | 31 +++ scripts/agentic_regeneration/configs/r2e.toml | 31 +++ .../agentic_regeneration/convert_traces.py | 252 ++++++++++++++++++ scripts/agentic_regeneration/pyproject.toml | 13 + .../integration/datagen/test_preprocessing.py | 24 ++ tests/unit/test_agentic_regeneration.py | 131 +++++++++ 12 files changed, 631 insertions(+) create mode 100644 docs/user_guide/tutorials/agentic_regeneration.md create mode 100644 scripts/agentic_regeneration/.gitignore create mode 100644 scripts/agentic_regeneration/.python-version create mode 100644 scripts/agentic_regeneration/configs/livecodebench.toml create mode 100644 scripts/agentic_regeneration/configs/prolog.toml create mode 100644 scripts/agentic_regeneration/configs/r2e.toml create mode 100644 scripts/agentic_regeneration/convert_traces.py create mode 100644 scripts/agentic_regeneration/pyproject.toml create mode 100644 tests/unit/test_agentic_regeneration.py diff --git a/docs/.nav.yml b/docs/.nav.yml index f22fb3f84..5dd8913fc 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -19,6 +19,7 @@ nav: - Train P-Eagle Model Offline: user_guide/tutorials/train_peagle_offline.md - Train MTP Model Online: user_guide/tutorials/train_mtp_online.md - Response Regeneration: user_guide/tutorials/response_regeneration.md + - Agentic Regeneration: user_guide/tutorials/agentic_regeneration.md - Evaluating Model Performance: user_guide/tutorials/evaluating_performance.md - Serve in vLLM: user_guide/tutorials/serve_vllm.md - Developer Guide: diff --git a/docs/user_guide/tutorials/agentic_regeneration.md b/docs/user_guide/tutorials/agentic_regeneration.md new file mode 100644 index 000000000..19be500dd --- /dev/null +++ b/docs/user_guide/tutorials/agentic_regeneration.md @@ -0,0 +1,107 @@ +# Agentic Regeneration with Verifiers + +Agentic regeneration should have a narrow integration boundary: + +1. the upstream Verifiers CLI runs an installed taskset and writes `config.toml` plus `traces.jsonl`; +2. one generic adapter converts each native root-to-leaf trace branch to exact `input_ids` and `loss_mask`; +3. Speculators replays those IDs through the same checkpoint to collect aligned hidden states. + +There is no Speculators environment wrapper or rollout loop. + +## Pinned UV environment + +The adapter has its own UV project under `scripts/agentic_regeneration/`. Its small `pyproject.toml` pins the compatibility boundary: Python 3.12, the stable `verifiers==0.2.0` release, and one exact research-environments commit. + +```bash +uv sync --project scripts/agentic_regeneration +``` + +This environment is intentionally separate from the main Speculators environment. Updating Verifiers or an upstream taskset is an explicit `pyproject.toml` change, while UV resolves ordinary transitive dependencies at install time. + +## Environment profiles + +Three native CLI profiles provide a useful progression: + +| Profile | What it exercises | Cost | +|---|---|---| +| `prolog.toml` | Multi-turn code editing, execution feedback, and hidden verification | Public SWI-Prolog image; recommended smoke test | +| `livecodebench.toml` | Python code generation with hidden execution tests | Shared Python image; useful correctness baseline | +| `r2e.toml` | Multi-turn inspection and repair of a real Python repository | Per-task repository image; production acceptance test | + +Prolog is the quickest genuinely agentic check: tasks are generated locally and the public sandbox image is small. LiveCodeBench is lighter coding data but is capped at one turn in the supplied profile. R2E-Gym is the target for realistic Python repair, and its image and dataset cost is inherent to reproducibly executing arbitrary repositories. + +All three tasksets and the coding harness are upstream. R2E uses Verifiers' built-in stable `default` harness, which supplies bash and edit tools; Speculators contains no R2E-specific code. + +## 1. Start Qwen3-8B in vLLM + +```bash +python scripts/launch_vllm.py Qwen/Qwen3-8B \ + --hidden-states-path output/agentic_regen/server_hidden_states \ + -- \ + --host 127.0.0.1 \ + --port 8000 \ + --gpu-memory-utilization 0.5 \ + --max-model-len 16384 +``` + +The profiles use Verifiers' `train` client and its Qwen3 renderer with `enable_thinking = false`. Rendering and tool-call parsing therefore happen client-side, while vLLM only serves exact token generation. Disabling thinking keeps this coding profile from spending its whole turn on an unparsed reasoning block. OpenAI chat-completions tool-parser flags are not part of this path. + +## 2. Run an upstream taskset + +Start with Prolog: + +```bash +VLLM_API_KEY=EMPTY uv run --project scripts/agentic_regeneration \ + eval @ scripts/agentic_regeneration/configs/prolog.toml +``` + +Switch only the final profile path for another environment: + +```text +scripts/agentic_regeneration/configs/livecodebench.toml +scripts/agentic_regeneration/configs/r2e.toml +``` + +The stable CLI owns dataset loading, Docker lifecycle, tools, turns, retries, and scoring. Each profile is a one-task smoke run; increase `num_tasks`, `num_rollouts`, and `max_concurrent` in TOML for collection. + +The `train` client is important: its native trace contains the exact rendered prompt and sampled completion token spans. A normal chat-completions response is not sufficient for exact hidden-state alignment. + +## 3. Convert native traces + +For the Prolog profile: + +```bash +uv run --project scripts/agentic_regeneration \ + python scripts/agentic_regeneration/convert_traces.py \ + --traces output/agentic_regen/prolog/traces.jsonl \ + --outfile output/agentic_regen/prolog/trajectories.jsonl +``` + +The adapter targets the stable Verifiers `WireTrace` format. It reconstructs every root-to-leaf branch, retains messages and tool-call linkage, and reads model, endpoint, and taskset metadata from the sibling `config.toml`. Stable 0.2.0 does not duplicate tool schemas in the JSON trace; exact-token replay does not need to render them again. + +Completed zero-reward traces are retained because unsuccessful actions are valid on-policy data. + +## 4. Prepare and replay exact IDs + +```bash +python scripts/prepare_data.py \ + --model Qwen/Qwen3-8B \ + --data output/agentic_regen/prolog/trajectories.jsonl \ + --seq-length 16384 \ + --minimum-valid-tokens 1 \ + --num-preprocessing-workers 1 \ + --output output/agentic_regen/prolog/preprocessed + +python scripts/data_generation_offline.py \ + --model Qwen/Qwen3-8B \ + --endpoint http://127.0.0.1:8000/v1 \ + --preprocessed-data output/agentic_regen/prolog/preprocessed \ + --output output/agentic_regen/prolog/eagle_data \ + --concurrency 1 \ + --validate-outputs \ + --fail-on-error +``` + +When a record contains both `input_ids` and `loss_mask`, preprocessing passes them through together rather than applying the chat template again. `--validate-outputs` then checks that vLLM returned those same IDs and an aligned hidden-state sequence. + +The resulting states are on-policy for the served checkpoint because the model sampled every assistant turn against real observations, and replay teacher-forces that exact branch through the same weights. Regenerate after a policy update. diff --git a/docs/user_guide/tutorials/index.md b/docs/user_guide/tutorials/index.md index 8ab166060..82bb0100b 100644 --- a/docs/user_guide/tutorials/index.md +++ b/docs/user_guide/tutorials/index.md @@ -44,6 +44,12 @@ Regenerate dataset responses using your target model for improved drafter alignm **Time required:** ~10 minutes +## [Agentic Regeneration with Verifiers](agentic_regeneration.md) + +Collect exact-token, on-policy agent trajectories and aligned vLLM hidden states. + +**Time required:** ~10 minutes + ## [Evaluating Model Performance](evaluating_performance.md) Benchmark and evaluate your trained speculator models. diff --git a/scripts/agentic_regeneration/.gitignore b/scripts/agentic_regeneration/.gitignore new file mode 100644 index 000000000..e74dcbf87 --- /dev/null +++ b/scripts/agentic_regeneration/.gitignore @@ -0,0 +1,2 @@ +.venv/ +uv.lock diff --git a/scripts/agentic_regeneration/.python-version b/scripts/agentic_regeneration/.python-version new file mode 100644 index 000000000..e4fba2183 --- /dev/null +++ b/scripts/agentic_regeneration/.python-version @@ -0,0 +1 @@ +3.12 diff --git a/scripts/agentic_regeneration/configs/livecodebench.toml b/scripts/agentic_regeneration/configs/livecodebench.toml new file mode 100644 index 000000000..6306a9e94 --- /dev/null +++ b/scripts/agentic_regeneration/configs/livecodebench.toml @@ -0,0 +1,32 @@ +model = "Qwen/Qwen3-8B" +num_tasks = 1 +num_rollouts = 1 +max_concurrent = 1 +max_turns = 1 +push = false +rich = false +output_dir = "output/agentic_regen/livecodebench" + +[taskset] +id = "livecodebench-v1" +difficulty = "easy" + +[harness] +id = "default" + +[harness.runtime] +type = "docker" + +[client] +type = "train" +base_url = "http://127.0.0.1:8000/v1" +api_key_var = "VLLM_API_KEY" +renderer_model_name = "Qwen/Qwen3-8B" + +[client.renderer] +name = "qwen3" +enable_thinking = false + +[sampling] +temperature = 0.0 +max_tokens = 4096 diff --git a/scripts/agentic_regeneration/configs/prolog.toml b/scripts/agentic_regeneration/configs/prolog.toml new file mode 100644 index 000000000..0d6fdaf83 --- /dev/null +++ b/scripts/agentic_regeneration/configs/prolog.toml @@ -0,0 +1,31 @@ +model = "Qwen/Qwen3-8B" +num_tasks = 1 +num_rollouts = 1 +max_concurrent = 1 +max_turns = 6 +push = false +rich = false +output_dir = "output/agentic_regen/prolog" + +[taskset] +id = "prolog-v1" + +[harness] +id = "default" + +[harness.runtime] +type = "docker" + +[client] +type = "train" +base_url = "http://127.0.0.1:8000/v1" +api_key_var = "VLLM_API_KEY" +renderer_model_name = "Qwen/Qwen3-8B" + +[client.renderer] +name = "qwen3" +enable_thinking = false + +[sampling] +temperature = 0.0 +max_tokens = 4096 diff --git a/scripts/agentic_regeneration/configs/r2e.toml b/scripts/agentic_regeneration/configs/r2e.toml new file mode 100644 index 000000000..ba5e43f3f --- /dev/null +++ b/scripts/agentic_regeneration/configs/r2e.toml @@ -0,0 +1,31 @@ +model = "Qwen/Qwen3-8B" +num_tasks = 1 +num_rollouts = 1 +max_concurrent = 1 +max_turns = 6 +push = false +rich = false +output_dir = "output/agentic_regen/r2e" + +[taskset] +id = "r2e-gym-v1" + +[harness] +id = "default" + +[harness.runtime] +type = "docker" + +[client] +type = "train" +base_url = "http://127.0.0.1:8000/v1" +api_key_var = "VLLM_API_KEY" +renderer_model_name = "Qwen/Qwen3-8B" + +[client.renderer] +name = "qwen3" +enable_thinking = false + +[sampling] +temperature = 0.0 +max_tokens = 4096 diff --git a/scripts/agentic_regeneration/convert_traces.py b/scripts/agentic_regeneration/convert_traces.py new file mode 100644 index 000000000..52ead9bcd --- /dev/null +++ b/scripts/agentic_regeneration/convert_traces.py @@ -0,0 +1,252 @@ +#!/usr/bin/env python3 +"""Convert native Verifiers traces into Speculators trajectory JSONL.""" + +import argparse +import json +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +import tomllib +import verifiers.v1 as vf + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Convert Verifiers eval traces into exact-token, " + "Speculators-compatible trajectories." + ) + ) + parser.add_argument( + "--traces", + nargs="+", + required=True, + help="Native Verifiers traces.jsonl files", + ) + parser.add_argument( + "--outfile", + default="output/agentic_regen/trajectories.jsonl", + help="Speculators-compatible trajectory JSONL", + ) + parser.add_argument( + "--model", + default=None, + help="Override model metadata (otherwise read from the run config)", + ) + parser.add_argument( + "--endpoint", + default=None, + help="Override endpoint metadata (otherwise read from the run config)", + ) + parser.add_argument( + "--environment", + default=None, + help="Override environment metadata (otherwise read from the run config)", + ) + parser.add_argument( + "--include-errors", + action="store_true", + help="Export partial errored traces in addition to completed traces", + ) + return parser.parse_args() + + +def _content_to_json(content: Any) -> Any: + if isinstance(content, list): + return [part.model_dump(mode="json") for part in content] + return content + + +def message_to_openai(message: vf.Message) -> dict[str, Any]: + """Convert a Verifiers message without losing tool-call linkage.""" + if isinstance(message, vf.AssistantMessage): + result: dict[str, Any] = { + "role": "assistant", + "content": message.content or "", + } + if message.reasoning_content: + result["reasoning_content"] = message.reasoning_content + if message.tool_calls: + result["tool_calls"] = [ + { + "id": call.id, + "type": "function", + "function": { + "name": call.name, + "arguments": call.arguments, + }, + } + for call in message.tool_calls + ] + return result + if isinstance(message, vf.ToolMessage): + result = { + "role": "tool", + "content": _content_to_json(message.content), + "tool_call_id": message.tool_call_id, + } + if message.name: + result["name"] = message.name + return result + return { + "role": message.role, + "content": _content_to_json(message.content), + } + + +def tool_to_openai(tool: vf.Tool) -> dict[str, Any]: + function: dict[str, Any] = { + "name": tool.name, + "description": tool.description, + "parameters": tool.parameters, + } + if tool.strict is not None: + function["strict"] = tool.strict + return {"type": "function", "function": function} + + +def trace_to_dataset_records( + trace: vf.Trace, *, model: str, endpoint: str, environment_id: str +) -> list[dict[str, Any]]: + """Export one record for every root-to-leaf branch in a Verifiers trace.""" + # Stable Verifiers stores exact rendered tokens and messages, but does not + # duplicate tool schemas on disk. Preserve them if a future release adds them. + tools = [ + tool_to_openai(tool) + for tool in getattr(trace, "tools", None) or [] + ] + usage = trace.usage.model_dump(mode="json") if trace.usage else None + records = [] + for branch in trace.branches: + conversations = [message_to_openai(message) for message in branch.messages] + if not conversations: + continue + token_ids = branch.token_ids or [] + sampled_mask = branch.sampled_mask or [] + has_exact_tokens = bool(token_ids) and len(token_ids) == len(sampled_mask) + records.append( + { + "id": f"{trace.id}:branch-{branch.index}", + "conversations": conversations, + "tools": tools, + "input_ids": token_ids if has_exact_tokens else None, + "loss_mask": sampled_mask if has_exact_tokens else None, + "metadata": { + "source": "verifiers", + "environment": environment_id, + "trace_id": trace.id, + "branch_index": branch.index, + "model": model, + "endpoint": endpoint, + "on_policy": True, + "reward": trace.reward, + "rewards": trace.rewards, + "metrics": trace.metrics, + "num_turns": trace.num_turns, + "stop_condition": trace.stop_condition, + "usage": usage, + "verifiers_token_ids_available": has_exact_tokens, + "token_count": len(token_ids), + "sampled_token_count": sum(sampled_mask), + }, + } + ) + return records + + +def _run_metadata(path: Path) -> tuple[str, str, str]: + """Read taskset, model, and endpoint from the CLI's sibling config.toml.""" + config_path = path.parent / "config.toml" + if not config_path.exists(): + return "", "", "" + try: + with config_path.open("rb") as file: + config = tomllib.load(file) + except (OSError, tomllib.TOMLDecodeError) as error: + raise ValueError(f"Invalid Verifiers run config: {config_path}") from error + + taskset = config.get("taskset", {}) + client = config.get("client", {}) + return ( + str(taskset.get("id", "")), + str(config.get("model", "")), + str(client.get("base_url", "")), + ) + + +def iter_native_traces( + paths: list[str], +) -> Iterator[tuple[vf.Trace, str, str, str]]: + """Read Verifiers 0.2.0 native trace records and their run metadata.""" + for path in paths: + trace_path = Path(path) + environment, model, endpoint = _run_metadata(trace_path) + with trace_path.open(encoding="utf-8") as file: + for line_number, line in enumerate(file, start=1): + if not line.strip(): + continue + try: + record = json.loads(line) + yield ( + vf.WireTrace.model_validate(record), + environment, + model, + endpoint, + ) + except (json.JSONDecodeError, ValueError) as error: + raise ValueError( + f"Invalid Verifiers record at {path}:{line_number}" + ) from error + + +def write_jsonl(path: str, rows: list[dict[str, Any]]) -> None: + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + with output.open("w", encoding="utf-8") as file: + for row in rows: + file.write(json.dumps(row, ensure_ascii=False) + "\n") + + +def main() -> None: + args = parse_args() + rows = [] + trace_count = 0 + skipped_count = 0 + for trace, native_environment, native_model, native_endpoint in ( + iter_native_traces(args.traces) + ): + trace_count += 1 + if not args.include_errors and ( + not trace.is_completed or trace.has_error + ): + skipped_count += 1 + continue + rows.extend( + trace_to_dataset_records( + trace, + model=args.model or native_model, + endpoint=args.endpoint or native_endpoint, + environment_id=args.environment or native_environment, + ) + ) + + if not rows: + raise RuntimeError("No trajectory branches were exported") + exact_rows = sum( + row["metadata"]["verifiers_token_ids_available"] for row in rows + ) + if exact_rows != len(rows): + raise RuntimeError( + "Some branches lack exact token IDs; rerun eval with client.type=train" + ) + + write_jsonl(args.outfile, rows) + print(f"Read traces: {trace_count} (skipped: {skipped_count})") + print(f"Exported branches: {len(rows)}") + print(f"Exact-token branches: {exact_rows}/{len(rows)}") + print(f"Trajectories: {args.outfile}") + + +if __name__ == "__main__": + main() diff --git a/scripts/agentic_regeneration/pyproject.toml b/scripts/agentic_regeneration/pyproject.toml new file mode 100644 index 000000000..532065529 --- /dev/null +++ b/scripts/agentic_regeneration/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "speculators-agentic-regeneration" +version = "0.1.0" +requires-python = ">=3.11,<3.14" +dependencies = [ + "verifiers==0.2.0", + "prolog-v1 @ git+https://github.com/PrimeIntellect-ai/research-environments.git@54f30788fb9cf16c541feea8afcd83bfb3f9fb60#subdirectory=environments/reasoning/prolog_v1", + "livecodebench-v1 @ git+https://github.com/PrimeIntellect-ai/research-environments.git@54f30788fb9cf16c541feea8afcd83bfb3f9fb60#subdirectory=environments/code/livecodebench_v1", + "r2e-gym-v1 @ git+https://github.com/PrimeIntellect-ai/research-environments.git@54f30788fb9cf16c541feea8afcd83bfb3f9fb60#subdirectory=environments/swe/r2e_gym_v1", +] + +[tool.uv] +package = false diff --git a/tests/integration/datagen/test_preprocessing.py b/tests/integration/datagen/test_preprocessing.py index 2d46d7290..fd105c863 100644 --- a/tests/integration/datagen/test_preprocessing.py +++ b/tests/integration/datagen/test_preprocessing.py @@ -716,6 +716,30 @@ def test_preprocess_batch_truncation(): assert len(results["loss_mask"][0]) <= max_length +@pytest.mark.sanity +def test_preprocess_batch_preserves_exact_rollout_tokens(): + """Exact rollout IDs and their mask bypass chat-template rendering.""" + results = _preprocess_batch( + { + "conversations": [ + [ + {"role": "user", "content": "Use the tool"}, + {"role": "assistant", "content": "done"}, + ] + ], + "input_ids": [[11, 12, 13, 14]], + "loss_mask": [[False, True, True, False]], + }, + processor=object(), + max_length=3, + assistant_pattern=None, + ) + + assert results["input_ids"][0].tolist() == [11, 12, 13] + assert results["loss_mask"][0].tolist() == [0, 1, 1] + assert results["seq_len"] == [3] + + @pytest.mark.sanity def test_preprocess_batch_uses_hf_assistant_mask(): """Test that HF assistant token mask is used when supported.""" diff --git a/tests/unit/test_agentic_regeneration.py b/tests/unit/test_agentic_regeneration.py new file mode 100644 index 000000000..539bbc55e --- /dev/null +++ b/tests/unit/test_agentic_regeneration.py @@ -0,0 +1,131 @@ +"""Unit tests for native Verifiers trace conversion and exact-token replay.""" + +import importlib.util +from pathlib import Path + +import pytest + +pytest.importorskip("verifiers.v1") + +import verifiers.v1 as vf + +SCRIPT_DIR = Path(__file__).resolve().parents[2] / "scripts" / "agentic_regeneration" + + +def _load_script(): + spec = importlib.util.spec_from_file_location( + "agentic_regeneration_converter", SCRIPT_DIR / "convert_traces.py" + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_trace_export_preserves_tool_call_linkage(): + script = _load_script() + trace = vf.Trace( + id="trace-1", + task=vf.TraceTask( + type="ArithmeticTask", + data=vf.TaskData(idx=0, prompt="Use the tool"), + ), + nodes=[ + vf.MessageNode( + message=vf.UserMessage(content="Use the tool"), + token_ids=[1, 2], + mask=[False, False], + ), + vf.MessageNode( + parent=0, + sampled=True, + token_ids=[3, 4], + mask=[True, True], + message=vf.AssistantMessage( + content=None, + tool_calls=[ + vf.ToolCall( + id="call-1", + name="multiply_and_add", + arguments='{"a":7,"b":8,"c":9}', + ) + ], + ), + ), + vf.MessageNode( + parent=1, + token_ids=[5], + mask=[False], + message=vf.ToolMessage( + tool_call_id="call-1", + content="65", + ), + ), + vf.MessageNode( + parent=2, + sampled=True, + token_ids=[6], + mask=[True], + message=vf.AssistantMessage(content="65"), + ), + ], + is_completed=True, + rewards={"correct_after_tool": 1.0}, + ) + + (row,) = script.trace_to_dataset_records( + trace, + model="Qwen/Qwen3-0.6B", + endpoint="http://localhost:8000/v1", + environment_id="agentic_regen_env", + ) + + assert [message["role"] for message in row["conversations"]] == [ + "user", + "assistant", + "tool", + "assistant", + ] + assert row["conversations"][1]["tool_calls"][0]["id"] == "call-1" + assert row["conversations"][2]["tool_call_id"] == "call-1" + assert row["tools"] == [] + assert row["metadata"]["on_policy"] is True + assert row["input_ids"] == [1, 2, 3, 4, 5, 6] + assert row["loss_mask"] == [False, False, True, True, False, True] + assert row["metadata"]["verifiers_token_ids_available"] is True + + +def test_native_trace_reader_uses_sibling_run_config(tmp_path): + script = _load_script() + trace = vf.Trace( + id="trace-1", + task=vf.TraceTask( + type="Task", data=vf.TaskData(idx=0, prompt="hello") + ), + nodes=[ + vf.MessageNode( + message=vf.UserMessage(content="hello"), + token_ids=[1], + mask=[False], + ) + ], + is_completed=True, + ) + path = tmp_path / "traces.jsonl" + path.write_text(trace.model_dump_json() + "\n", encoding="utf-8") + (tmp_path / "config.toml").write_text( + 'model = "Qwen/Qwen3-8B"\n' + "[taskset]\n" + 'id = "upstream-v1"\n' + "[client]\n" + 'base_url = "http://localhost:8000/v1"\n', + encoding="utf-8", + ) + + ((loaded_trace, environment_id, model, endpoint),) = ( + script.iter_native_traces([str(path)]) + ) + + assert environment_id == "upstream-v1" + assert model == "Qwen/Qwen3-8B" + assert endpoint == "http://localhost:8000/v1" + assert loaded_trace.branches[0].token_ids == [1]