| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| from .backends.image import DemoImageBackend |
| from .backends.text import DemoTextBackend |
| from .orchestrator import ForestOrchestrator |
| from .trace import TraceRecorder |
|
|
| SCENARIOS = ( |
| ("work", "starting a new design job and worried about belonging"), |
| ("study", "preparing for a difficult exam after one disappointing result"), |
| ("creative", "returning to a novel after months of feeling stuck"), |
| ) |
|
|
|
|
| def build_trace_dataset(output: Path) -> int: |
| output.unlink(missing_ok=True) |
| output.parent.mkdir(parents=True, exist_ok=True) |
| row_count = 0 |
| for index, (category, situation) in enumerate(SCENARIOS): |
| trace_id = f"forest-{index + 1:02d}" |
| temporary = output.with_name(f".{trace_id}.jsonl") |
| temporary.unlink(missing_ok=True) |
| orchestrator = ForestOrchestrator( |
| text_backend=DemoTextBackend(), |
| image_backend=DemoImageBackend(width=128, height=96), |
| trace_recorder=TraceRecorder(temporary), |
| ) |
| list(orchestrator.generate("Sample visitor", situation, seed=3407 + index * 10)) |
| records = [ |
| { |
| "stage": "scenario", |
| "scenario_category": category, |
| "public_synthetic_input": situation, |
| }, |
| *[ |
| json.loads(line) |
| for line in temporary.read_text(encoding="utf-8").splitlines() |
| ], |
| ] |
| with output.open("a", encoding="utf-8") as trace_file: |
| for sequence, record in enumerate(records): |
| row = { |
| "trace_id": trace_id, |
| "sequence": sequence, |
| "stage": record["stage"], |
| "model": record.get("model", ""), |
| "payload_json": json.dumps(record, ensure_ascii=True, sort_keys=True), |
| } |
| trace_file.write(json.dumps(row, ensure_ascii=True) + "\n") |
| row_count += 1 |
| temporary.unlink() |
| return row_count |
|
|
