mlinegar/ThinkingTrees

2

stars

50

commits

Python

primary language

Jul 2, 2026

updated

README

ThinkingTrees: Oracle-Preserving Summarization (OPS)

Hierarchical summarization with verifiable information preservation guarantees. Build recursive summarization trees that maintain task-critical information through probabilistic auditing and DSPy-based optimization.

ThinkingTrees And treepo

This repository now exposes two public faces:

  • ThinkingTrees: the full platform for long-document OPS pipelines, task plugins, training, and deployment.
  • treepo: a focused PyTorch package under treepo/ for method-level simulations and reports.

The treepo package is the canonical home for the new HyperLogLog streaming/cardinality work. Typical commands:

cd treepo
pip install -e ".[torch]"
treepo-bench suite cardinality-paper --out-root ../outputs/cardinality --jobs 4
treepo-bench report cardinality --output-root ../outputs/cardinality

Supported Runtime-Eval Quick Start

The official lightweight path uses the LongBench v2 fixture and mock inference, so it does not require local model servers:

source venv/bin/activate

python scripts/run_runtime_eval.py plan \
  --config config/runtime_eval/longbench_v2_smoke.yaml \
  --experiment-id longbench_v2_smoke \
  --json

python scripts/run_runtime_eval.py init \
  --config config/runtime_eval/longbench_v2_smoke.yaml \
  --output-dir outputs/runtime_eval \
  --experiment-id longbench_v2_smoke

python scripts/run_runtime_eval.py run \
  --experiment-dir outputs/runtime_eval/longbench_v2_smoke \
  --mock-llm \
  --max-problems 1

python scripts/run_runtime_eval.py aggregate \
  --experiment-dir outputs/runtime_eval/longbench_v2_smoke

Runtime methods use paper-facing roles:

  • scorer: final task scoring / answer prediction.
  • summarizer: tree-summary generation, defaulting to scorer if omitted.
  • embedder: retrieval and embedding evidence selection.
  • state_model: deterministic or learned state machinery.
  • oracle: benchmark labels or trusted evaluator provenance.

The canonical experiment artifacts are experiment_manifest.json, experiment_status.json, artifacts.json, results.jsonl, and, for active runtime paths, compact calls.jsonl traces without prompts or full contexts.

For live vLLM/SGLang/OpenAI-compatible endpoints, start servers first and set the scorer, summarizer, embedder, and state_model sections in config/runtime_eval/*.yaml.

Supported Entrypoints

New work should start from the supported experiment front doors:

  • scripts/run_runtime_eval.py for LongBench/RULER-style runtime evaluation.
  • scripts/run_experiment.py for planning, launching, collecting, and reporting canonical experiment sidecars around supported scripts.
  • scripts/run_longbench_batched_example.py and scripts/run_manifesto_batched_example.py for batched document pipelines.
  • scripts/run_method_compare.py, scripts/report_method_compare.py, and scripts/run_method_compare_lbv2.py for method grids.
  • src/training/run_pipeline.py, scripts/train_neural_operators.py, and scripts/train_ctreepo.py for training/operator paths.

The maintained inventory is config/runtime_umbrella_entrypoints.yaml. It classifies supported scripts explicitly and treats older simulations, one-off reports, and historical demos as legacy unless listed there.

Markov Contextual-Sufficiency Probe

The paper hazard-panel contextual controls are available through the ctreepo sim run contextual-sbijax CLI. The current exact-zero path is the repo-owned local-law trainer; package NASS/NASSS remain approximate baselines. After installing the optional JAX/SBI extras, use the saved MarkovOPSDataBundle panels directly:

source venv/bin/activate
python -m pip install -e ".[contextual_sbi]"

XLA_PYTHON_CLIENT_PREALLOCATE=false ctreepo sim run contextual-sbijax \
  --data-source markov \
  --load-data-bundle outputs/_bundles/markov_hazard_panels/paper_hazard_panel_v1_t128/seed_0/base_bundle.json \
  --sbijax-trainer learned_local_laws \
  --sbijax-method nasss \
  --sbijax-package-theta markov_exact_sketch \
  --sbijax-input-encoding markov_exact_sketch \
  --train-docs 1024 \
  --val-docs 256 \
  --test-docs 256 \
  --fragment-len 1 \
  --context-samples-per-doc 1 \
  --response-signature-contexts 16 \
  --response-signature-slices 8 \
  --embedding-dim 32 \
  --state-dim 25 \
  --hidden-dim 128 \
  --learning-rate 0.0003 \
  --n-iter 1000 \
  --batch-size 128 \
  --local-law-weight 1.0 \
  --local-law-leaf-weight 1.0 \
  --local-law-merge-weight 1.0 \
  --local-law-idempotence-weight 1.0 \
  --local-law-contextual-weight 1.0 \
  --seed 0 \
  --output-root outputs/contextual_sbijax_exact_zero_smoke

Use ctreepo-contextual-sbijax as the equivalent standalone console script after reinstalling editable entrypoints. See docs/contextual_sbijax_walkthrough.md and docs/markov_hazard_panels.md for the exact-sketch oracle checks and leaf-length grid commands. Paper-facing leaf-length grids should include the full t128 ladder 1, 2, 4, 8, 16, 32, 64; the small leaves are required edge/capacity checks, not optional smoke rungs.

The command above is the exact-zero recovery path. Use --sbijax-trainer package --sbijax-method nasss only when you want the pre-resolution package baseline that plateaus on its contrastive objective. See docs/contextual_sbijax_optimize_to_zero_resolved_2026-05-05.md for the resolution write-up and full leaf-grid artifacts. The current post-resolution status page is docs/markov_sim_status.md. The post-resolution ablation handoff is docs/markov_contextual_sufficiency_ablation_handoff_2026-05-05.md; the full 93-row table report is outputs/markov_contextual_ablation_grid_report_20260505.md. Those ablations show that local-law / Markov-sketch supervision is still the sufficiency selector: NASSS helps as a low-weight auxiliary, learned merge and learned decoder work inside the local-law lane, and the standalone CleanUnifiedNO general f/g path has not reached exact recovery. The newest JAX regime-one-hot recovery grid adds local_law_summary_family="regime_transition_sum" and fixes the large-leaf count-extraction failure without feeding the exact Markov sketch as input.

Useful repo-readiness gates:

python scripts/run_experiment.py list --json
python scripts/audit_runtime_umbrella_coverage.py --fail-on-unclassified
python scripts/check_repo_release_hygiene.py --json

Training Quick Start

# Activate environment
source venv/bin/activate

# Start inference servers
./scripts/start_dual_servers.sh

# Run training pipeline (default task + dataset: RILE scoring on manifestos)
./scripts/run_training_pipeline.sh \
  --output-dir outputs/train_$(date +%Y%m%d_%H%M) \
  --train-samples 100 \
  --optimizer bootstrap_random_search

# Full training example (large-model-only path; GenRM/TOT flags are deprecated)
./scripts/run_training_pipeline.sh \
  --output-dir outputs/train_$(date +%Y%m%d_%H%M) \
  --train-samples 100 \
  --val-samples 30 \
  --test-samples 30 \
  --opt-model-port 8001 \
  --optimizer bootstrap_random_search \
  --optimizer-budget heavy \
  --n-iterations 2

# Init trees are filtered by prompt token budget (set with --max-init-prompt-tokens)

# Run with generic summarization task (still on manifestos by default)
./scripts/run_training_pipeline.sh \
  --task summarization \
  --output-dir outputs/summarization_test

LawStress Benchmark (MVP)

Generate synthetic local-law stress data for general information extraction (C1/C2/C3-aware) and evaluate staged. This workflow is not a STEM/math/coding problem generator.

# 1) Generate benchmark fixtures + pipeline-consumable JSONL
python scripts/generate_manifesto_lawstress.py \
  --output-dir outputs/lawstress_mvp \
  --teacher-base-url http://localhost:8000/v1 \
  --teacher-model "$TEACHER_MODEL_PATH"

# 2a) Stage 1: summarization only (small model)
python scripts/eval_manifesto_lawstress.py \
  --records outputs/lawstress_mvp/lawstress_records.jsonl \
  --output-dir outputs/lawstress_eval \
  --mode summarize_only \
  --summarizer-model qwen3.5-4b

# 2b) Stage 2: teacher scoring only (GenRM disabled)
python scripts/eval_manifesto_lawstress.py \
  --records outputs/lawstress_mvp/lawstress_records.jsonl \
  --output-dir outputs/lawstress_eval \
  --mode score_and_judge_only \
  --scorer-model "$TEACHER_MODEL_PATH" \
  --disable-genrm

Teacher Trace Bootstrap (Real Anchors)

Generate training traces from real manifesto anchors:

  • sample real manifesto text with known RILE
  • generate score-preserving English expansion with the teacher
  • produce 2-hop summaries + structured extraction traces
# Optional: launch 397B teacher on port 8000
./scripts/start_vllm.sh qwen3.5-397b-a17b-nvfp4 --port 8000 --cuda-devices 0,1,2,3

# Generate traces
python scripts/generate_manifesto_teacher_traces.py \
  --output-dir outputs/teacher_trace_bootstrap \
  --train-size 120 \
  --val-size 30 \
  --test-size 30 \
  --teacher-base-url http://localhost:8000/v1 \
  --teacher-model "$TEACHER_MODEL_PATH" \
  --scorer-base-url http://localhost:8000/v1 \
  --scorer-model "$TEACHER_MODEL_PATH"

Optional: feed generated docs back through the existing JSONL dataset path:

./scripts/run_training_pipeline.sh \
  --task manifesto_rile \
  --dataset jsonl \
  --dataset-path outputs/teacher_trace_bootstrap/benchmark_docs.jsonl \
  --train-samples 120 \
  --val-samples 30 \
  --test-samples 30

Method Stack v1 (Equal-Maturity UX)

The pipeline now exposes a consistent interface for:

  • LLM prompt optimization (Phase 2)
  • Embedding proxy heads (ridge, linear_sgd, mil_sgd) (Phase 1.25)
  • Neural operators (CTreePO, mergeable_sketch) (Phase 1.3)
  • Generator fine-tuning with LoRA/full-FT toggle (Phase 3.25/3.5)

Quick examples:

# Embedding proxy with explicit error policy
./scripts/run_training_pipeline.sh \
  --adaptive-embedding-proxy \
  --adaptive-embedding-head-method ridge \
  --embedding-proxy-fail-on-error

# Neural operators + hybrid representation auto-wire
./scripts/run_training_pipeline.sh \
  --train-neural-operators \
  --neural-operators-which both \
  --hybrid-oracle-seeded-ensemble

# Generator fine-tuning (LoRA)
./scripts/run_training_pipeline.sh \
  --train-generator \
  --generator-method dpo \
  --generator-use-lora

# One-command compare (fast-smoke default)
python scripts/run_method_compare.py --output-root outputs/method_compare_smoke
python scripts/report_method_compare.py --manifest outputs/method_compare_smoke/method_compare_manifest.json

Generate-First TreePO Stack (Unified StateTree Runner)

For the unified TreePO API (contract-driven symbolic, text, and learned-state settings), see:

  • docs/treepo_generate_first_stack.md

Runnable demos:

  • ./venv/bin/python scripts/run_treepo_stack_markov_demo.py
  • ./venv/bin/python scripts/run_treepo_stack_generate_demo.py --engine sglang --base-url http://localhost:30000 --surface generate
  • ./venv/bin/python scripts/run_paper_unified_examples.py --output-dir outputs/paper_unified_examples

Architecture

ThinkingTrees/
├── config/
│   └── settings.yaml              # Model configs, generation params
├── src/
│   ├── core/                      # Generic building blocks
│   │   ├── data_models.py         # Node, Tree, AuditResult
│   │   ├── documents.py           # DocumentSample, DocumentResult
│   │   ├── llm_client.py          # LLMClient (vLLM/OpenAI)
│   │   ├── signatures.py          # DSPy signatures (generic)
│   │   ├── strategy.py            # SummarizationStrategy protocol + registry
│   │   ├── scoring.py             # OracleScore, ScoringOracle
│   │   ├── ops_checks.py          # CheckType, CheckConfig
│   │   ├── batch_processor.py     # Async batched LLM client + request pooling
│   │   ├── batch_orchestrator.py  # Global pipelined tree batching across documents
│   │   └── output_parser.py       # Case-insensitive LLM output parsing
│   │
│   ├── tree/                      # Tree building and verification
│   │   ├── builder.py             # TreeBuilder (async-first)
│   │   ├── auditor.py             # Probabilistic verification
│   │   ├── labeled.py             # LabeledTree, LabeledDataset
│   │   └── verification.py        # TreeVerifier, OracleNodeVerifier
│   │
│   ├── training/                  # Training and optimization
│   │   ├── run_pipeline.py        # Main training entry point
│   │   ├── optimization/          # DSPy optimizers (GEPA, MIPRO, Bootstrap)
│   │   ├── preference/            # Preference learning
│   │   ├── judges/                # Pairwise comparison judges
│   │   ├── metrics/               # Evaluation metrics
│   │   └── data_sources/          # Training data sources
│   │
│   ├── tasks/                     # Task plugins
│   │   ├── base.py                # AbstractTask, ScaleDefinition
│   │   ├── registry.py            # Task discovery
│   │   ├── scoring.py             # Generic ScoringTask
│   │   ├── document_analysis.py   # Content preservation (0 to 1)
│   │   └── manifesto/             # RILE scoring building blocks
│   │
│   ├── datasets/                  # Dataset plugins
│   │   ├── base.py                # DatasetPlugin protocol
│   │   ├── manifesto.py           # Manifesto dataset
│   │   └── jsonl.py               # Generic JSONL dataset
│   │
│   ├── pipelines/                 # Task/dataset-agnostic pipelines
│   │   └── batched.py             # Batched inference pipeline
│   │
│   └── preprocessing/             # Document processing
│       └── chunker.py             # DocumentChunker
│
├── scripts/
│   ├── start_dual_servers.sh      # Start inference servers
│   ├── run_training_pipeline.sh   # Training wrapper
│   ├── generate_manifesto_teacher_traces.py  # Real-anchor teacher trace generation
│   └── stop_small_servers.sh      # Server shutdown
│
└── experiments/                   # Experiment scripts

Core Concepts

Node

The atomic unit of the summarization tree:

@dataclass
class Node:
    id: str                          # Unique identifier
    level: int                       # 0 = leaf, higher = more summarized
    raw_text_span: Optional[str]     # Original text (leaves only)
    summary: str                     # Summary at this node
    left_child: Optional[Node]       # Left subtree
    right_child: Optional[Node]      # Right subtree
    audit_result: AuditResult        # Verification status

Building Blocks Pattern

Tasks are composed from generic building blocks, not hardcoded:

from src.tasks.base import ScoringTask, ScaleDefinition
from src.core import ScaleScorer, GenericSummarizer

# Define your scale
MY_SCALE = ScaleDefinition(
    name="sentiment",
    min_value=-1.0,
    max_value=1.0,
    description="Sentiment score",
)

# Compose a task from building blocks
task = ScoringTask(
    name="sentiment",
    scale=MY_SCALE,
    rubric="Preserve sentiment indicators...",
    predictor_factory=lambda: ScaleScorer(MySentimentSignature),
)

Example using RILE building blocks from src/tasks/manifesto/:

from src.tasks.base import ScoringTask
from src.tasks.manifesto import (
    RILE_SCALE,                  # ScaleDefinition(-100, +100)
    RILE_PRESERVATION_RUBRIC,   # Domain rubric
    ManifestoDataset,           # Data loading
    RILEScorer,                 # Domain scorer
)

rile_task = ScoringTask(
    name="rile",
    scale=RILE_SCALE,
    rubric=RILE_PRESERVATION_RUBRIC,
    data_loader_factory=lambda: ManifestoDataset(),
    predictor_factory=lambda: RILEScorer(),
)

Available building blocks in src/core/:

  • ScaleScorer - Generic DSPy scorer for any bounded scale
  • PairwiseScorer - Generic pairwise comparison scorer
  • GenericSummarizer - Configurable summarization module
  • GenericMerger - Configurable merge module

Normalization and Metrics

Internal optimization uses normalized 0-1 units even when tasks have a real-world scale:

  • DSPy metrics expect higher-is-better in [0, 1]; OracleScore.score follows this.
  • Tournament preference labels are derived from normalized errors (lower is better), not raw scores.
  • Raw task values (e.g., RILE -100 to +100) are preserved for reporting and stored alongside normalized errors.
  • Tie margins are expressed in normalized units; use the task scale range to convert raw margins.

For tasks with a scale, normalization follows:

normalized_error = abs(predicted - ground_truth) / scale.range

Dataset Plugins

Datasets define where documents come from:

from src.datasets import get_dataset

dataset = get_dataset("manifesto")
samples = dataset.load_samples(limit=100)

OPS Laws (Verified by Auditor)

  1. Sufficiency (C1): oracle(summary) ≈ oracle(original)
  2. Idempotence (C2): oracle(summarize(S)) ≈ oracle(S)
  3. Merge Consistency (C3): oracle(merge) ≈ aggregate(oracle(children))

CLI Reference (src/training/run_pipeline.py)

Server Options

FlagDefaultDescription
--port8000vLLM port for summarizer/inference
--opt-model-portNoneOptional prompt-tuning LM (set to GenRM port, e.g. 8001)

Data Options

FlagDefaultDescription
--train-samples33Number of training samples
--val-samples11Number of validation samples
--test-samples11Number of test samples
--rounds3Reserved (currently unused)

Concurrency

FlagDefaultDescription
--concurrent-docs20Documents processed in parallel
--concurrent-requests200Concurrent LLM requests
--num-threads64Parallel metric evaluations

Caching

  • vLLM prefix caching (APC) is controlled by vllm.enable_prefix_caching in config/settings.yaml and is enabled by default in the server scripts.
  • DSPy response caching is enabled by default; pass --no-cache to disable it for a run.
  • Oracle memoization is used during iterative optimization via create_cached_oracle_metric (per-run in-memory cache of oracle predictions).
  • Oracle pre-caching seeds that cache with predictions for the current trainset by default; pass --no-precache to skip it.
  • Caching is independent of generation temperature; disable caching if you want maximum variability.

Optimizer

FlagDefaultDescription
--optimizerbootstrap_random_searchOptimizer (gepa, bootstrap, bootstrap_random_search, mipro, labeled_fewshot)
--optimizer-budgetheavyBudget level for GEPA/MIPRO
--max-metric-callsNoneExplicit metric-call budget (overrides budget)

Iterative Optimization

FlagDefaultDescription
--n-iterations1Iterations (0=until convergence)
--convergence-threshold0.01Early stop threshold
--convergence-patience3Early stop patience
--skip-oracle-optFalseSkip oracle/scorer optimization

Legacy GenRM/TOT (Deprecated)

FlagDefaultDescription
--enable-genrmblockedDeprecated; use local-law bootstrap (teacher scorer + proxy/GEPA), no GenRM
--start-genrm (wrapper)blockedDeprecated; wrapper exits with error
--train-comparison-moduleblockedDeprecated; wrapper exits with error

Judge Optimization (Deprecated GenRM/TOT Path)

FlagDefaultDescription
--optimize-judgeblockedDeprecated; use local-law bootstrap path
--judge-optimization-budgetlightJudge optimization budget
--use-dspy-strategyFalseReserved (currently unused)
--load-optimized-judgeNoneLoad a pre-optimized judge

Tournament of Tournaments (Deprecated)

FlagDefaultDescription
--tournament-of-tournamentsblockedDeprecated; use local-law bootstrap path
--tot-max-iterations5Max ToT iterations
--tot-convergence-threshold0.01ToT convergence threshold
--tot-convergence-patience2ToT convergence patience
--tot-samples-per-iteration50Samples per ToT iteration
--tot-judge-test-split0.2Holdout split for judge accuracy
--tot-shuffle-samplesTrueShuffle samples each iteration
--tot-random-seed42RNG seed for ToT sampling

Resume and Output

FlagDefaultDescription
--resumeFalseResume from checkpoints
--output-dirrequiredOutput directory

Inference Only

FlagDefaultDescription
--load-scorer-pathNoneLoad scorer module and skip optimization
--inference-onlyFalseRun inference only (requires scorer path)

Scale Configuration

FlagDefaultDescription
--scale-min-100.0Minimum score value
--scale-max100.0Maximum score value

Task/Dataset Selection

FlagDefaultDescription
--tasksettings.yaml defaultTask plugin (e.g., manifesto_rile, document_analysis)
--datasetsettings.yaml defaultDataset plugin (e.g., manifesto, jsonl)
--dataset-pathNonePath for file-based datasets (jsonl)

Models

ModelPortUse Case
Nemotron-30B-FP88000Default inference
Qwen3.5-397B-A17B-NVFP48001Large teacher/scorer (when launched as second server)
DiffusionGemma-26B-A4B-IT-NVFP48004Optional diffusion LLM smoke/evaluation; see docs/diffusiongemma_vllm.md

Development

# Run tests
pytest tests/ -v

# Check syntax
python3 -m py_compile src/**/*.py

# View training logs
tail -f outputs/*/training.log

References

  • AGENTS.md: Quick reference for AI agents
  • doc/architecture.md: Detailed system design
  • config/settings.yaml: All configuration options

Contributors

mlinegar

50 commits

mlinegar/ThinkingTrees

2

stars

50

commits

Python

primary language

Jul 2, 2026

updated

README

ThinkingTrees: Oracle-Preserving Summarization (OPS)

Hierarchical summarization with verifiable information preservation guarantees. Build recursive summarization trees that maintain task-critical information through probabilistic auditing and DSPy-based optimization.

ThinkingTrees And treepo

This repository now exposes two public faces:

  • ThinkingTrees: the full platform for long-document OPS pipelines, task plugins, training, and deployment.
  • treepo: a focused PyTorch package under treepo/ for method-level simulations and reports.

The treepo package is the canonical home for the new HyperLogLog streaming/cardinality work. Typical commands:

cd treepo
pip install -e ".[torch]"
treepo-bench suite cardinality-paper --out-root ../outputs/cardinality --jobs 4
treepo-bench report cardinality --output-root ../outputs/cardinality

Supported Runtime-Eval Quick Start

The official lightweight path uses the LongBench v2 fixture and mock inference, so it does not require local model servers:

source venv/bin/activate

python scripts/run_runtime_eval.py plan \
  --config config/runtime_eval/longbench_v2_smoke.yaml \
  --experiment-id longbench_v2_smoke \
  --json

python scripts/run_runtime_eval.py init \
  --config config/runtime_eval/longbench_v2_smoke.yaml \
  --output-dir outputs/runtime_eval \
  --experiment-id longbench_v2_smoke

python scripts/run_runtime_eval.py run \
  --experiment-dir outputs/runtime_eval/longbench_v2_smoke \
  --mock-llm \
  --max-problems 1

python scripts/run_runtime_eval.py aggregate \
  --experiment-dir outputs/runtime_eval/longbench_v2_smoke

Runtime methods use paper-facing roles:

  • scorer: final task scoring / answer prediction.
  • summarizer: tree-summary generation, defaulting to scorer if omitted.
  • embedder: retrieval and embedding evidence selection.
  • state_model: deterministic or learned state machinery.
  • oracle: benchmark labels or trusted evaluator provenance.

The canonical experiment artifacts are experiment_manifest.json, experiment_status.json, artifacts.json, results.jsonl, and, for active runtime paths, compact calls.jsonl traces without prompts or full contexts.

For live vLLM/SGLang/OpenAI-compatible endpoints, start servers first and set the scorer, summarizer, embedder, and state_model sections in config/runtime_eval/*.yaml.

Supported Entrypoints

New work should start from the supported experiment front doors:

  • scripts/run_runtime_eval.py for LongBench/RULER-style runtime evaluation.
  • scripts/run_experiment.py for planning, launching, collecting, and reporting canonical experiment sidecars around supported scripts.
  • scripts/run_longbench_batched_example.py and scripts/run_manifesto_batched_example.py for batched document pipelines.
  • scripts/run_method_compare.py, scripts/report_method_compare.py, and scripts/run_method_compare_lbv2.py for method grids.
  • src/training/run_pipeline.py, scripts/train_neural_operators.py, and scripts/train_ctreepo.py for training/operator paths.

The maintained inventory is config/runtime_umbrella_entrypoints.yaml. It classifies supported scripts explicitly and treats older simulations, one-off reports, and historical demos as legacy unless listed there.

Markov Contextual-Sufficiency Probe

The paper hazard-panel contextual controls are available through the ctreepo sim run contextual-sbijax CLI. The current exact-zero path is the repo-owned local-law trainer; package NASS/NASSS remain approximate baselines. After installing the optional JAX/SBI extras, use the saved MarkovOPSDataBundle panels directly:

source venv/bin/activate
python -m pip install -e ".[contextual_sbi]"

XLA_PYTHON_CLIENT_PREALLOCATE=false ctreepo sim run contextual-sbijax \
  --data-source markov \
  --load-data-bundle outputs/_bundles/markov_hazard_panels/paper_hazard_panel_v1_t128/seed_0/base_bundle.json \
  --sbijax-trainer learned_local_laws \
  --sbijax-method nasss \
  --sbijax-package-theta markov_exact_sketch \
  --sbijax-input-encoding markov_exact_sketch \
  --train-docs 1024 \
  --val-docs 256 \
  --test-docs 256 \
  --fragment-len 1 \
  --context-samples-per-doc 1 \
  --response-signature-contexts 16 \
  --response-signature-slices 8 \
  --embedding-dim 32 \
  --state-dim 25 \
  --hidden-dim 128 \
  --learning-rate 0.0003 \
  --n-iter 1000 \
  --batch-size 128 \
  --local-law-weight 1.0 \
  --local-law-leaf-weight 1.0 \
  --local-law-merge-weight 1.0 \
  --local-law-idempotence-weight 1.0 \
  --local-law-contextual-weight 1.0 \
  --seed 0 \
  --output-root outputs/contextual_sbijax_exact_zero_smoke

Use ctreepo-contextual-sbijax as the equivalent standalone console script after reinstalling editable entrypoints. See docs/contextual_sbijax_walkthrough.md and docs/markov_hazard_panels.md for the exact-sketch oracle checks and leaf-length grid commands. Paper-facing leaf-length grids should include the full t128 ladder 1, 2, 4, 8, 16, 32, 64; the small leaves are required edge/capacity checks, not optional smoke rungs.

The command above is the exact-zero recovery path. Use --sbijax-trainer package --sbijax-method nasss only when you want the pre-resolution package baseline that plateaus on its contrastive objective. See docs/contextual_sbijax_optimize_to_zero_resolved_2026-05-05.md for the resolution write-up and full leaf-grid artifacts. The current post-resolution status page is docs/markov_sim_status.md. The post-resolution ablation handoff is docs/markov_contextual_sufficiency_ablation_handoff_2026-05-05.md; the full 93-row table report is outputs/markov_contextual_ablation_grid_report_20260505.md. Those ablations show that local-law / Markov-sketch supervision is still the sufficiency selector: NASSS helps as a low-weight auxiliary, learned merge and learned decoder work inside the local-law lane, and the standalone CleanUnifiedNO general f/g path has not reached exact recovery. The newest JAX regime-one-hot recovery grid adds local_law_summary_family="regime_transition_sum" and fixes the large-leaf count-extraction failure without feeding the exact Markov sketch as input.

Useful repo-readiness gates:

python scripts/run_experiment.py list --json
python scripts/audit_runtime_umbrella_coverage.py --fail-on-unclassified
python scripts/check_repo_release_hygiene.py --json

Training Quick Start

# Activate environment
source venv/bin/activate

# Start inference servers
./scripts/start_dual_servers.sh

# Run training pipeline (default task + dataset: RILE scoring on manifestos)
./scripts/run_training_pipeline.sh \
  --output-dir outputs/train_$(date +%Y%m%d_%H%M) \
  --train-samples 100 \
  --optimizer bootstrap_random_search

# Full training example (large-model-only path; GenRM/TOT flags are deprecated)
./scripts/run_training_pipeline.sh \
  --output-dir outputs/train_$(date +%Y%m%d_%H%M) \
  --train-samples 100 \
  --val-samples 30 \
  --test-samples 30 \
  --opt-model-port 8001 \
  --optimizer bootstrap_random_search \
  --optimizer-budget heavy \
  --n-iterations 2

# Init trees are filtered by prompt token budget (set with --max-init-prompt-tokens)

# Run with generic summarization task (still on manifestos by default)
./scripts/run_training_pipeline.sh \
  --task summarization \
  --output-dir outputs/summarization_test

LawStress Benchmark (MVP)

Generate synthetic local-law stress data for general information extraction (C1/C2/C3-aware) and evaluate staged. This workflow is not a STEM/math/coding problem generator.

# 1) Generate benchmark fixtures + pipeline-consumable JSONL
python scripts/generate_manifesto_lawstress.py \
  --output-dir outputs/lawstress_mvp \
  --teacher-base-url http://localhost:8000/v1 \
  --teacher-model "$TEACHER_MODEL_PATH"

# 2a) Stage 1: summarization only (small model)
python scripts/eval_manifesto_lawstress.py \
  --records outputs/lawstress_mvp/lawstress_records.jsonl \
  --output-dir outputs/lawstress_eval \
  --mode summarize_only \
  --summarizer-model qwen3.5-4b

# 2b) Stage 2: teacher scoring only (GenRM disabled)
python scripts/eval_manifesto_lawstress.py \
  --records outputs/lawstress_mvp/lawstress_records.jsonl \
  --output-dir outputs/lawstress_eval \
  --mode score_and_judge_only \
  --scorer-model "$TEACHER_MODEL_PATH" \
  --disable-genrm

Teacher Trace Bootstrap (Real Anchors)

Generate training traces from real manifesto anchors:

  • sample real manifesto text with known RILE
  • generate score-preserving English expansion with the teacher
  • produce 2-hop summaries + structured extraction traces
# Optional: launch 397B teacher on port 8000
./scripts/start_vllm.sh qwen3.5-397b-a17b-nvfp4 --port 8000 --cuda-devices 0,1,2,3

# Generate traces
python scripts/generate_manifesto_teacher_traces.py \
  --output-dir outputs/teacher_trace_bootstrap \
  --train-size 120 \
  --val-size 30 \
  --test-size 30 \
  --teacher-base-url http://localhost:8000/v1 \
  --teacher-model "$TEACHER_MODEL_PATH" \
  --scorer-base-url http://localhost:8000/v1 \
  --scorer-model "$TEACHER_MODEL_PATH"

Optional: feed generated docs back through the existing JSONL dataset path:

./scripts/run_training_pipeline.sh \
  --task manifesto_rile \
  --dataset jsonl \
  --dataset-path outputs/teacher_trace_bootstrap/benchmark_docs.jsonl \
  --train-samples 120 \
  --val-samples 30 \
  --test-samples 30

Method Stack v1 (Equal-Maturity UX)

The pipeline now exposes a consistent interface for:

  • LLM prompt optimization (Phase 2)
  • Embedding proxy heads (ridge, linear_sgd, mil_sgd) (Phase 1.25)
  • Neural operators (CTreePO, mergeable_sketch) (Phase 1.3)
  • Generator fine-tuning with LoRA/full-FT toggle (Phase 3.25/3.5)

Quick examples:

# Embedding proxy with explicit error policy
./scripts/run_training_pipeline.sh \
  --adaptive-embedding-proxy \
  --adaptive-embedding-head-method ridge \
  --embedding-proxy-fail-on-error

# Neural operators + hybrid representation auto-wire
./scripts/run_training_pipeline.sh \
  --train-neural-operators \
  --neural-operators-which both \
  --hybrid-oracle-seeded-ensemble

# Generator fine-tuning (LoRA)
./scripts/run_training_pipeline.sh \
  --train-generator \
  --generator-method dpo \
  --generator-use-lora

# One-command compare (fast-smoke default)
python scripts/run_method_compare.py --output-root outputs/method_compare_smoke
python scripts/report_method_compare.py --manifest outputs/method_compare_smoke/method_compare_manifest.json

Generate-First TreePO Stack (Unified StateTree Runner)

For the unified TreePO API (contract-driven symbolic, text, and learned-state settings), see:

  • docs/treepo_generate_first_stack.md

Runnable demos:

  • ./venv/bin/python scripts/run_treepo_stack_markov_demo.py
  • ./venv/bin/python scripts/run_treepo_stack_generate_demo.py --engine sglang --base-url http://localhost:30000 --surface generate
  • ./venv/bin/python scripts/run_paper_unified_examples.py --output-dir outputs/paper_unified_examples

Architecture

ThinkingTrees/
├── config/
│   └── settings.yaml              # Model configs, generation params
├── src/
│   ├── core/                      # Generic building blocks
│   │   ├── data_models.py         # Node, Tree, AuditResult
│   │   ├── documents.py           # DocumentSample, DocumentResult
│   │   ├── llm_client.py          # LLMClient (vLLM/OpenAI)
│   │   ├── signatures.py          # DSPy signatures (generic)
│   │   ├── strategy.py            # SummarizationStrategy protocol + registry
│   │   ├── scoring.py             # OracleScore, ScoringOracle
│   │   ├── ops_checks.py          # CheckType, CheckConfig
│   │   ├── batch_processor.py     # Async batched LLM client + request pooling
│   │   ├── batch_orchestrator.py  # Global pipelined tree batching across documents
│   │   └── output_parser.py       # Case-insensitive LLM output parsing
│   │
│   ├── tree/                      # Tree building and verification
│   │   ├── builder.py             # TreeBuilder (async-first)
│   │   ├── auditor.py             # Probabilistic verification
│   │   ├── labeled.py             # LabeledTree, LabeledDataset
│   │   └── verification.py        # TreeVerifier, OracleNodeVerifier
│   │
│   ├── training/                  # Training and optimization
│   │   ├── run_pipeline.py        # Main training entry point
│   │   ├── optimization/          # DSPy optimizers (GEPA, MIPRO, Bootstrap)
│   │   ├── preference/            # Preference learning
│   │   ├── judges/                # Pairwise comparison judges
│   │   ├── metrics/               # Evaluation metrics
│   │   └── data_sources/          # Training data sources
│   │
│   ├── tasks/                     # Task plugins
│   │   ├── base.py                # AbstractTask, ScaleDefinition
│   │   ├── registry.py            # Task discovery
│   │   ├── scoring.py             # Generic ScoringTask
│   │   ├── document_analysis.py   # Content preservation (0 to 1)
│   │   └── manifesto/             # RILE scoring building blocks
│   │
│   ├── datasets/                  # Dataset plugins
│   │   ├── base.py                # DatasetPlugin protocol
│   │   ├── manifesto.py           # Manifesto dataset
│   │   └── jsonl.py               # Generic JSONL dataset
│   │
│   ├── pipelines/                 # Task/dataset-agnostic pipelines
│   │   └── batched.py             # Batched inference pipeline
│   │
│   └── preprocessing/             # Document processing
│       └── chunker.py             # DocumentChunker
│
├── scripts/
│   ├── start_dual_servers.sh      # Start inference servers
│   ├── run_training_pipeline.sh   # Training wrapper
│   ├── generate_manifesto_teacher_traces.py  # Real-anchor teacher trace generation
│   └── stop_small_servers.sh      # Server shutdown
│
└── experiments/                   # Experiment scripts

Core Concepts

Node

The atomic unit of the summarization tree:

@dataclass
class Node:
    id: str                          # Unique identifier
    level: int                       # 0 = leaf, higher = more summarized
    raw_text_span: Optional[str]     # Original text (leaves only)
    summary: str                     # Summary at this node
    left_child: Optional[Node]       # Left subtree
    right_child: Optional[Node]      # Right subtree
    audit_result: AuditResult        # Verification status

Building Blocks Pattern

Tasks are composed from generic building blocks, not hardcoded:

from src.tasks.base import ScoringTask, ScaleDefinition
from src.core import ScaleScorer, GenericSummarizer

# Define your scale
MY_SCALE = ScaleDefinition(
    name="sentiment",
    min_value=-1.0,
    max_value=1.0,
    description="Sentiment score",
)

# Compose a task from building blocks
task = ScoringTask(
    name="sentiment",
    scale=MY_SCALE,
    rubric="Preserve sentiment indicators...",
    predictor_factory=lambda: ScaleScorer(MySentimentSignature),
)

Example using RILE building blocks from src/tasks/manifesto/:

from src.tasks.base import ScoringTask
from src.tasks.manifesto import (
    RILE_SCALE,                  # ScaleDefinition(-100, +100)
    RILE_PRESERVATION_RUBRIC,   # Domain rubric
    ManifestoDataset,           # Data loading
    RILEScorer,                 # Domain scorer
)

rile_task = ScoringTask(
    name="rile",
    scale=RILE_SCALE,
    rubric=RILE_PRESERVATION_RUBRIC,
    data_loader_factory=lambda: ManifestoDataset(),
    predictor_factory=lambda: RILEScorer(),
)

Available building blocks in src/core/:

  • ScaleScorer - Generic DSPy scorer for any bounded scale
  • PairwiseScorer - Generic pairwise comparison scorer
  • GenericSummarizer - Configurable summarization module
  • GenericMerger - Configurable merge module

Normalization and Metrics

Internal optimization uses normalized 0-1 units even when tasks have a real-world scale:

  • DSPy metrics expect higher-is-better in [0, 1]; OracleScore.score follows this.
  • Tournament preference labels are derived from normalized errors (lower is better), not raw scores.
  • Raw task values (e.g., RILE -100 to +100) are preserved for reporting and stored alongside normalized errors.
  • Tie margins are expressed in normalized units; use the task scale range to convert raw margins.

For tasks with a scale, normalization follows:

normalized_error = abs(predicted - ground_truth) / scale.range

Dataset Plugins

Datasets define where documents come from:

from src.datasets import get_dataset

dataset = get_dataset("manifesto")
samples = dataset.load_samples(limit=100)

OPS Laws (Verified by Auditor)

  1. Sufficiency (C1): oracle(summary) ≈ oracle(original)
  2. Idempotence (C2): oracle(summarize(S)) ≈ oracle(S)
  3. Merge Consistency (C3): oracle(merge) ≈ aggregate(oracle(children))

CLI Reference (src/training/run_pipeline.py)

Server Options

FlagDefaultDescription
--port8000vLLM port for summarizer/inference
--opt-model-portNoneOptional prompt-tuning LM (set to GenRM port, e.g. 8001)

Data Options

FlagDefaultDescription
--train-samples33Number of training samples
--val-samples11Number of validation samples
--test-samples11Number of test samples
--rounds3Reserved (currently unused)

Concurrency

FlagDefaultDescription
--concurrent-docs20Documents processed in parallel
--concurrent-requests200Concurrent LLM requests
--num-threads64Parallel metric evaluations

Caching

  • vLLM prefix caching (APC) is controlled by vllm.enable_prefix_caching in config/settings.yaml and is enabled by default in the server scripts.
  • DSPy response caching is enabled by default; pass --no-cache to disable it for a run.
  • Oracle memoization is used during iterative optimization via create_cached_oracle_metric (per-run in-memory cache of oracle predictions).
  • Oracle pre-caching seeds that cache with predictions for the current trainset by default; pass --no-precache to skip it.
  • Caching is independent of generation temperature; disable caching if you want maximum variability.

Optimizer

FlagDefaultDescription
--optimizerbootstrap_random_searchOptimizer (gepa, bootstrap, bootstrap_random_search, mipro, labeled_fewshot)
--optimizer-budgetheavyBudget level for GEPA/MIPRO
--max-metric-callsNoneExplicit metric-call budget (overrides budget)

Iterative Optimization

FlagDefaultDescription
--n-iterations1Iterations (0=until convergence)
--convergence-threshold0.01Early stop threshold
--convergence-patience3Early stop patience
--skip-oracle-optFalseSkip oracle/scorer optimization

Legacy GenRM/TOT (Deprecated)

FlagDefaultDescription
--enable-genrmblockedDeprecated; use local-law bootstrap (teacher scorer + proxy/GEPA), no GenRM
--start-genrm (wrapper)blockedDeprecated; wrapper exits with error
--train-comparison-moduleblockedDeprecated; wrapper exits with error

Judge Optimization (Deprecated GenRM/TOT Path)

FlagDefaultDescription
--optimize-judgeblockedDeprecated; use local-law bootstrap path
--judge-optimization-budgetlightJudge optimization budget
--use-dspy-strategyFalseReserved (currently unused)
--load-optimized-judgeNoneLoad a pre-optimized judge

Tournament of Tournaments (Deprecated)

FlagDefaultDescription
--tournament-of-tournamentsblockedDeprecated; use local-law bootstrap path
--tot-max-iterations5Max ToT iterations
--tot-convergence-threshold0.01ToT convergence threshold
--tot-convergence-patience2ToT convergence patience
--tot-samples-per-iteration50Samples per ToT iteration
--tot-judge-test-split0.2Holdout split for judge accuracy
--tot-shuffle-samplesTrueShuffle samples each iteration
--tot-random-seed42RNG seed for ToT sampling

Resume and Output

FlagDefaultDescription
--resumeFalseResume from checkpoints
--output-dirrequiredOutput directory

Inference Only

FlagDefaultDescription
--load-scorer-pathNoneLoad scorer module and skip optimization
--inference-onlyFalseRun inference only (requires scorer path)

Scale Configuration

FlagDefaultDescription
--scale-min-100.0Minimum score value
--scale-max100.0Maximum score value

Task/Dataset Selection

FlagDefaultDescription
--tasksettings.yaml defaultTask plugin (e.g., manifesto_rile, document_analysis)
--datasetsettings.yaml defaultDataset plugin (e.g., manifesto, jsonl)
--dataset-pathNonePath for file-based datasets (jsonl)

Models

ModelPortUse Case
Nemotron-30B-FP88000Default inference
Qwen3.5-397B-A17B-NVFP48001Large teacher/scorer (when launched as second server)
DiffusionGemma-26B-A4B-IT-NVFP48004Optional diffusion LLM smoke/evaluation; see docs/diffusiongemma_vllm.md

Development

# Run tests
pytest tests/ -v

# Check syntax
python3 -m py_compile src/**/*.py

# View training logs
tail -f outputs/*/training.log

References

  • AGENTS.md: Quick reference for AI agents
  • doc/architecture.md: Detailed system design
  • config/settings.yaml: All configuration options

Contributors

mlinegar

50 commits

Languages

Python

73.3%

TeX

12.8%

Lean

11.6%

Shell

2.4%