Upload folder using huggingface_hub

README.md

@@ -1,34 +1,68 @@
 ---
-title: Small Calculator Model
+title: Calculator LLM
 emoji: 🧮
 colorFrom: blue
 colorTo: purple
 sdk: gradio
-sdk_version: 6.3.0
+sdk_version: "5.9.1"
 app_file: app.py
 pinned: false
-license: apache-2.0
+license: mit
 ---
 
-# Calculator LLM
+# 🧮 Calculator LLM
 
-A tiny LLM that converts English math phrases to answers. Built from scratch to learn how language models work.
+A tiny transformer model (~105K parameters) that solves English math problems.
 
-## What it does
+## Try It
 
-Converts text like:
-```
-"two plus three" -> "five"
-"seven minus four" -> "three"
-"six times eight" -> "forty eight"
-```
+Enter a math problem in English like:
+- "two plus three"
+- "seven times eight"
+- "twenty minus five"
 
-## Try it
+The model will output the answer in English!
 
-Type a math phrase in English (e.g., "seven times eight") and watch the model respond.
+## Examples
 
-## Learn More
+| Input | Output |
+|-------|--------|
+| two plus three | five |
+| seven times eight | fifty six |
+| twenty minus five | fifteen |
+| nine times nine | eighty one |
 
-This model is part of an educational tutorial: [Build Your First LLM](https://sid.sh/learn/build-your-first-llm)
+## Built From Scratch
 
-Source code: [GitHub](https://github.com/slahiri/small_calculator_model)
+This model was built following the tutorial at [sid.sh/learn/build-your-first-llm](https://sid.sh/learn/build-your-first-llm)
+
+Same architecture as GPT (attention, feed-forward, transformer blocks), just much smaller!
+
+## Model Details
+
+| Property | Value |
+|----------|-------|
+| Parameters | ~105K |
+| Layers | 2 transformer blocks |
+| Embedding | 64 dimensions |
+| Attention Heads | 4 |
+| Vocabulary | 36 tokens |
+| Operations | plus, minus, times |
+| Number Range | 0-99 |
+
+## Architecture
+
+This is a decoder-only transformer with:
+- Token embeddings + sinusoidal positional encoding
+- 2 transformer blocks (multi-head attention + feed-forward)
+- Causal masking for autoregressive generation
+- Layer normalization and residual connections
+
+## Training
+
+Trained on 5,000 randomly generated math problems for 20 epochs.
+
+## Links
+
+- 📚 [Tutorial: Build Your First LLM](https://sid.sh/learn/build-your-first-llm)
+- 💻 [Full Notebook on GitHub](https://github.com/slahiri/blog/blob/main/public/notebooks/full_calculator_llm.ipynb)

app.py

@@ -0,0 +1,76 @@
+"""
+Calculator LLM Demo - A tiny transformer that solves English math problems.
+Built from scratch following: https://sid.sh/learn/build-your-first-llm
+"""
+
+import gradio as gr
+from calculator_llm import load_model, generate
+
+# Load the trained model
+print("Loading model...")
+model, tokenizer, vocab = load_model(".")
+print("Model loaded!")
+
+
+def solve_math(problem: str) -> str:
+    """Solve an English math problem."""
+    if not problem.strip():
+        return ""
+
+    # Ensure the problem ends with 'equals'
+    problem = problem.lower().strip()
+    if not problem.endswith("equals"):
+        problem = problem + " equals"
+
+    # Generate the answer
+    result = generate(model, tokenizer, vocab, problem)
+
+    # Extract just the answer part
+    answer = result.replace(problem, "").strip()
+
+    return answer
+
+
+# Example problems
+examples = [
+    ["two plus three"],
+    ["seven times eight"],
+    ["fifteen minus six"],
+    ["twenty plus thirty"],
+    ["nine times nine"],
+    ["fifty minus twenty five"],
+    ["twelve plus seven"],
+    ["eighty one minus forty"],
+]
+
+# Create the Gradio interface
+demo = gr.Interface(
+    fn=solve_math,
+    inputs=gr.Textbox(
+        label="Math Problem (in English)",
+        placeholder="e.g., two plus three",
+        info="Enter numbers 0-99 with operations: plus, minus, times",
+    ),
+    outputs=gr.Textbox(label="Answer"),
+    title="🧮 Calculator LLM",
+    description="""
+A tiny transformer model (~105K parameters) that solves English math problems.
+
+**Built from scratch** following the tutorial at [sid.sh/learn/build-your-first-llm](https://sid.sh/learn/build-your-first-llm)
+
+Same architecture as GPT, just much smaller!
+
+| Property | Value |
+|----------|-------|
+| Parameters | ~105K |
+| Layers | 2 |
+| Attention Heads | 4 |
+| Embedding Dim | 64 |
+""",
+    examples=examples,
+    theme=gr.themes.Soft(),
+    cache_examples=True,
+)
+
+if __name__ == "__main__":
+    demo.launch()

calculator_llm.py

@@ -0,0 +1,254 @@
+"""
+Calculator LLM - A tiny transformer for solving English math problems.
+Built from scratch following: https://sid.sh/learn/build-your-first-llm
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import json
+
+
+class PositionalEncoding(nn.Module):
+    """Adds positional information to embeddings using sine/cosine waves."""
+
+    def __init__(self, embed_dim, max_seq_len=512, dropout=0.1):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        pe = torch.zeros(max_seq_len, embed_dim)
+        position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, embed_dim, 2).float() * (-math.log(10000.0) / embed_dim)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        x = x + self.pe[:, : x.size(1), :]
+        return self.dropout(x)
+
+
+class TokenEmbedding(nn.Module):
+    """Converts token IDs to embedding vectors with positional encoding."""
+
+    def __init__(self, vocab_size, embed_dim, max_seq_len, dropout=0.1):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.token_embedding = nn.Embedding(vocab_size, embed_dim)
+        self.pos_encoding = PositionalEncoding(embed_dim, max_seq_len, dropout)
+        self.scale = math.sqrt(embed_dim)
+
+    def forward(self, x):
+        x = self.token_embedding(x) * self.scale
+        x = self.pos_encoding(x)
+        return x
+
+
+class MultiHeadAttention(nn.Module):
+    """Multi-head self-attention mechanism."""
+
+    def __init__(self, embed_dim, num_heads, dropout=0.1):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.scale = math.sqrt(self.head_dim)
+
+    def forward(self, x, mask=None):
+        batch_size, seq_len, _ = x.shape
+        Q = (
+            self.q_proj(x)
+            .view(batch_size, seq_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+        )
+        K = (
+            self.k_proj(x)
+            .view(batch_size, seq_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+        )
+        V = (
+            self.v_proj(x)
+            .view(batch_size, seq_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+        )
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, float("-inf"))
+        attn_weights = F.softmax(scores, dim=-1)
+        attn_weights = self.dropout(attn_weights)
+        attn_output = torch.matmul(attn_weights, V)
+        attn_output = (
+            attn_output.transpose(1, 2)
+            .contiguous()
+            .view(batch_size, seq_len, self.embed_dim)
+        )
+        return self.out_proj(attn_output), attn_weights
+
+
+class FeedForward(nn.Module):
+    """Position-wise feed-forward network."""
+
+    def __init__(self, embed_dim, ff_dim, dropout=0.1):
+        super().__init__()
+        self.linear1 = nn.Linear(embed_dim, ff_dim)
+        self.linear2 = nn.Linear(ff_dim, embed_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        return self.linear2(self.dropout(F.relu(self.linear1(x))))
+
+
+class TransformerBlock(nn.Module):
+    """A single transformer decoder block."""
+
+    def __init__(self, embed_dim, num_heads, ff_dim, dropout=0.1):
+        super().__init__()
+        self.attention = MultiHeadAttention(embed_dim, num_heads, dropout)
+        self.norm1 = nn.LayerNorm(embed_dim)
+        self.feed_forward = FeedForward(embed_dim, ff_dim, dropout)
+        self.norm2 = nn.LayerNorm(embed_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, mask=None):
+        attn_output, attn_weights = self.attention(x, mask)
+        x = self.norm1(x + self.dropout(attn_output))
+        ff_output = self.feed_forward(x)
+        x = self.norm2(x + self.dropout(ff_output))
+        return x, attn_weights
+
+
+def create_causal_mask(seq_len):
+    """Create a causal mask to prevent attending to future tokens."""
+    mask = torch.tril(torch.ones(seq_len, seq_len))
+    return mask.unsqueeze(0).unsqueeze(0)
+
+
+class CalculatorLLM(nn.Module):
+    """A tiny transformer LLM for solving English math problems."""
+
+    def __init__(
+        self, vocab_size, embed_dim, num_heads, num_layers, ff_dim, max_seq_len, dropout=0.1
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.max_seq_len = max_seq_len
+        self.embedding = TokenEmbedding(vocab_size, embed_dim, max_seq_len, dropout)
+        self.layers = nn.ModuleList(
+            [
+                TransformerBlock(embed_dim, num_heads, ff_dim, dropout)
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm = nn.LayerNorm(embed_dim)
+        self.output_proj = nn.Linear(embed_dim, vocab_size)
+
+    def forward(self, x, mask=None):
+        if mask is None:
+            seq_len = x.size(1)
+            mask = create_causal_mask(seq_len).to(x.device)
+        x = self.embedding(x)
+        for layer in self.layers:
+            x, _ = layer(x, mask)
+        x = self.norm(x)
+        return self.output_proj(x)
+
+
+class Tokenizer:
+    """Converts text to token IDs and back."""
+
+    def __init__(self, vocab):
+        self.vocab = vocab
+        self.id_to_word = {v: k for k, v in vocab.items()}
+
+    def normalize(self, text):
+        text = text.lower().strip()
+        text = text.replace("+", " plus ").replace("-", " minus ")
+        text = (
+            text.replace("*", " times ").replace("x", " times ").replace("=", " equals ")
+        )
+        tens = [
+            "twenty", "thirty", "forty", "fifty", "sixty", "seventy", "eighty", "ninety"
+        ]
+        ones = ["one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
+        for ten in tens:
+            for one in ones:
+                text = text.replace(f"{ten}{one}", f"{ten} {one}")
+        return " ".join(text.split())
+
+    def encode(self, text, add_special_tokens=True):
+        text = self.normalize(text)
+        ids = [self.vocab["[START]"]] if add_special_tokens else []
+        for word in text.split():
+            ids.append(self.vocab.get(word, self.vocab["[UNK]"]))
+        if add_special_tokens:
+            ids.append(self.vocab["[END]"])
+        return ids
+
+    def decode(self, ids, skip_special_tokens=True):
+        special = {"[PAD]", "[START]", "[END]", "[UNK]"}
+        words = [
+            self.id_to_word.get(id, "[UNK]")
+            for id in ids
+            if not (skip_special_tokens and self.id_to_word.get(id, "[UNK]") in special)
+        ]
+        return " ".join(words)
+
+
+def load_model(model_dir, device="cpu"):
+    """Load a saved Calculator LLM model."""
+    with open(f"{model_dir}/config.json") as f:
+        config = json.load(f)
+    with open(f"{model_dir}/vocab.json") as f:
+        vocab = json.load(f)
+
+    model = CalculatorLLM(
+        vocab_size=config["vocab_size"],
+        embed_dim=config["embed_dim"],
+        num_heads=config["num_heads"],
+        num_layers=config["num_layers"],
+        ff_dim=config["ff_dim"],
+        max_seq_len=config["max_seq_len"],
+        dropout=config["dropout"],
+    )
+    model.load_state_dict(
+        torch.load(f"{model_dir}/model.pt", map_location=device, weights_only=True)
+    )
+    model.to(device)
+    model.eval()
+
+    tokenizer = Tokenizer(vocab)
+    return model, tokenizer, vocab
+
+
+def generate(model, tokenizer, vocab, prompt, device="cpu", max_new_tokens=10):
+    """Generate text from a prompt."""
+    model.eval()
+    tokens = tokenizer.encode(prompt, add_special_tokens=True)[:-1]
+    input_ids = torch.tensor([tokens]).to(device)
+
+    with torch.no_grad():
+        for _ in range(max_new_tokens):
+            logits = model(input_ids)
+            next_token = logits[0, -1, :].argmax().item()
+            if next_token == vocab["[END]"]:
+                break
+            input_ids = torch.cat(
+                [input_ids, torch.tensor([[next_token]]).to(device)], dim=1
+            )
+
+    return tokenizer.decode(input_ids[0].tolist())
+
+
+if __name__ == "__main__":
+    # Example usage
+    model, tokenizer, vocab = load_model(".")
+    result = generate(model, tokenizer, vocab, "two plus three equals")
+    print(f"Result: {result}")

config.json

@@ -0,0 +1,9 @@
+{
+  "vocab_size": 36,
+  "embed_dim": 64,
+  "num_heads": 4,
+  "num_layers": 2,
+  "ff_dim": 256,
+  "max_seq_len": 16,
+  "dropout": 0.1
+}

model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c729d0dee5708cb1ee6f63509a3bea2556e0dafdc63e7b80b421932394e226f
+size 434201

requirements.txt

@@ -0,0 +1,2 @@
+torch
+gradio

vocab.json

@@ -0,0 +1,38 @@
+{
+  "[PAD]": 0,
+  "[START]": 1,
+  "[END]": 2,
+  "[UNK]": 3,
+  "zero": 4,
+  "one": 5,
+  "two": 6,
+  "three": 7,
+  "four": 8,
+  "five": 9,
+  "six": 10,
+  "seven": 11,
+  "eight": 12,
+  "nine": 13,
+  "ten": 14,
+  "eleven": 15,
+  "twelve": 16,
+  "thirteen": 17,
+  "fourteen": 18,
+  "fifteen": 19,
+  "sixteen": 20,
+  "seventeen": 21,
+  "eighteen": 22,
+  "nineteen": 23,
+  "twenty": 24,
+  "thirty": 25,
+  "forty": 26,
+  "fifty": 27,
+  "sixty": 28,
+  "seventy": 29,
+  "eighty": 30,
+  "ninety": 31,
+  "plus": 32,
+  "minus": 33,
+  "times": 34,
+  "equals": 35
+}