README.md

---
license: mit
language:
- en
tags:
- text-generation
- transformer
- gpt
- javascript
- nodejs
- bpe
- tinystories
datasets:
- roneneldan/TinyStories
pipeline_tag: text-generation
library_name: mni-ml-framework
inference: false
---

# mni-ml/transformer

A **12.3M-parameter** decoder-only Transformer (GPT-style) trained in **Node.js** with
[`@mni-ml/framework`](https://www.npmjs.com/package/@mni-ml/framework) on the
[TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories) corpus,
using a HuggingFace-style **ByteLevel BPE** tokenizer (vocab 4096).

Source code, training scripts, and data-prep utilities live at
[github.com/mni-ml/transformer](https://github.com/mni-ml/transformer).

> The HF inference widget is disabled for this model. It uses a custom Node.js
> runtime (`@mni-ml/framework`), not `transformers`, so the widget cannot load it.
> See [Running locally](#running-locally) below.

## Files

| File | Size | Description |
|------|------|-------------|
| `model-final.json` | ~249 MB | Final checkpoint: weights, config, and optimizer state, loaded by `@mni-ml/framework` |
| `tokenizer.json`   | ~266 KB | HuggingFace-format ByteLevel BPE tokenizer (vocab 4096, special token `<\|endoftext\|>`) |

## Architecture

Standard GPT-style decoder-only Transformer with pre-norm blocks, causal
self-attention, learnable position embeddings, and weight-tied output head.

| Hyperparameter | Value |
|----------------|-------|
| Parameters     | 12,322,816 |
| Layers (`n_layer`) | 6 |
| Attention heads (`n_head`) | 6 |
| Embedding dim (`n_embd`) | 384 |
| Head dim | 64 |
| Context window (`block_size`) | 256 tokens |
| Vocab size | 4,096 |
| Activation | GELU |
| Normalization | LayerNorm (pre-norm), ε = 1e-5 |

The full config is also embedded in `model-final.json` under the `config` key and
is read automatically by the generate scripts.

## Running locally

Because this model uses a custom JS runtime, you need **three pieces** to run
inference: the npm framework, and two source files (`src/generate.js` and
`src/bpe.js`) from the GitHub repo.

### Prerequisites

- **Node.js ≥ 22.18** (required by `@mni-ml/framework`)
- `git` (to grab the source files) and `hf` CLI (to download the weights)

### Step-by-step

```bash
# 1. Clone the source repo (needed for src/generate.js + src/bpe.js)
git clone https://github.com/mni-ml/transformer.git
cd transformer

# 2. Install the JS runtime
npm install

# 3. Download the checkpoint + tokenizer into ./out
hf download mni-ml/transformer model-final.json tokenizer.json --local-dir ./out

# 4. Generate
node src/generate.js out/model-final.json "<|endoftext|>" 400 0.9 out/tokenizer.json
```

CLI arguments to `generate.js`:

```
node src/generate.js <checkpoint> <prompt> <max_new_tokens> <temperature> <tokenizer_path>
```

> ⚠️ **The 5th argument (`tokenizer_path`) is effectively required when using
> this public checkpoint.** `model-final.json` internally records the path
> `/app/data/tokenizer.json` (the training container's path), which will not
> exist on your machine. Always pass `out/tokenizer.json` (or wherever you
> downloaded it) as the 5th arg.

Temperature `0` gives greedy decoding; values `> 0` do temperature sampling.
The prompt is encoded with the BPE tokenizer, so any UTF-8 string works;
`<\|endoftext\|>` is the only special token.

### GPU (optional)

If you install a matching `@mni-ml/framework-*` native package that exposes
`native.flashAttention`:

```bash
node src/generate_gpu.js out/model-final.json "<|endoftext|>" 400 0.9 out/tokenizer.json
```

### Quick sanity check

```bash
node src/generate.js out/model-final.json "Once upon a time" 100 0.8 out/tokenizer.json
```

Expected output style: short, simple, children's-story English (since the
training corpus is TinyStories).

## Intended use

Small research / educational model that demonstrates training a Transformer
end-to-end in JavaScript. It is fluent on short children's-story-style English
and is **not** a general-purpose chat or instruction model.

- Suitable for: short-form story continuation, JS/Node learning demos, tokenizer experiments.
- Not suitable for: factual Q&A, code generation, non-English text, long-context tasks (256-token window), safety-critical use.

## Training data

[TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories) — a synthetic
corpus of short English children's stories, originally generated by GPT-3.5 / GPT-4
and designed for training small language models. The BPE tokenizer in
`tokenizer.json` was trained on the same corpus via
`scripts/prepare_tinystories.py` in the source repo.

## Training procedure

- **Framework:** [`@mni-ml/framework`](https://www.npmjs.com/package/@mni-ml/framework) v0.3.4 (Node.js)
- **Entry point:** `src/train.js` (CPU) or `src/train_gpu.js` (GPU)
- **Objective:** next-token cross-entropy

| Hyperparameter | Value |
|----------------|-------|
| Optimizer | AdamW |
| β₁, β₂ | 0.9, 0.95 |
| Weight decay | 0.1 |
| Max grad norm | 1.0 |
| Peak LR | 3e-4 |
| Min LR | 6e-5 |
| LR schedule | Linear warmup (200 steps) → cosine decay |
| Max iterations | 7,500 |
| Batch size | 8 |
| Gradient accumulation | 4 (→ effective batch 32) |
| Dropout | 0.1 (training only) |

## Limitations and biases

- Trained only on TinyStories, so outputs mimic simple children's stories and will
  hallucinate or produce nonsense for anything outside that domain.
- TinyStories is itself GPT-generated, so any biases or artifacts of the generating
  models can propagate here.
- 256-token context window is very short.
- No RLHF, no instruction tuning, no safety alignment.
- English-only.

## License

MIT — see the source repository for details.

## Citation

```bibtex
@misc{mni-ml-transformer,
  title  = {mni-ml/transformer: a 12M-parameter Transformer trained in Node.js},
  author = {mni-ml},
  year   = {2026},
  url    = {https://github.com/mni-ml/transformer}
}

@article{eldan2023tinystories,
  title   = {TinyStories: How Small Can Language Models Be and Still Speak Coherent English?},
  author  = {Eldan, Ronen and Li, Yuanzhi},
  journal = {arXiv preprint arXiv:2305.07759},
  year    = {2023}
}
```