Initial upload of TinyLLM

.gitignore

@@ -0,0 +1,2 @@
+# Content of .gitignore
+upload_repo.py

README.md

@@ -0,0 +1,133 @@
+---
+license: mit
+language: en
+tags:
+- llm
+- pytorch
+- custom-model
+- causal-lm
+- character-level
+- math
+- tiny-model
+model_type: tiny-causal-llm
+datasets:
+- custom
+pipeline_tag: text-generation
+---
+# TinyLLM: Character-Level Math Solver
+
+## Model Description
+
+**TinyLLM** is a highly compact, character-level **Causal Language Model** (based on the standard Transformer decoder architecture) trained specifically to solve single-digit math problems.
+
+This model serves as a minimalist, educational example of how a standard LLM architecture can be trained from scratch on a very small, custom dataset.
+
+### Key Features
+* **Architecture:** Causal Transformer Decoder.
+* **Task:** Character-level text generation (autoregressive).
+* **Input/Output:** Solves problems formatted as `N op N` and generates the answer, e.g., `4 + 5 = 9<EOS>`.
+* **Custom Code Required:** This is a custom PyTorch model and requires custom code (`model.py`, `tokenizer.py`) to be loaded by users.
+
+---
+##  How to Use (Inference)
+
+To load and run this custom model, users must download the entire repository structure and use the provided custom code, specifically the `TinyLLM` class defined in **`model.py`** and the `CharacterTokenizer` in **`tokenizer.py`**.
+
+### 1. Installation
+
+First, ensure you have the required libraries installed:
+```bash
+pip install torch huggingface-hub
+from huggingface_hub import snapshot_download
+import torch
+import os
+import sys
+
+# 1. Configuration: REPLACE with your repository ID
+MODEL_ID = "anujbhatt4ai/tiny-math-llm" 
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# 2. Download all files (code and weights)
+local_path = snapshot_download(repo_id=MODEL_ID)
+
+# 3. Import Custom Classes
+# The downloaded path must be added to sys.path to allow custom imports
+sys.path.append(local_path) 
+from model import TinyLLM
+from tokenizer import CharacterTokenizer, generate_v1_data
+
+# 4. Setup and Load Model
+def load_tiny_llm():
+    # In this minimal case, we hardcode the known config values
+    vocab_size = 22
+    block_size = 14
+    
+    # Initialize the model with the exact trained parameters
+    model = TinyLLM(
+        vocab_size=vocab_size, 
+        block_size=block_size, 
+        n_embed=64, n_head=4, n_layer=4, dropout=0.1
+    ).to(DEVICE)
+
+    # Load the trained weights
+    weights_path = os.path.join(local_path, "pytorch_model.bin")
+    model.load_state_dict(torch.load(weights_path, map_location=DEVICE))
+    model.eval()
+    
+    # Initialize the tokenizer
+    raw_data = generate_v1_data()
+    tokenizer = CharacterTokenizer(raw_data)
+    
+    return model, tokenizer
+
+# Use the loaded model and tokenizer in your own generation logic
+model, tokenizer = load_tiny_llm()
+print("Model loaded and ready for math inference!")
+
+**Block 4: Training Details and Repository Files**
+
+`markdown
+##  Training Details
+
+### Architecture Configuration
+
+The `TinyLLM` is configured with the following parameters, derived from the `config.json` and `model.py` files:
+
+| Parameter | Value | Description |
+| :--- | :--- | :--- |
+| **`vocab_size`** | `22` | The size of the character vocabulary. |
+| **`block_size`** | `14` | The maximum sequence length (context window). |
+| **`n_embed`** | `64` | Embedding dimension. |
+| **`n_head`** | `4` | Number of attention heads. |
+| **`n_layer`** | `4` | Number of Transformer decoder blocks. |
+| **`dropout`** | `0.1` | Dropout rate. |
+
+### Training Hyperparameters (from `train.py`)
+
+| Parameter | Value |
+| :--- | :--- |
+| **`BATCH_SIZE`** | `32` |
+| **`LEARNING_RATE`** | `1e-3` (AdamW) |
+| **`EPOCHS`** | `100` |
+| **`DEVICE`** | `cuda` if available, else `cpu` |
+
+### Dataset
+
+The model was trained on an **exhaustive set of all single-digit math problems** (addition, subtraction, multiplication, and non-remainder division) where the result is also a single digit (0-9). The **`dataset.py`** file contains the logic for the essential sequence shift used for language modeling training.
+
+---
+
+## Repository Files
+
+This flat repository contains all the source code needed for complete reproducibility.
+
+| File Name | Description |
+| :--- | :--- |
+| **`pytorch_model.bin`** | The trained model weights. |
+| **`config.json`** | Model configuration/hyperparameters. |
+| **`model.py`** | **Core Logic:** Custom `TinyLLM` architecture definition. |
+| **`tokenizer.py`** | **Core Logic:** Custom `CharacterTokenizer` and data generator. |
+| **`dataset.py`** | Defines the `MathDataset` class and sequence shift logic. |
+| **`train.py`** | The complete training script and final hyperparameters. |
+| **`custom_run.py`** (or `run.py`) | Example script demonstrating how to use the model for generation. |
+| **`README.md`** | This model card and documentation. |

config.json

@@ -0,0 +1,11 @@
+{
+  "n_embed": 64,
+  "n_head": 4,
+  "n_layer": 4,
+  "dropout": 0.1,
+  "vocab_size": 22,
+  "block_size": 14,  
+  "architectures": ["TinyLLM"],
+  "model_type": "tiny-causal-llm",
+  "_from_model_config": true
+}

custom_run.py

@@ -0,0 +1,154 @@
+import torch
+import torch.nn.functional as F
+import os
+import sys
+
+# --- Ensure src folder is in the path for imports ---
+# This helps the script find model.py, tokenizer.py, etc.
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src')))
+
+# --- Import all project components ---
+from src.tokenizer import generate_v1_data, CharacterTokenizer
+from src.model import TinyLLM, n_embed, n_head, n_layer, dropout # Also import hyperparams
+
+# --- Configuration (CHECK THIS PATH!) ---
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+# Use the file name confirmed in your last successful training run
+WEIGHTS_PATH = 'data/tinyllm_v1_weights1.pt' 
+
+
+@torch.no_grad()
+def generate(model, idx, max_new_tokens):
+    """
+    Takes a sequence of indices (idx) and generates max_new_tokens new indices 
+    using the model autoregressively.
+    """
+    model.eval() # Set model to evaluation mode
+    
+    # idx is (B, T) array of indices in the current context
+    for _ in range(max_new_tokens):
+        # Crop context to the model's block size (block_size will be set below)
+        block_size = model.block_size
+        idx_cond = idx[:, -block_size:] 
+        
+        # Get predictions
+        logits, _ = model(idx_cond)
+        
+        # Focus only on the last time step (the next token)
+        logits = logits[:, -1, :] 
+        
+        # Apply softmax to get probabilities
+        probs = F.softmax(logits, dim=-1)
+        
+        # Sample from the distribution
+        idx_next = torch.multinomial(probs, num_samples=1) 
+        
+        # Append sampled index to the running sequence
+        idx = torch.cat((idx, idx_next), dim=1) 
+    
+    return idx
+
+
+def setup_inference():
+    """Sets up the model, tokenizer, and loads weights for inference."""
+    try:
+        # 1. Setup Data Pipeline to determine sequence lengths
+        raw_data = generate_v1_data()
+        tokenizer = CharacterTokenizer(raw_data)
+        max_len = max(len(s) for s in raw_data)
+        
+        # FIX: Ensure block_size matches the model's training size (14)
+        # block_size is the maximum sequence length (T) the model can handle
+        block_size = max_len # Use max_len directly to get the 14 size for the V1 dataset
+        
+        # 2. Initialize Model Architecture
+        model = TinyLLM(
+            vocab_size=tokenizer.vocab_size,
+            n_embed=n_embed,
+            n_head=n_head,
+            n_layer=n_layer,
+            block_size=block_size,
+            dropout=dropout
+        ).to(DEVICE)
+        
+        # 3. Load Trained Weights
+        model.load_state_dict(torch.load(WEIGHTS_PATH, map_location=DEVICE))
+        print(f"\nSuccessfully loaded model weights from {WEIGHTS_PATH}")
+        
+        return model, tokenizer, block_size
+    
+    except FileNotFoundError:
+        print(f"Error: Weights file not found at {WEIGHTS_PATH}. Please run train.py first.")
+        return None, None, None
+    except RuntimeError as e:
+        print(f"Runtime Error during loading: {e}")
+        print("Please ensure your src/model.py hyperparameters match the saved weights.")
+        return None, None, None
+
+
+def solve_problem(model, tokenizer, question_str, block_size):
+    """Encodes a question, generates the answer, and prints the result."""
+    
+    # 1. Encode the question string (e.g., "5 + 3")
+    context_tokens = tokenizer.encode(question_str)
+    # Add an extra space before the = for clean formatting
+    context_tokens.append(tokenizer.encode(' ')[0]) 
+    
+    # Convert list of token IDs to a PyTorch tensor (1, T)
+    idx = torch.tensor([context_tokens], dtype=torch.long, device=DEVICE)
+    
+    # 2. Generate the rest of the sequence (the "= ANS" part)
+    # The max_len is the length of the expected output: = 9 (4 characters)
+    max_new_tokens = block_size - idx.shape[1] 
+    
+    if max_new_tokens <= 0:
+        print("Error: Input sequence is too long.")
+        return
+
+    # Generate tokens
+    generated_idx = generate(model, idx, max_new_tokens=max_new_tokens)
+    
+    # 3. Decode the result and print
+    generated_sequence = tokenizer.decode(generated_idx[0].tolist())
+    
+    print(f"Question: '{question_str}'")
+    print(f"Model Output: '{generated_sequence}'")
+
+
+# --- Main Interactive User Loop ---
+if __name__ == '__main__':
+    model, tokenizer, block_size = setup_inference()
+    
+    if model is not None:
+        print("\n--- TinyLLM Math Chatbot Initialized ---")
+        print("Enter a single-digit math problem (e.g., 4 + 5, 8 / 2).")
+        print("Type 'exit' to quit.")
+        
+        while True:
+            # 1. Get user input
+            question_str = input("Input: ")
+            
+            if question_str.lower() == 'exit':
+                break
+                
+            # 2. Basic Input Validation
+            question_str = question_str.strip()
+            parts = question_str.split()
+            
+            # Simple check for format N op N and single digits
+            is_valid = (
+                len(parts) == 3 and 
+                parts[0].isdigit() and len(parts[0]) == 1 and
+                parts[2].isdigit() and len(parts[2]) == 1 and
+                parts[1] in ['+', '-', '*', '/']
+            )
+
+            if not is_valid:
+                print("Error: Please enter a problem in the format 'N op N' with single-digit operands (e.g., 2 + 3).\n")
+                continue
+
+            # 3. Solve the problem using the trained model
+            solve_problem(model, tokenizer, question_str, block_size)
+            print("-" * 30)
+        
+        print("\n--- Chatbot Shutting Down ---")

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:81314013353fa80296b075a6a7a45e0db34ddc5584bbcac4728160cd5c60e3ad
+size 832685

src/dataset.py

@@ -0,0 +1,48 @@
+import torch
+from torch.utils.data import Dataset
+from typing import List, Tuple
+
+class MathDataset(Dataset):
+    """
+    A custom PyTorch Dataset to handle the encoded math problem sequences.
+    It performs the crucial language model shift (X is the input, Y is X shifted by one)
+    and handles padding.
+    """
+    def __init__(self, data: List[str], tokenizer, max_len: int):
+        self.data = data
+        self.tokenizer = tokenizer
+        self.max_len = max_len
+        self.pad_token_id = tokenizer.pad_token_id # Use the ID stored in the tokenizer
+
+    def __len__(self):
+        # Returns the total number of problems in the dataset
+        return len(self.data)
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        # 1. Get the raw encoded sequence (list of IDs)
+        raw_text = self.data[idx]
+        sequence_ids = self.tokenizer.encode(raw_text)
+        
+        # 2. Sequence Shift: The core of Language Modeling
+        # X (Input): The Transformer sees this. (e.g., [7, +, 2, =])
+        # Y (Target): The Transformer must predict this at the next step. (e.g., [+, 2, =, 9])
+        
+        # X: All tokens except the final <EOS> token (or final answer token)
+        # We cut off the last token because there is no token for the model to predict AFTER it.
+        x = sequence_ids[:-1] 
+        
+        # Y: All tokens except the first one. This is the sequence X is trying to predict.
+        # This is the "correct next token" for every position in X.
+        y = sequence_ids[1:]  
+        
+        # 3. Padding
+        # All sequences in a batch must have the same length (T or block_size).
+        
+        padding_length = self.max_len - len(x)
+        
+        # Pad the sequences X and Y with the <PAD> token ID
+        x_padded = x + [self.pad_token_id] * padding_length
+        y_padded = y + [self.pad_token_id] * padding_length
+        
+        # 4. Convert to PyTorch Tensors (dtype=torch.long is standard for integer IDs)
+        return torch.tensor(x_padded, dtype=torch.long), torch.tensor(y_padded, dtype=torch.long)

src/model.py

@@ -0,0 +1,170 @@
+from huggingface_hub import PyTorchModelHubMixin 
+
+
+# ... (rest of your model code)
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+# --- Hyperparameters (You can adjust these later) ---
+# For a "Tiny" LLM, we keep the size very small.
+n_embed = 64      # C: Embedding dimension (size of the vector representing a character)
+n_head = 4        # H: Number of attention heads
+n_layer = 4       # Number of repeating Transformer blocks
+dropout = 0.1     # Dropout rate
+
+# --- 1. Causal Self-Attention (The "Attention is All You Need" Component) ---
+
+class CausalSelfAttention(nn.Module):
+    """A multi-head masked self-attention module."""
+
+    def __init__(self, n_embed, n_head, block_size, dropout):
+        super().__init__()
+        
+        self.n_embed = n_embed
+        self.n_head = n_head
+        self.head_size = n_embed // n_head
+        
+        # Combined projection for Q, K, and V (more efficient)
+        self.c_attn = nn.Linear(n_embed, 3 * n_embed, bias=False)
+        # Output projection
+        self.c_proj = nn.Linear(n_embed, n_embed, bias=False)
+        self.attn_dropout = nn.Dropout(dropout)
+        self.resid_dropout = nn.Dropout(dropout)
+
+        # Causal Mask (tril = lower triangular matrix)
+        # This mask prevents a token from attending to future tokens (autoregressive)
+        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size))
+                                    .view(1, 1, block_size, block_size))
+
+    def forward(self, x):
+        B, T, C = x.shape  # Batch size, Sequence length (Time), Embedding dimension (Channel)
+
+        # 1. Compute Q, K, V and split (efficiently)
+        # q, k, v are (B, T, C)
+        qkv = self.c_attn(x)
+        q, k, v = qkv.split(self.n_embed, dim=2) 
+        
+        # 2. Reshape for Multi-Head Attention (B, T, C) -> (B, H, T, Head_size)
+        # We prepare the tensors so that each head processes a smaller chunk of the dimension C
+        k = k.view(B, T, self.n_head, self.head_size).transpose(1, 2) 
+        q = q.view(B, T, self.n_head, self.head_size).transpose(1, 2)
+        v = v.view(B, T, self.n_head, self.head_size).transpose(1, 2)
+        
+        # 3. Scaled Dot-Product Attention: (B, H, T, T)
+        # wei = (q @ k.transpose(-2, -1)) / sqrt(Head_size)
+        wei = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_size))
+        
+        # 4. Apply Causal Mask
+        # Set attention scores to -inf for future tokens (where tril == 0)
+        wei = wei.masked_fill(self.tril[:,:,:T,:T] == 0, float('-inf'))
+        
+        # 5. Softmax and Dropout
+        wei = F.softmax(wei, dim=-1)
+        wei = self.attn_dropout(wei)
+
+        # 6. Compute Weighted Sum of Values: (B, H, T, Head_size)
+        out = wei @ v 
+
+        # 7. Re-assemble heads: (B, H, T, Head_size) -> (B, T, C)
+        out = out.transpose(1, 2).contiguous().view(B, T, C)
+
+        # 8. Final Linear Projection
+        out = self.resid_dropout(self.c_proj(out))
+        return out
+
+# --- 2. Feed Forward Network (FFN) ---
+
+class FeedForward(nn.Module):
+    """A two-layer MLP for processing attention output."""
+    def __init__(self, n_embed, dropout):
+        super().__init__()
+        self.net = nn.Sequential(
+            # Standard ratio is 4x the embedding size
+            nn.Linear(n_embed, 4 * n_embed),
+            nn.GELU(), # Modern activation function (smoother than ReLU)
+            nn.Linear(4 * n_embed, n_embed),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+# --- 3. Transformer Block (The Repeating Unit) ---
+
+class TransformerBlock(nn.Module):
+    """A standard Transformer decoder block with Attention and FFN."""
+
+    def __init__(self, n_embed, n_head, block_size, dropout):
+        super().__init__()
+        # LayerNorm applied BEFORE the sub-layer (Pre-Norm style)
+        self.ln_1 = nn.LayerNorm(n_embed) 
+        self.attn = CausalSelfAttention(n_embed, n_head, block_size, dropout)
+        self.ln_2 = nn.LayerNorm(n_embed)
+        self.ffn = FeedForward(n_embed, dropout)
+
+    def forward(self, x):
+        # 1. Attention with Residual Connection and LayerNorm
+        x = x + self.attn(self.ln_1(x))
+        # 2. FFN with Residual Connection and LayerNorm
+        x = x + self.ffn(self.ln_2(x))
+        return x
+
+# --- 4. The Final TinyLLM Model ---
+
+class TinyLLM(nn.Module, PyTorchModelHubMixin):
+    """The complete Decoder-Only Transformer model."""
+    
+    def __init__(self, vocab_size, n_embed, n_head, n_layer, block_size, dropout):
+        super().__init__()
+        
+        self.block_size = block_size
+        
+        self.token_embedding_table = nn.Embedding(vocab_size, n_embed)
+        # Positional Encoding: A fixed table for position information
+        self.position_embedding_table = nn.Embedding(block_size, n_embed)
+        
+        # Stack of Transformer Blocks
+        self.blocks = nn.Sequential(*[
+            TransformerBlock(n_embed, n_head, block_size, dropout) 
+            for _ in range(n_layer)
+        ])
+        
+        self.ln_f = nn.LayerNorm(n_embed) # Final LayerNorm
+        # Linear layer to map the embedding vector back to the vocabulary space
+        self.lm_head = nn.Linear(n_embed, vocab_size)
+
+    def forward(self, idx, targets=None):
+        # idx is the input tensor X of shape (B, T)
+        B, T = idx.shape
+        
+        # 1. Token and Positional Embeddings
+        # Token embedding: (B, T, C)
+        tok_emb = self.token_embedding_table(idx) 
+        # Position embedding: (T, C) -> expanded to (B, T, C)
+        pos = torch.arange(T, device=idx.device) 
+        pos_emb = self.position_embedding_table(pos) 
+        
+        # 2. Combine (Add) Embeddings
+        x = tok_emb + pos_emb # (B, T, C)
+        
+        # 3. Pass through Transformer Blocks
+        x = self.blocks(x) # (B, T, C)
+        
+        # 4. Final LayerNorm and Linear Head
+        x = self.ln_f(x)
+        logits = self.lm_head(x) # (B, T, vocab_size)
+
+        loss = None
+        if targets is not None:
+            # Reshape for CrossEntropyLoss: (B*T, vocab_size) and (B*T)
+            B, T, C = logits.shape
+            logits = logits.view(B*T, C)
+            targets = targets.view(B*T)
+            
+            # Compute the negative log-likelihood loss
+            loss = F.cross_entropy(logits, targets)
+
+        return logits, loss

src/tokenizer.py

@@ -0,0 +1,81 @@
+import random
+
+def generate_v1_data():
+    """Generates all exhaustive single-digit math problems."""
+    data = []
+    
+    # Operators and their functions
+    ops = {'+': lambda a, b: a + b, 
+           '-': lambda a, b: a - b, 
+           '*': lambda a, b: a * b, 
+           '/': lambda a, b: a / b}
+    
+    # Iterate through all single-digit pairs (0-9)
+    for a in range(10):
+        for b in range(10):
+            for op_char, op_func in ops.items():
+                
+                # Check for constraints: Single-Digit Answer (0-9) & Validity
+                
+                if op_char == '+':
+                    result = op_func(a, b)
+                    # Constraint: Sum must be a single digit (<= 9)
+                    if result <= 9:
+                        problem = f"{a} + {b} = {result}"
+                        data.append(problem)
+                        
+                elif op_char == '-':
+                    result = op_func(a, b)
+                    # Constraint: Result must be non-negative (>= 0) and <= 9
+                    if 0 <= result <= 9:
+                        problem = f"{a} - {b} = {result}"
+                        data.append(problem)
+                    
+                elif op_char == '*':
+                    result = op_func(a, b)
+                    # Constraint: Product must be a single digit (<= 9)
+                    if result <= 9:
+                        problem = f"{a} * {b} = {result}"
+                        data.append(problem)
+                        
+                elif op_char == '/':
+                    # Cannot divide by zero
+                    if b == 0:
+                        continue 
+                    result = op_func(a, b)
+                    # Constraint: Result must be a whole number (no remainder) AND a single digit (<= 9)
+                    if a % b == 0 and result <= 9:
+                        # Use int() to remove potential float from division result
+                        problem = f"{a} / {b} = {int(result)}" 
+                        data.append(problem)
+                        
+    # IMPORTANT: Shuffle and add <EOS> marker
+    random.shuffle(data)
+    final_data = [d + "<EOS>" for d in data]
+    
+    return final_data
+
+class CharacterTokenizer:
+    """A simple character-level tokenizer for the math problems."""
+    
+    def __init__(self, data):
+        # 1. Build the unique vocabulary from the entire dataset
+        # We need to make sure the data is generated first!
+        chars = sorted(list(set("".join(data))))
+        
+        # Add a Padding token for PyTorch batching
+        if '<PAD>' not in chars:
+             chars.append('<PAD>') 
+
+        self.stoi = {ch: i for i, ch in enumerate(chars)}
+        self.itos = {i: ch for i, ch in enumerate(chars)}
+        self.vocab_size = len(chars)
+        self.pad_token_id = self.stoi['<PAD>']
+
+    def encode(self, s):
+        """Encodes a string into a list of integers."""
+        return [self.stoi[c] for c in s]
+
+    def decode(self, l):
+        """Decodes a list of integers back into a string."""
+        return "".join([self.itos[i] for i in l])

train.py

@@ -0,0 +1,100 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+# --- Import all project components ---
+from src.tokenizer import generate_v1_data, CharacterTokenizer
+from src.dataset import MathDataset
+from src.model import TinyLLM, n_embed, n_head, n_layer, dropout # Also import hyperparams
+
+# --- Hyperparameters for Training ---
+BATCH_SIZE = 32
+LEARNING_RATE = 1e-3 # Standard starting learning rate for Adam
+EPOCHS = 100         # Number of full passes over the dataset (Adjust as needed)
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
+def setup_data_pipeline(batch_size=BATCH_SIZE):
+    """Sets up the data generation, tokenization, and PyTorch DataLoaders."""
+    
+    # 1. Generate data and initialize tokenizer
+    raw_data = generate_v1_data()
+    tokenizer = CharacterTokenizer(raw_data)
+    max_len = max(len(s) for s in raw_data)
+    
+    # 2. Create Dataset and DataLoader
+    train_dataset = MathDataset(raw_data, tokenizer, max_len)
+    train_dataloader = DataLoader(
+        train_dataset, 
+        batch_size=batch_size, 
+        shuffle=True, 
+        drop_last=True
+    )
+    
+    print(f"Total problems: {len(raw_data)}")
+    print(f"Vocabulary Size: {tokenizer.vocab_size}")
+    print(f"Max Sequence Length (T): {max_len}")
+    print(f"Device: {DEVICE}")
+    
+    return train_dataloader, tokenizer.vocab_size, max_len
+
+
+def train(dataloader, vocab_size, block_size):
+    """Initializes the model and runs the full training loop."""
+    
+    # 1. Initialize Model, Optimizer, and move to Device
+    model = TinyLLM(
+        vocab_size=vocab_size,
+        n_embed=n_embed,
+        n_head=n_head,
+        n_layer=n_layer,
+        block_size=block_size,
+        dropout=dropout
+    ).to(DEVICE)
+    
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
+    
+    print(f"TinyLLM Parameters: {sum(p.numel() for p in model.parameters())/1e3:.1f}K")
+    print(f"Starting training for {EPOCHS} epochs...")
+
+    # 2. Training Loop
+    for epoch in range(EPOCHS):
+        model.train() # Set model to training mode
+        total_loss = 0
+        
+        for batch_idx, (X, Y) in enumerate(dataloader):
+            # Move data to the selected device (CPU or CUDA)
+            X, Y = X.to(DEVICE), Y.to(DEVICE)
+            
+            # Forward pass: calculate logits and loss
+            logits, loss = model(X, targets=Y)
+            total_loss += loss.item()
+            
+            # Backward pass: calculate gradients and update weights
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            
+            # Log progress every 100 batches (adjust frequency if needed)
+            if batch_idx % 100 == 0 and batch_idx > 0:
+                 print(f"  Epoch {epoch}/{EPOCHS} | Batch {batch_idx}/{len(dataloader)} | Loss: {loss.item():.4f}")
+
+        avg_loss = total_loss / len(dataloader)
+        print(f"--- Epoch {epoch+1} Complete --- Average Loss: {avg_loss:.4f}")
+
+        # If the loss is very low, the model has likely memorized the math.
+        if avg_loss < 0.01:
+            print("Loss is very low. Stopping training early.")
+            break
+
+    # 3. Save the trained model
+    torch.save(model.state_dict(), 'data/tinyllm_v1_weights1.pt')
+    print("\nTraining complete! Model weights saved to data/tinyllm_v1_weights1.pt")
+
+
+if __name__ == '__main__':
+    # 1. Setup the data
+    dataloader, vocab_size, max_len = setup_data_pipeline()
+    
+    # 2. Start the training process
+    train(dataloader, vocab_size, max_len)