# kitft/natural_language_autoencoders architecture

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

- Interactive diagram: https://gitdiagram.com/kitft/natural_language_autoencoders
- Repository: https://github.com/kitft/natural_language_autoencoders

## Overview

NLA trains paired activation verbalizer (AV) and reconstructor (AR) language models: data generation turns residual-stream vectors into explanation and reconstruction datasets, SFT initializes the models, and GRPO RL improves AV generations using AR reconstruction error. During rollout, the trainer injects activation embeddings and sends them to SGLang; training and scoring are integrated through Miles. `nla_inference.py` provides a separate checkpoint-based generation and AR-scoring path. The README documents the broader training capabilities; sampled excerpts confirm selected orchestration, rollout, training, storage, and inference behavior, so uncertain internal wiring is omitted.

Main entry points: [`nla/storage.py`](https://github.com/kitft/natural_language_autoencoders/blob/main/nla/storage.py) (Parquet storage), [`nla/datagen/run_pipeline.py`](https://github.com/kitft/natural_language_autoencoders/blob/main/nla/datagen/run_pipeline.py) (Data pipeline), [`nla_inference.py`](https://github.com/kitft/natural_language_autoencoders/blob/main/nla_inference.py) (NLA client), [`nla/train_actor.py`](https://github.com/kitft/natural_language_autoencoders/blob/main/nla/train_actor.py) (FSDP actor)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Data generation"]
  node_extract["Activation extraction<br/>[stage0_extract.py]"]
  node_split["Dataset splitting<br/>[stage1_split.py]"]
  node_explain["API explanations"]
  node_build["Dataset assembly<br/>[stage3_build.py]"]
  node_pipeline["Data pipeline<br/>[run_pipeline.py]"]
end

subgraph group_training["Training"]
  node_source["Training data source<br/>[data_source.py]"]
  node_sft["SFT rollout<br/>[sft_actor.py]"]
  node_fsdp["FSDP actor<br/>[train_actor.py]"]
  node_megatron["Megatron actor<br/>[train_actor.py]"]
  node_rollout["NLA generation<br/>[nla_generate.py]"]
  node_reward["Reconstruction reward<br/>[reward.py]"]
  node_loss["Training losses<br/>[loss.py]"]
end

subgraph group_inference["Inference"]
  node_checkpoint["AV / AR checkpoint"]
  node_client["NLA client<br/>[nla_inference.py]"]
  node_critic["AR scoring<br/>[nla_inference.py]"]
end

subgraph group_integration["External services"]
  node_corpus["Text corpus"]
  node_miles["Miles trainer"]
  node_sglang["SGLang serving"]
end

subgraph group_contracts["Shared contracts"]
  node_store[("Parquet storage<br/>[storage.py]")]
  node_schema["Metadata schema<br/>[schema.py]"]
end

node_user(("Researcher"))

node_user -->|"configures"| node_pipeline
node_corpus -->|"provides text"| node_extract
node_pipeline -->|"runs stage"| node_extract
node_pipeline -->|"runs stage"| node_split
node_pipeline -->|"runs stage"| node_explain
node_pipeline -->|"runs stage"| node_build
node_extract -->|"writes activations"| node_store
node_split -->|"writes splits"| node_store
node_explain -->|"writes explanations"| node_store
node_build -->|"writes datasets"| node_store
node_store -->|"reads datasets"| node_source
node_schema -->|"defines sidecars"| node_store
node_source -->|"supplies batches"| node_miles
node_miles -->|"trains actor"| node_fsdp
node_miles -.->|"trains actor"| node_megatron
node_miles -->|"uses rollout"| node_sft
node_miles -->|"dispatches generation"| node_rollout
node_rollout -->|"sends embeddings"| node_sglang
node_miles -->|"computes reward"| node_reward
node_fsdp -->|"uses loss"| node_loss
node_megatron -.->|"uses loss"| node_loss
node_fsdp -->|"saves models"| node_checkpoint
node_megatron -.->|"saves models"| node_checkpoint
node_user -->|"submits vector"| node_client
node_client -->|"loads AV"| node_checkpoint
node_client -->|"requests generation"| node_sglang
node_sglang -->|"returns text"| node_client
node_user -->|"submits vector and text"| node_critic
node_critic -->|"loads AR"| node_checkpoint
node_critic -->|"returns score"| node_user

click node_extract "https://github.com/kitft/natural_language_autoencoders/blob/main/nla/datagen/stage0_extract.py"
click node_split "https://github.com/kitft/natural_language_autoencoders/blob/main/nla/datagen/stage1_split.py"
click node_explain "https://github.com/kitft/natural_language_autoencoders/blob/main/nla/datagen/stage2_api_explain.py"
click node_build "https://github.com/kitft/natural_language_autoencoders/blob/main/nla/datagen/stage3_build.py"
click node_pipeline "https://github.com/kitft/natural_language_autoencoders/blob/main/nla/datagen/run_pipeline.py"
click node_store "https://github.com/kitft/natural_language_autoencoders/blob/main/nla/storage.py"
click node_schema "https://github.com/kitft/natural_language_autoencoders/blob/main/nla/schema.py"
click node_source "https://github.com/kitft/natural_language_autoencoders/blob/main/nla/data_source.py"
click node_sft "https://github.com/kitft/natural_language_autoencoders/blob/main/nla/rollout/sft_actor.py"
click node_fsdp "https://github.com/kitft/natural_language_autoencoders/blob/main/nla/train_actor.py"
click node_megatron "https://github.com/kitft/natural_language_autoencoders/blob/main/nla/megatron/train_actor.py"
click node_rollout "https://github.com/kitft/natural_language_autoencoders/blob/main/nla/rollout/nla_generate.py"
click node_reward "https://github.com/kitft/natural_language_autoencoders/blob/main/nla/reward.py"
click node_loss "https://github.com/kitft/natural_language_autoencoders/blob/main/nla/loss.py"
click node_client "https://github.com/kitft/natural_language_autoencoders/blob/main/nla_inference.py"
click node_critic "https://github.com/kitft/natural_language_autoencoders/blob/main/nla_inference.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_extract,node_split,node_explain,node_build,node_pipeline toneBlue
class node_source,node_sft,node_fsdp,node_megatron,node_rollout,node_reward,node_loss toneAmber
class node_checkpoint,node_client,node_critic toneMint
class node_corpus,node_miles,node_sglang toneRose
class node_store,node_schema,node_user toneIndigo
```

## Components

### Data generation

- [Activation extraction](https://github.com/kitft/natural_language_autoencoders/blob/main/nla/datagen/stage0_extract.py): `nla/datagen/stage0_extract.py`
- [Dataset splitting](https://github.com/kitft/natural_language_autoencoders/blob/main/nla/datagen/stage1_split.py): `nla/datagen/stage1_split.py`
- [API explanations](https://github.com/kitft/natural_language_autoencoders/blob/main/nla/datagen/stage2_api_explain.py): `nla/datagen/stage2_api_explain.py`
- [Dataset assembly](https://github.com/kitft/natural_language_autoencoders/blob/main/nla/datagen/stage3_build.py): `nla/datagen/stage3_build.py`
- [Data pipeline](https://github.com/kitft/natural_language_autoencoders/blob/main/nla/datagen/run_pipeline.py): `nla/datagen/run_pipeline.py`

### Training

- [Training data source](https://github.com/kitft/natural_language_autoencoders/blob/main/nla/data_source.py): `nla/data_source.py`
- [SFT rollout](https://github.com/kitft/natural_language_autoencoders/blob/main/nla/rollout/sft_actor.py): `nla/rollout/sft_actor.py`
- [FSDP actor](https://github.com/kitft/natural_language_autoencoders/blob/main/nla/train_actor.py): `nla/train_actor.py`
- [Megatron actor](https://github.com/kitft/natural_language_autoencoders/blob/main/nla/megatron/train_actor.py): `nla/megatron/train_actor.py`
- [NLA generation](https://github.com/kitft/natural_language_autoencoders/blob/main/nla/rollout/nla_generate.py): `nla/rollout/nla_generate.py`
- [Reconstruction reward](https://github.com/kitft/natural_language_autoencoders/blob/main/nla/reward.py): `nla/reward.py`
- [Training losses](https://github.com/kitft/natural_language_autoencoders/blob/main/nla/loss.py): `nla/loss.py`

### Inference

- **AV / AR checkpoint**
- [NLA client](https://github.com/kitft/natural_language_autoencoders/blob/main/nla_inference.py): `nla_inference.py`
- [AR scoring](https://github.com/kitft/natural_language_autoencoders/blob/main/nla_inference.py): `nla_inference.py`

### External services

- **Text corpus**
- **Miles trainer**
- **SGLang serving**

### Shared contracts

- [Parquet storage](https://github.com/kitft/natural_language_autoencoders/blob/main/nla/storage.py): `nla/storage.py`
- [Metadata schema](https://github.com/kitft/natural_language_autoencoders/blob/main/nla/schema.py): `nla/schema.py`

### Other components

- **Researcher**

## Connections

- Researcher → Data pipeline: configures
- Text corpus → Activation extraction: provides text
- Data pipeline → Activation extraction: runs stage
- Data pipeline → Dataset splitting: runs stage
- Data pipeline → API explanations: runs stage
- Data pipeline → Dataset assembly: runs stage
- Activation extraction → Parquet storage: writes activations
- Dataset splitting → Parquet storage: writes splits
- API explanations → Parquet storage: writes explanations
- Dataset assembly → Parquet storage: writes datasets
- Parquet storage → Training data source: reads datasets
- Metadata schema → Parquet storage: defines sidecars
- Training data source → Miles trainer: supplies batches
- Miles trainer → FSDP actor: trains actor
- Miles trainer → Megatron actor: trains actor
- Miles trainer → SFT rollout: uses rollout
- Miles trainer → NLA generation: dispatches generation
- NLA generation → SGLang serving: sends embeddings
- Miles trainer → Reconstruction reward: computes reward
- FSDP actor → Training losses: uses loss
- Megatron actor → Training losses: uses loss
- FSDP actor → AV / AR checkpoint: saves models
- Megatron actor → AV / AR checkpoint: saves models
- Researcher → NLA client: submits vector
- NLA client → AV / AR checkpoint: loads AV
- NLA client → SGLang serving: requests generation
- SGLang serving → NLA client: returns text
- Researcher → AR scoring: submits vector and text
- AR scoring → AV / AR checkpoint: loads AR
- AR scoring → Researcher: returns score