# huahenry/branchmi architecture

> Architecture diagram of the GitHub repository huahenry/branchmi, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-28). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/huahenry/branchmi
- Repository: https://github.com/huahenry/branchmi

## Overview

BranchMI Pilot is a counterfactual study tool that tests whether token-level BranchMI scores identify reasoning steps that affect a model’s final answer. The run workflow loads configured problems, generates baseline reasoning, selects checkpoints, probes alternative continuations, compares them with oracle outcomes, and analyzes score quality into a report and continue/stop recommendation. The README is TBD, so capability details come from the sampled source and configuration; some experiment internals are only partially excerpted.

Main entry points: [`branchmi_pilot/src/branchmi_pilot/experiment.py`](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/experiment.py) (Pilot experiment), [`branchmi_pilot/src/branchmi_pilot/analysis.py`](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/analysis.py) (Run analysis), [`branchmi_pilot/src/branchmi_pilot/cli.py`](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/cli.py) (Run CLI), [`branchmi_pilot/src/branchmi_pilot/backend.py`](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/backend.py) (Transformers backend)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_run["Experiment run"]
  node_cli["Run CLI<br/>[cli.py]"]
  node_config["Configuration<br/>[config.py]"]
  node_experiment["Pilot experiment<br/>[experiment.py]"]
  node_dataset["Problem loading<br/>[data.py]"]
  node_run_io["Run records<br/>[io_utils.py]"]
end

subgraph group_inference["Model inference"]
  node_prompting["Prompt construction<br/>[prompting.py]"]
  node_backend["Transformers backend<br/>[backend.py]"]
  node_baseline["Baseline generation<br/>[experiment.py]"]
  node_checkpoints["Checkpoint selection<br/>[scoring.py]"]
  node_counterfactual["Branch probes<br/>[experiment.py]"]
end

subgraph group_evaluation["Answer evaluation"]
  node_oracle["Oracle comparison<br/>[experiment.py]"]
  node_answers["Answer evaluation<br/>[answers.py]"]
  node_scores["BranchMI scoring<br/>[scoring.py]"]
end

subgraph group_analysis["Study analysis"]
  node_analysis["Run analysis<br/>[analysis.py]"]
  node_reports["Metrics and report<br/>[analysis.py]"]
  node_analysis_io["Analysis I/O<br/>[io_utils.py]"]
end

node_user(("Researcher"))
node_dataset_source[("Dataset source")]
node_model["Pretrained model"]

node_user -->|"invokes"| node_cli
node_cli -->|"loads config"| node_config
node_cli -->|"starts run"| node_experiment
node_cli -->|"requests analysis"| node_analysis
node_experiment -->|"loads problems"| node_dataset
node_dataset -->|"loads data"| node_dataset_source
node_experiment -->|"creates backend"| node_backend
node_backend -->|"loads model"| node_model
node_experiment -->|"builds prompts"| node_prompting
node_experiment -->|"generates baseline"| node_baseline
node_baseline -->|"generates tokens"| node_backend
node_experiment -->|"selects positions"| node_checkpoints
node_experiment -->|"runs probes"| node_counterfactual
node_counterfactual -->|"generates branches"| node_backend
node_experiment -->|"generates oracle"| node_oracle
node_oracle -->|"generates continuation"| node_backend
node_experiment -->|"evaluates answers"| node_answers
node_experiment -->|"computes BranchMI"| node_scores
node_experiment -->|"writes records"| node_run_io
node_cli -->|"passes config"| node_experiment
node_experiment -->|"analyzes run"| node_analysis
node_analysis -->|"reads records"| node_analysis_io
node_analysis -->|"writes summary"| node_analysis_io
node_analysis -->|"computes metrics"| node_scores
node_analysis -->|"produces report"| node_reports

click node_cli "https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/cli.py"
click node_config "https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/config.py"
click node_experiment "https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/experiment.py"
click node_dataset "https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/data.py"
click node_run_io "https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/io_utils.py"
click node_prompting "https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/prompting.py"
click node_backend "https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/backend.py"
click node_baseline "https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/experiment.py"
click node_checkpoints "https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/scoring.py"
click node_counterfactual "https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/experiment.py"
click node_oracle "https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/experiment.py"
click node_answers "https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/answers.py"
click node_scores "https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/scoring.py"
click node_analysis "https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/analysis.py"
click node_reports "https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/analysis.py"
click node_analysis_io "https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/io_utils.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_cli,node_config,node_experiment,node_dataset,node_run_io toneBlue
class node_prompting,node_backend,node_baseline,node_checkpoints,node_counterfactual,node_dataset_source toneAmber
class node_oracle,node_answers,node_scores toneMint
class node_analysis,node_reports,node_analysis_io toneRose
class node_user,node_model toneIndigo
```

## Components

### Experiment run

- [Run CLI](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/cli.py): `branchmi_pilot/src/branchmi_pilot/cli.py`
- [Configuration](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/config.py): `branchmi_pilot/src/branchmi_pilot/config.py`
- [Pilot experiment](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/experiment.py): `branchmi_pilot/src/branchmi_pilot/experiment.py`
- [Problem loading](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/data.py): `branchmi_pilot/src/branchmi_pilot/data.py`
- [Run records](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/io_utils.py): `branchmi_pilot/src/branchmi_pilot/io_utils.py`

### Model inference

- [Prompt construction](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/prompting.py): `branchmi_pilot/src/branchmi_pilot/prompting.py`
- [Transformers backend](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/backend.py): `branchmi_pilot/src/branchmi_pilot/backend.py`
- [Baseline generation](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/experiment.py): `branchmi_pilot/src/branchmi_pilot/experiment.py`
- [Checkpoint selection](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/scoring.py): `branchmi_pilot/src/branchmi_pilot/scoring.py`
- [Branch probes](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/experiment.py): `branchmi_pilot/src/branchmi_pilot/experiment.py`

### Answer evaluation

- [Oracle comparison](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/experiment.py): `branchmi_pilot/src/branchmi_pilot/experiment.py`
- [Answer evaluation](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/answers.py): `branchmi_pilot/src/branchmi_pilot/answers.py`
- [BranchMI scoring](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/scoring.py): `branchmi_pilot/src/branchmi_pilot/scoring.py`

### Study analysis

- [Run analysis](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/analysis.py): `branchmi_pilot/src/branchmi_pilot/analysis.py`
- [Metrics and report](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/analysis.py): `branchmi_pilot/src/branchmi_pilot/analysis.py`
- [Analysis I/O](https://github.com/huahenry/branchmi/blob/main/branchmi_pilot/src/branchmi_pilot/io_utils.py): `branchmi_pilot/src/branchmi_pilot/io_utils.py`

### Other components

- **Researcher**
- **Dataset source**
- **Pretrained model**

## Connections

- Researcher → Run CLI: invokes
- Run CLI → Configuration: loads config
- Run CLI → Pilot experiment: starts run
- Run CLI → Run analysis: requests analysis
- Pilot experiment → Problem loading: loads problems
- Problem loading → Dataset source: loads data
- Pilot experiment → Transformers backend: creates backend
- Transformers backend → Pretrained model: loads model
- Pilot experiment → Prompt construction: builds prompts
- Pilot experiment → Baseline generation: generates baseline
- Baseline generation → Transformers backend: generates tokens
- Pilot experiment → Checkpoint selection: selects positions
- Pilot experiment → Branch probes: runs probes
- Branch probes → Transformers backend: generates branches
- Pilot experiment → Oracle comparison: generates oracle
- Oracle comparison → Transformers backend: generates continuation
- Pilot experiment → Answer evaluation: evaluates answers
- Pilot experiment → BranchMI scoring: computes BranchMI
- Pilot experiment → Run records: writes records
- Run CLI → Pilot experiment: passes config
- Pilot experiment → Run analysis: analyzes run
- Run analysis → Analysis I/O: reads records
- Run analysis → Analysis I/O: writes summary
- Run analysis → BranchMI scoring: computes metrics
- Run analysis → Metrics and report: produces report