# vizuaraailabs/nano-gpt-oss architecture

> Architecture diagram of the GitHub repository vizuaraailabs/nano-gpt-oss, 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/vizuaraailabs/nano-gpt-oss
- Repository: https://github.com/vizuaraailabs/nano-gpt-oss

## Overview

Nano-GPT-OSS is a language-model project with GPT-OSS and GPT-2 architectures, TinyStories training data preparation, and prompt-based text generation. The documented training workflow is not fully wired in the sampled entry point: `train.py` builds and saves a model and generates text, while `training/trainer.py` contains a separate optimization and evaluation loop. Edges reflect sampled calls; other README capabilities are represented where the tree supports them.

Main entry points: [`architecture/gptoss.py`](https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/architecture/gptoss.py) (GPT-OSS transformer), [`inference.py`](https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/inference.py) (Text generation), [`training/trainer.py`](https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/training/trainer.py) (Trainer and evaluator), [`training/data_loader.py`](https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/training/data_loader.py) (Dataset loader)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Data preparation"]
  node_data_loader["Dataset loader<br/>[data_loader.py]"]
  node_tokenizer["Tokenizer<br/>[tokenizer.py]"]
end

subgraph group_models["Model architectures"]
  node_gptoss["GPT-OSS transformer<br/>[gptoss.py]"]
  node_attention["Attention and RoPE<br/>[gptoss.py]"]
  node_moe["MoE SwiGLU block<br/>[gptoss.py]"]
  node_gpt2["GPT-2 transformer<br/>[gpt2.py]"]
end

subgraph group_training["Training and evaluation"]
  node_train_entry["Train entry point<br/>[train.py]"]
  node_trainer["Trainer and evaluator<br/>[trainer.py]"]
end

subgraph group_generation["Text generation"]
  node_inference["Text generation<br/>[inference.py]"]
end

node_user(("User"))
node_huggingface["Hugging Face datasets"]
node_wandb["Weights &amp; Biases"]

node_data_loader -->|"loads stories"| node_huggingface
node_data_loader -->|"encodes text"| node_tokenizer
node_train_entry -->|"imports loaders"| node_data_loader
node_train_entry -->|"constructs model"| node_gptoss
node_train_entry -->|"generates text"| node_inference
node_gptoss -->|"uses attention"| node_attention
node_gptoss -->|"uses MoE"| node_moe
node_user -->|"provides prompt"| node_inference
node_inference -->|"encodes and decodes"| node_tokenizer
node_inference -->|"requests logits"| node_gptoss
node_trainer -->|"logs metrics"| node_wandb
node_trainer -->|"generates sample"| node_inference
node_trainer -->|"trains and evaluates"| node_gptoss
node_trainer -->|"consumes loaders"| node_data_loader

click node_data_loader "https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/training/data_loader.py"
click node_tokenizer "https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/architecture/tokenizer.py"
click node_train_entry "https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/train.py"
click node_trainer "https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/training/trainer.py"
click node_gptoss "https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/architecture/gptoss.py"
click node_attention "https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/architecture/gptoss.py"
click node_moe "https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/architecture/gptoss.py"
click node_gpt2 "https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/architecture/gpt2.py"
click node_inference "https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/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_data_loader,node_tokenizer,node_user toneBlue
class node_gptoss,node_attention,node_moe,node_gpt2 toneAmber
class node_train_entry,node_trainer toneMint
class node_inference toneRose
class node_huggingface,node_wandb toneIndigo
```

## Components

### Data preparation

- [Dataset loader](https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/training/data_loader.py): `training/data_loader.py`
- [Tokenizer](https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/architecture/tokenizer.py): `architecture/tokenizer.py`

### Model architectures

- [GPT-OSS transformer](https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/architecture/gptoss.py): `architecture/gptoss.py`
- [Attention and RoPE](https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/architecture/gptoss.py): `architecture/gptoss.py`
- [MoE SwiGLU block](https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/architecture/gptoss.py): `architecture/gptoss.py`
- [GPT-2 transformer](https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/architecture/gpt2.py): `architecture/gpt2.py`

### Training and evaluation

- [Train entry point](https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/train.py): `train.py`
- [Trainer and evaluator](https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/training/trainer.py): `training/trainer.py`

### Text generation

- [Text generation](https://github.com/vizuaraailabs/nano-gpt-oss/blob/main/inference.py): `inference.py`

### Other components

- **User**
- **Hugging Face datasets**
- **Weights & Biases**

## Connections

- Dataset loader → Hugging Face datasets: loads stories
- Dataset loader → Tokenizer: encodes text
- Train entry point → Dataset loader: imports loaders
- Train entry point → GPT-OSS transformer: constructs model
- Train entry point → Text generation: generates text
- GPT-OSS transformer → Attention and RoPE: uses attention
- GPT-OSS transformer → MoE SwiGLU block: uses MoE
- User → Text generation: provides prompt
- Text generation → Tokenizer: encodes and decodes
- Text generation → GPT-OSS transformer: requests logits
- Trainer and evaluator → Weights & Biases: logs metrics
- Trainer and evaluator → Text generation: generates sample
- Trainer and evaluator → GPT-OSS transformer: trains and evaluates
- Trainer and evaluator → Dataset loader: consumes loaders