README.md

---
license: mit
language:
- en
tags:
- mlx
- nanoLLM
- causal-lm
- text-generation
datasets:
- chonkie-ai/wikipedia-100k
pipeline_tag: text-generation
library_name: mlx
widget:
- text: "Cats are"
  example_title: "Cats"
---

# NanoLLM 20.2M Wiki

`nanollm_wiki_20.2M` is a lightweight, decoder-only transformer model designed for educational purposes, fast training experimentation, and local generation on Apple Silicon using Apple's **MLX** framework. 

The model is trained on a 100k subset of English Wikipedia articles (`chonkie-ai/wikipedia-100k`) to demonstrate basic autoregressive text generation capabilities with minimal compute requirements.

## Model Details

- **Developed by:** [samairtimer](https://huggingface.co/samairtimer)
- **Model Type:** Decoder-only Transformer (Custom NanoLLM Architecture)
- **Language(s):** English
- **License:** MIT
- **Framework:** Apple MLX
- **Tokenizer:** GPT-2 (via `tiktoken`, vocab size 50,257)

## Model Architecture & Hyperparameters

Unlike standard Llama or GPT architectures, this model is a **highly simplified attention-only transformer** that omits the typical Feed-Forward Neural Network (FFN) sub-blocks and layer normalization to keep training extremely fast and lightweight.

| Parameter | Value | Detail |
|---|---|---|
| **`vocab_size`** | 50,257 | GPT-2 Tiktoken Vocabulary |
| **`maxlen`** (Context Length) | 128 | Maximum context window |
| **`embed_dim`** | 192 | Hidden size dimension |
| **`num_transformer_blocks`** | 6 | Number of decoder layers |
| **`num_heads`** | 6 | Multi-head attention heads (32-dim heads) |
| **`feed_forward_dim`** | 512 | Configured but omitted in runtime computation |

### Parameter Count Breakdown

The model contains exactly **20,208,000 parameters (~20.2M)**. 

$$
\text{Total Parameters} = \text{Embeddings} + \text{Transformer Blocks} + \text{Output Projection}
$$

1. **Embedding Layer:**
   - Token Embeddings: $50,257 \times 192 = 9,649,344$
   - Positional Embeddings: $128 \times 192 = 24,576$
   - *Subtotal:* **9,673,920 parameters**

2. **Transformer Blocks (6 Layers):**
   - Each layer contains only a Multi-head Attention block (no biases, no feed-forward projection):
     - Query Projection: $192 \times 192 = 36,864$
     - Key Projection: $192 \times 192 = 36,864$
     - Value Projection: $192 \times 192 = 36,864$
     - Output Projection: $192 \times 192 = 36,864$
     - *Subtotal per layer:* $147,456$ parameters
   - *Total for 6 layers:* $147,456 \times 6 = \mathbf{884,736}$ **parameters**

3. **Output Layer (LM Head):**
   - Untied weights, linear layer with `bias=False`: $192 \times 50,257 = \mathbf{9,649,344}$ **parameters**

$$\text{Grand Total} = 9,673,920 + 884,736 + 9,649,344 = 20,208,000 \text{ parameters}$$

---

## Training Recipe

### Dataset
- **Name:** `chonkie-ai/wikipedia-100k`
- **Size:** 100,000 high-quality parsed English Wikipedia samples (~449 MB of raw text).
- **Format:** Combined articles delimited by the `<|endoftext|>` token.

### Training Configuration
- **Optimizer:** AdamW (`learning_rate=3e-4`)
- **Batch Size:** 32
- **Sequence Length:** 128 tokens
- **Epochs:** 5 epochs (~15,600 steps total)
- **Precision:** float32 (saving to `.safetensors` size of ~80.8 MB)
- **Hardware:** Apple Silicon Mac GPU (accelerated via MLX Metal backend)

### Loss Curve
* **Initial Loss:** 10.8269
* **End of Epoch 1 (Step 3100):** 5.3371
* **End of Epoch 2 (Step 6200):** 5.0588
* **End of Epoch 3 (Step 9350):** 4.8099
* **End of Epoch 4 (Step 12450):** 4.4191
* **End of Epoch 5 (Step 15600):** 4.3014

---

## How to Use

To use this model for text generation, define the architecture matching the parameters and load the saved `.safetensors` weights using MLX.

### 1. Requirements

Install the dependencies:
```bash
pip install mlx tiktoken huggingface_hub
```

### 2. Inference Code

```python
import os
import mlx.core as mx
import mlx.nn as nn
import tiktoken
from huggingface_hub import hf_hub_download

# Define model architecture identical to training
class TokenAndPositionEmbedding(nn.Module):
    def __init__(self, maxlen: int, vocab_size: int, embed_dim: int):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, embed_dim)
        self.pos_emb = nn.Embedding(maxlen, embed_dim)

    def __call__(self, x):
        seq_len = x.shape[1]
        positions = mx.arange(seq_len)[None, :]
        return self.token_emb(x) + self.pos_emb(positions)

class TransformerBlock(nn.Module):
    def __init__(self, emed_dim: int, num_heads: int, ff_dim: int):
        super().__init__()
        self.attention = nn.MultiHeadAttention(emed_dim, num_heads)
        
    def __call__(self, x, mask=None):
        attn_out = self.attention(x, x, x, mask=mask)
        return x + attn_out

class NanoLLM(nn.Module):
    def __init__(self, maxlen: int, vocab_size: int, embed_dim: int, num_heads: int, feed_forward_dim: int, num_transformer_blocks: int):
        super().__init__()
        self.maxlen = maxlen
        self.embedding = TokenAndPositionEmbedding(maxlen, vocab_size, embed_dim)
        self.transformer_blocks = [
            TransformerBlock(embed_dim, num_heads, feed_forward_dim)
            for _ in range(num_transformer_blocks)
        ]
        self.output_layer = nn.Linear(embed_dim, vocab_size, bias=False)

    def __call__(self, token_ids):
        seq_len = token_ids.shape[1]
        mask = nn.MultiHeadAttention.create_additive_causal_mask(seq_len)
        x = self.embedding(token_ids)
        for block in self.transformer_blocks:
            x = block(x, mask=mask)
        return self.output_layer(x)

# 1. Download weights from Hugging Face
repo_id = "samairtimer/nanollm_wiki_20.2M"
weights_path = hf_hub_download(repo_id=repo_id, filename="wikipedia_checkpoint.safetensors")

# 2. Initialize and load weights
tokenizer = tiktoken.get_encoding("gpt2")
model = NanoLLM(
    maxlen=128,
    vocab_size=tokenizer.n_vocab,
    embed_dim=192,
    num_heads=6,
    feed_forward_dim=512,
    num_transformer_blocks=6
)
model.load_weights(weights_path)
print("Model loaded successfully!")

# 3. Autoregressive Generation function
def generate(model, tokenizer, prompt, max_new_tokens=100, temperature=0.6):
    tokens = tokenizer.encode(prompt)
    x = mx.array(tokens)[None, :]
    end_token_id = tokenizer.encode('<|endoftext|>', allowed_special={'<|endoftext|>'})[0]
    
    print(prompt, end="", flush=True)
    for _ in range(max_new_tokens):
        if x.shape[1] > model.maxlen:
            x = x[:, -model.maxlen:]
        
        logits = model(x)
        next_token_logits = logits[0, -1, :]
        
        if temperature == 0.0:
            next_token = mx.argmax(next_token_logits, axis=-1).item()
        else:
            next_token = mx.random.categorical(next_token_logits / temperature, axis=-1).item()
            
        if next_token == end_token_id:
            break
            
        word = tokenizer.decode([next_token])
        print(word, end="", flush=True)
        x = mx.concatenate([x, mx.array([[next_token]])], axis=1)
    print("\n")

# Run generation
generate(model, tokenizer, "Cats are", max_new_tokens=50, temperature=0.4)
```

---

## Limitations & Biases

- **Extremely Small Scale:** At 20.2M parameters, the model is not capable of complex reasoning, logical deduction, or factual lookup.
- **Attention-Only:** The omission of MLP (Feed-Forward) layers limits the key-value lookup and relational storage capacity of the network, resulting in higher recurrence of repetitive structures.
- **Short Context Window:** The 128-token context window restricts its capacity to maintain long-term coherence.
- **Factual Inaccuracies:** The model is trained on a small subset of Wikipedia (100k entries) for 5 epochs. Any generated output should be treated strictly as synthetic text generation rather than historical or scientific fact.