# internlm/sim-cot architecture

> Architecture diagram of the GitHub repository internlm/sim-cot, 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/internlm/sim-cot
- Repository: https://github.com/internlm/sim-cot

## Overview

SIM-CoT is a research framework for training and evaluating language models that reason through implicit latent tokens, with step-level supervision intended to stabilize and make those latents interpretable. The README documents two execution paths: Coconut and CODI. In each, a runner prepares data and trains or evaluates a model; during training, an auxiliary decoder supervises latent reasoning steps, while inference omits that decoder. CODI also includes a latent-probing path for decoding and inspecting latent tokens. The graph reflects the sampled source and documented workflows; unsampled evaluation wiring is left general.

Main entry points: [`Coconut/run.py`](https://github.com/internlm/sim-cot/blob/main/Coconut/run.py) (Coconut runner), [`CODI/train.py`](https://github.com/internlm/sim-cot/blob/main/CODI/train.py) (CODI trainer), [`Coconut/coconut.py`](https://github.com/internlm/sim-cot/blob/main/Coconut/coconut.py) (Coconut model), [`CODI/probe_latent_token.py`](https://github.com/internlm/sim-cot/blob/main/CODI/probe_latent_token.py) (Latent probe)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_coconut["Coconut path"]
  node_coconut_run["Coconut runner<br/>[run.py]"]
  node_coconut_data["Latent datasets<br/>[dataset.py]"]
  node_coconut_model["Coconut model<br/>[coconut.py]"]
  node_coconut_decoder["Step decoder<br/>[coconut.py]"]
end

subgraph group_codi["CODI path"]
  node_codi_train["CODI trainer<br/>[train.py]"]
  node_codi_data["Supervised examples<br/>[train.py]"]
  node_codi_model["CODI model<br/>[model.py]"]
  node_latent_probe["Latent probe"]
end

subgraph group_resources["Model and data"]
  node_base_model{{"Pretrained LM"}}
  node_reasoning_data["Reasoning data"]
end

subgraph group_outputs["Results"]
  node_checkpoints["Model checkpoints"]
  node_decoded_steps["Decoded latent steps"]
  node_evaluation["Evaluation results"]
end

node_user(("Researcher"))

node_user -->|"configures run"| node_coconut_run
node_coconut_run -->|"builds datasets"| node_coconut_data
node_coconut_run -->|"trains or evaluates"| node_coconut_model
node_coconut_data -->|"loads examples"| node_reasoning_data
node_coconut_model -->|"wraps language model"| node_base_model
node_coconut_model -->|"supervises steps"| node_coconut_decoder
node_coconut_run -->|"saves model"| node_checkpoints
node_coconut_run -->|"evaluates accuracy"| node_evaluation
node_user -->|"starts training"| node_codi_train
node_codi_train -->|"formats examples"| node_codi_data
node_codi_train -->|"trains model"| node_codi_model
node_codi_data -->|"uses examples"| node_reasoning_data
node_codi_model -->|"loads language model"| node_base_model
node_codi_train -->|"saves model"| node_checkpoints
node_user -->|"requests inspection"| node_latent_probe
node_latent_probe -->|"runs generation"| node_codi_model
node_latent_probe -->|"decodes latents"| node_decoded_steps

click node_coconut_run "https://github.com/internlm/sim-cot/blob/main/Coconut/run.py"
click node_coconut_data "https://github.com/internlm/sim-cot/blob/main/Coconut/dataset.py"
click node_coconut_model "https://github.com/internlm/sim-cot/blob/main/Coconut/coconut.py"
click node_coconut_decoder "https://github.com/internlm/sim-cot/blob/main/Coconut/coconut.py"
click node_codi_train "https://github.com/internlm/sim-cot/blob/main/CODI/train.py"
click node_codi_data "https://github.com/internlm/sim-cot/blob/main/CODI/train.py"
click node_codi_model "https://github.com/internlm/sim-cot/blob/main/CODI/src/model.py"
click node_latent_probe "https://github.com/internlm/sim-cot/blob/main/CODI/probe_latent_token.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_coconut_run,node_coconut_data,node_coconut_model,node_coconut_decoder toneBlue
class node_codi_train,node_codi_data,node_codi_model,node_latent_probe toneAmber
class node_base_model,node_reasoning_data toneMint
class node_checkpoints,node_decoded_steps,node_evaluation toneRose
class node_user toneIndigo
```

## Components

### Coconut path

- [Coconut runner](https://github.com/internlm/sim-cot/blob/main/Coconut/run.py): `Coconut/run.py`
- [Latent datasets](https://github.com/internlm/sim-cot/blob/main/Coconut/dataset.py): `Coconut/dataset.py`
- [Coconut model](https://github.com/internlm/sim-cot/blob/main/Coconut/coconut.py): `Coconut/coconut.py`
- [Step decoder](https://github.com/internlm/sim-cot/blob/main/Coconut/coconut.py): `Coconut/coconut.py`

### CODI path

- [CODI trainer](https://github.com/internlm/sim-cot/blob/main/CODI/train.py): `CODI/train.py`
- [Supervised examples](https://github.com/internlm/sim-cot/blob/main/CODI/train.py): `CODI/train.py`
- [CODI model](https://github.com/internlm/sim-cot/blob/main/CODI/src/model.py): `CODI/src/model.py`
- [Latent probe](https://github.com/internlm/sim-cot/blob/main/CODI/probe_latent_token.py): `CODI/probe_latent_token.py`

### Model and data

- **Pretrained LM**
- **Reasoning data**

### Results

- **Model checkpoints**
- **Decoded latent steps**
- **Evaluation results**

### Other components

- **Researcher**

## Connections

- Researcher → Coconut runner: configures run
- Coconut runner → Latent datasets: builds datasets
- Coconut runner → Coconut model: trains or evaluates
- Latent datasets → Reasoning data: loads examples
- Coconut model → Pretrained LM: wraps language model
- Coconut model → Step decoder: supervises steps
- Coconut runner → Model checkpoints: saves model
- Coconut runner → Evaluation results: evaluates accuracy
- Researcher → CODI trainer: starts training
- CODI trainer → Supervised examples: formats examples
- CODI trainer → CODI model: trains model
- Supervised examples → Reasoning data: uses examples
- CODI model → Pretrained LM: loads language model
- CODI trainer → Model checkpoints: saves model
- Researcher → Latent probe: requests inspection
- Latent probe → CODI model: runs generation
- Latent probe → Decoded latent steps: decodes latents