# angelazzz-611/reasoning_models_probing architecture
> Architecture diagram of the GitHub repository angelazzz-611/reasoning_models_probing, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-29). It is an AI-made overview: check details against the source.
- Interactive diagram: https://gitdiagram.com/angelazzz-611/reasoning_models_probing
- Repository: https://github.com/angelazzz-611/reasoning_models_probing
## Overview
1. **What this repo does**
- This is an offline research pipeline for probing reasoning models’ hidden states to predict whether an intermediate reasoning step is correct.
- The system is not a serving stack; it is a data-generation, labeling, representation-extraction, training, and evaluation workflow.
2. **Main pipeline stages**
- The end-to-end flow is: generate model rollouts → segment rollouts into reasoning chunks → label chunk answers for correctness → extract hidden-state representations → train a probe classifier → evaluate on held-out data.
- Each stage produces a new artifact on disk, and later stages consume those artifacts rather than re-running earlier model inference.
3. **Raw reasoning data layer**
- `initial_cot/` contains pre-generated chain-of-thought rollout files in JSONL format for multiple model families and datasets.
- These files are the starting point for the pipeline and encode the reasoning traces used to derive training examples for the probe.
4. **Reasoning segmentation subsystem**
- `src/get_reasoning_chunks.py` and the shell wrapper `get_reasoning_chunks.sh` break raw CoT traces into smaller reasoning chunks.
- The README indicates this step uses spaCy (`en_core_web_sm`) for text segmentation, and it supports multiprocessing because this can be CPU-heavy.
- The output is a merged segmented dataset that becomes the unit of labeling and representation extraction.
5. **Intermediate answer labeling subsystem**
- `src/label_answer_correctness.py` labels each chunk with whether its intermediate answer is correct.
- This stage uses the Gemini API as an external LLM service for answer extraction and correctness labeling, and can merge chunks that do not contain an answer into later chunks.
- The output is a labeled JSONL dataset that aligns reasoning chunks with correctness targets.
6. **Hidden-state extraction subsystem**
- `src/get_representation.py` runs the target reasoning model on the labeled chunks and saves hidden-state representations.
- The README shows the model is loaded from a local path, and extraction is batched and sharded into multiple files via `file_id` and `file_size` to manage storage and shuffle quality.
- This is the boundary between the base reasoning model and the downstream probe: the probe never sees text alone, only chunk representations.
7. **Probe training subsystem**
- `src/train_predictor_with_class_weights.py` trains the self-verification probe, which is a classifier predicting correctness from hidden states.
- The training code supports class-weight style imbalance handling via `alpha_imbalance_penalty`, plus standard hyperparameters such as learning rate, weight decay, threshold, and optional hidden size.
- `train_probe.sh` appears to orchestrate grid search across probe configurations, storing checkpoints and experiment metadata under `grid_search/`.
8. **Probe evaluation subsystem**
- `src/test_predictor_with_class_weights.py` evaluates trained probes on held-out embeddings.
- It can test a single checkpoint or select top-k probes from grid-search results using a validation metric such as `best_val_acc`.
- Test outputs are stored separately from training artifacts, keeping evaluation isolated from model selection.
9. **Shared training/data utilities**
- `src/dataloader.py`, `src/probe_model.py`, `src/compute_metrics.py`, and `src/utils.py` are the supporting runtime layer for loading embeddings/labels, defining the probe network, computing metrics, and sharing helper logic.
- These modules are the internal boundary between data format concerns and model-training concerns.
10. **Math-specific grading support**
- `math_eval_tools/` is a bundled utility package for math answer normalization and grading.
- It likely supports correctness labeling and/or evaluation for math-style tasks where answers need canonicalization before comparison.
11. **Orchestration and automation**
- The repo relies heavily on shell scripts such as `generate_reasoning.sh`, `get_reasoning_chunks.sh`, `label_answer.sh`, `train_probe.sh`, and `test_probe.sh`.
- These scripts define the operational sequencing of the pipeline and make the system reproducible across datasets and model variants.
12. **Key storage boundaries**
- The pipeline is intentionally file-based: raw rollouts in `initial_cot/`, segmented/merged chunks in processed directories, labeled examples in `labeled_cot/`, model embeddings in `model_embeds/`, trained checkpoints in `grid_search/`, and evaluation outputs in `test_result/`.
- That separation is important for architecture diagrams because each directory marks a distinct artifact boundary and a distinct compute stage.
13. **External dependencies that materially shape the architecture**
- Hugging Face-style local model loading is implied for rollout generation and hidden-state extraction.
- spaCy is required for chunk segmentation.
- Gemini API is required for intermediate answer extraction/labeling.
- PyTorch is the likely training runtime for the probe classifier and checkpointing workflow.
14. **Repo-level architectural takeaway**
- The repository implements a staged offline ML pipeline, not a monolithic application.
- The central architectural idea is to convert reasoning traces into labeled hidden-state datasets, then learn a lightweight probe that estimates correctness from internal activations rather than from the final answer text alone.
Main entry points: [`src/test_predictor_with_class_weights.py`](https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/test_predictor_with_class_weights.py) (Evaluate probe), [`src/train_predictor_with_class_weights.py`](https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/train_predictor_with_class_weights.py) (Train probe), [`src/label_answer_correctness.py`](https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/label_answer_correctness.py) (Label correctness), [`src/get_reasoning_chunks.py`](https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/get_reasoning_chunks.py) (Chunk reasoning)
## Diagram (Mermaid)
```mermaid
flowchart TD
subgraph group_group_pipeline["Offline pipeline"]
node_node_generate_reasoning["Generate rollouts
script"]
node_node_chunking["Chunk reasoning
segmentation"]
node_node_labeling["Label correctness
LLM labeling"]
node_node_repr["Extract hidden states
representation extraction"]
node_node_train["Train probe
classifier training"]
node_node_test["Evaluate probe
evaluation"]
end
subgraph group_group_data["Artifact stores"]
node_node_initial_cot["Raw CoT
dataset"]
node_node_segmented["Segmented data
intermediate dataset"]
node_node_labeled["Labeled chunks
training dataset"]
node_node_embeds[("Model embeds
embedding shards")]
node_node_grid[("Grid search
experiments")]
node_node_results["Test results
metrics output"]
node_node_math_eval["Math grading
grading library"]
end
subgraph group_group_runtime["Training runtime"]
node_node_dataloader["Data loader
runtime util
[dataloader.py]"]
node_node_probe_model["Probe model
classifier
[probe_model.py]"]
node_node_metrics["Metrics
runtime util
[compute_metrics.py]"]
node_node_utils["Shared utils
helpers
[utils.py]"]
end
subgraph group_group_scripts["Orchestration"]
node_node_shells{{"Pipeline scripts
shell orchestration"}}
end
node_node_generate_reasoning -->|"writes"| node_node_initial_cot
node_node_initial_cot -->|"input"| node_node_chunking
node_node_chunking -->|"produces"| node_node_segmented
node_node_segmented -->|"input"| node_node_labeling
node_node_labeling -->|"produces"| node_node_labeled
node_node_labeled -->|"input"| node_node_repr
node_node_repr -->|"produces"| node_node_embeds
node_node_embeds -->|"input"| node_node_train
node_node_train -->|"writes"| node_node_grid
node_node_grid -->|"selects"| node_node_test
node_node_embeds -->|"input"| node_node_test
node_node_test -->|"writes"| node_node_results
node_node_dataloader -->|"supports"| node_node_train
node_node_dataloader -->|"supports"| node_node_test
node_node_probe_model -->|"implements"| node_node_train
node_node_probe_model -->|"implements"| node_node_test
node_node_metrics -->|"supports"| node_node_test
node_node_utils -->|"supports"| node_node_train
node_node_utils -->|"supports"| node_node_test
node_node_math_eval -->|"supports"| node_node_labeling
node_node_math_eval -->|"supports"| node_node_test
node_node_shells -->|"orchestrates"| node_node_generate_reasoning
node_node_shells -->|"orchestrates"| node_node_chunking
node_node_shells -->|"orchestrates"| node_node_labeling
node_node_shells -->|"orchestrates"| node_node_repr
node_node_shells -->|"orchestrates"| node_node_train
node_node_shells -->|"orchestrates"| node_node_test
click node_node_generate_reasoning "https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/generate_reasoning.py"
click node_node_initial_cot "https://github.com/angelazzz-611/reasoning_models_probing/tree/main/initial_cot"
click node_node_chunking "https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/get_reasoning_chunks.py"
click node_node_labeling "https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/label_answer_correctness.py"
click node_node_repr "https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/get_representation.py"
click node_node_train "https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/train_predictor_with_class_weights.py"
click node_node_test "https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/test_predictor_with_class_weights.py"
click node_node_dataloader "https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/dataloader.py"
click node_node_probe_model "https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/probe_model.py"
click node_node_metrics "https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/compute_metrics.py"
click node_node_utils "https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/utils.py"
click node_node_math_eval "https://github.com/angelazzz-611/reasoning_models_probing/tree/main/math_eval_tools"
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_node_generate_reasoning,node_node_chunking,node_node_labeling,node_node_repr,node_node_train,node_node_test toneBlue
class node_node_initial_cot,node_node_segmented,node_node_labeled,node_node_embeds,node_node_grid,node_node_results,node_node_math_eval toneAmber
class node_node_dataloader,node_node_probe_model,node_node_metrics,node_node_utils toneMint
class node_node_shells toneRose
```
## Components
### Offline pipeline
File-based stages that turn reasoning rollouts into probe-ready artifacts.
- [Generate rollouts](https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/generate_reasoning.py): script · Creates raw chain-of-thought rollouts from target reasoning models. · `src/generate_reasoning.py`
- [Chunk reasoning](https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/get_reasoning_chunks.py): segmentation · Splits raw reasoning traces into smaller reasoning chunks using spaCy and multiprocessing. · `src/get_reasoning_chunks.py`
- [Label correctness](https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/label_answer_correctness.py): LLM labeling · Extracts intermediate answers and assigns correctness labels with Gemini, optionally merging answerless chunks. · `src/label_answer_correctness.py`
- [Extract hidden states](https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/get_representation.py): representation extraction · Runs the reasoning model on labeled chunks and saves hidden-state embeddings in shards. · `src/get_representation.py`
- [Train probe](https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/train_predictor_with_class_weights.py): classifier training · Trains the self-verification probe to predict correctness from hidden states. · `src/train_predictor_with_class_weights.py`
- [Evaluate probe](https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/test_predictor_with_class_weights.py): evaluation · Scores trained probes on held-out embeddings and selects best checkpoints. · `src/test_predictor_with_class_weights.py`
### Artifact stores
On-disk datasets and checkpoints that separate each stage of the workflow.
- [Raw CoT](https://github.com/angelazzz-611/reasoning_models_probing/tree/main/initial_cot): dataset · Pre-generated rollout JSONL files used as the starting point for the pipeline. · `initial_cot`
- **Segmented data**: intermediate dataset · Merged chunk-level reasoning dataset produced from raw rollouts.
- **Labeled chunks**: training dataset · Chunk-level JSONL dataset with correctness targets for probe learning.
- **Model embeds**: embedding shards · Stored hidden-state representations used as probe inputs.
- **Grid search**: experiments · Stores probe checkpoints and tuning metadata for training runs.
- **Test results**: metrics output · Separated evaluation outputs for held-out probe performance.
- [Math grading](https://github.com/angelazzz-611/reasoning_models_probing/tree/main/math_eval_tools): grading library · Normalizes and grades math answers for correctness-focused tasks. · `math_eval_tools`
### Training runtime
Shared code used by probe training and evaluation.
- [Data loader](https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/dataloader.py): runtime util · Loads embeddings, labels, and split artifacts for training and evaluation. · `src/dataloader.py`
- [Probe model](https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/probe_model.py): classifier · Defines the lightweight correctness predictor used by the probe pipeline. · `src/probe_model.py`
- [Metrics](https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/compute_metrics.py): runtime util · Computes evaluation metrics and validation summaries. · `src/compute_metrics.py`
- [Shared utils](https://github.com/angelazzz-611/reasoning_models_probing/blob/main/src/utils.py): helpers · Common helper logic used across training and evaluation code. · `src/utils.py`
### Orchestration
Shell entrypoints that sequence the offline workflow.
- **Pipeline scripts**: shell orchestration · Shell wrappers that drive generation, chunking, labeling, training, and testing.
## Connections
- Generate rollouts → Raw CoT: writes · Saves rollout JSONL artifacts for downstream processing.
- Raw CoT → Chunk reasoning: input · Raw reasoning traces are segmented into chunks.
- Chunk reasoning → Segmented data: produces · Creates merged segmented reasoning datasets.
- Segmented data → Label correctness: input · Chunked reasoning becomes the labeling target.
- Label correctness → Labeled chunks: produces · Outputs labeled intermediate-answer examples.
- Labeled chunks → Extract hidden states: input · Labeled chunks are fed to the reasoning model for embedding extraction.
- Extract hidden states → Model embeds: produces · Writes sharded hidden-state representations to disk.
- Model embeds → Train probe: input · Embeddings are the training features for the probe.
- Train probe → Grid search: writes · Stores checkpoints and grid-search artifacts.
- Grid search → Evaluate probe: selects · Evaluation can load top-ranked checkpoints from search results.
- Model embeds → Evaluate probe: input · Held-out embeddings are used for evaluation.
- Evaluate probe → Test results: writes · Saves final evaluation outputs separately from training artifacts.
- Data loader → Train probe: supports · Loads training splits and embedding datasets.
- Data loader → Evaluate probe: supports · Loads evaluation splits and checkpoints.
- Probe model → Train probe: implements · Defines the classifier architecture used during training.
- Probe model → Evaluate probe: implements · Shared model definition is reused for checkpoint loading.
- Metrics → Evaluate probe: supports · Computes accuracy and validation summaries.
- Shared utils → Train probe: supports · Shared training helpers and configuration logic.
- Shared utils → Evaluate probe: supports · Shared evaluation helpers and artifact handling.
- Math grading → Label correctness: supports · Canonicalizes and grades math answers for labeling.
- Math grading → Evaluate probe: supports · May support math-style evaluation normalization.
- Pipeline scripts → Generate rollouts: orchestrates · Shell entrypoints drive the staged workflow.
- Pipeline scripts → Chunk reasoning: orchestrates · Shell entrypoints drive the staged workflow.
- Pipeline scripts → Label correctness: orchestrates · Shell entrypoints drive the staged workflow.
- Pipeline scripts → Extract hidden states: orchestrates · Shell entrypoints drive the staged workflow.
- Pipeline scripts → Train probe: orchestrates · Shell entrypoints drive the staged workflow.
- Pipeline scripts → Evaluate probe: orchestrates · Shell entrypoints drive the staged workflow.