Apex-64M: Small Language Model (SLM) & Chat Assistant

A complete, end-to-end 63.8M Parameter GPT-Style Decoder Transformer built and trained 100% from scratch in pure PyTorch, followed by Supervised Fine-Tuning (SFT) on Databricks Dolly 15k to create an instruction-following assistant.

---

Highlights

• Built 100% From Scratch: Zero external pre-trained weights used. Pre-trained on ~1.28 Billion tokens of the C4 English corpus from random weight initialization N(0, 0.02).
• Instruction Aligned (SFT): Fine-tuned on Databricks Dolly 15k with prompt loss masking (ignore_index = -1) and 1-token autoregressive shifting.
• Multi-Turn Chat Memory: Sliding-window context buffer supporting multi-turn dialogue within 512 tokens.
• Dual-Layer Guardrails: Keyword heuristics + neural toxicity filtering (unitary/toxic-bert) for input and output safety.
• Ultra Lightweight: Weighs only ~128 MB (FP16) / ~255 MB (FP32). Runs at lightning speed on standard laptop CPUs with <300 MB RAM!

---

Repository Structure

• ARCHITECTURE.md - Complete Architectural Deep Dive, mathematical derivations, and parameter accounting.
• model.py - 63.8M Transformer architecture with Pre-LN, Flash Attention (F.scaled_dot_product_attention), and tied embedding weights.
• config.py - Hyperparameter dataclasses for GPTConfig and TrainConfig.
• tokenizer.py - OpenAI tiktoken BPE wrapper (50,257 vocabulary size).
• dataset.py - Continuous token packing pipeline with memory-mapped caching.
• train.py - Pre-training loop with AMP FP16, Cosine LR scheduling, and multi-cloud checkpointing.
• checkpoint.py - Robust checkpoint saving, loading, and Drive mirroring.
• evaluate.py - Perplexity evaluation and generation benchmarking.
• app.py - Gradio Web UI for live chatting with the Apex model.

---

Model Specifications

Parameter	Value
Model Name	Apex-64M
Total Parameters	63,823,360 (~63.82M)
Embedding Dimension (d_model)	512
Transformer Blocks (n_layer)	12
Attention Heads (n_head)	8 (head_dim = 64)
Feed-Forward Dimension (d_ffn)	2048 (4 x d_model)
Context Length	512 tokens
Vocabulary Size	50,257 (tiktoken GPT-2 / r50k_base)
Attention Mechanism	Fused FlashAttention (F.scaled_dot_product_attention)
Weight Sharing	Token Embedding <-> LM Output Head (wte.weight)
Pre-Training Loss	4.02 on C4 English (1.28B tokens)
SFT Loss	~3.17 on Databricks Dolly 15k (3 Epochs)

---

Quickstart: Chat with Apex-64M

1. Installation

git clone https://github.com/kapilverse/SMALL_MODEL.git
cd SMALL_MODEL
pip install torch transformers datasets tiktoken gradio

2. Interactive CLI Chat (with Memory & Safety)

import torch
from config import GPTConfig
from model import SmallGPT
from tokenizer import get_tokenizer

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = get_tokenizer()
eot_id = tokenizer.eot_token

config = GPTConfig()
model = SmallGPT(config)

# Load SFT Weights
state_dict = torch.load("sft_model.pt", map_location=device)
model.load_state_dict(state_dict)
model = model.to(device).eval()

def ask(question: str, temperature: float = 0.7, top_p: float = 0.9):
    prompt = f"### Instruction:\n{question}\n\n### Response:\n"
    tokens = tokenizer.encode(prompt)
    idx = torch.tensor(tokens, dtype=torch.long, device=device).unsqueeze(0)
    
    with torch.no_grad():
        for _ in range(150):
            idx_cond = idx[:, -config.context_length:]
            logits, _ = model(idx_cond)
            logits = logits[:, -1, :] / max(temperature, 1e-5)
            
            # Top-P sampling
            sorted_logits, sorted_indices = torch.sort(logits, descending=True)
            cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
            sorted_indices_to_remove = cumulative_probs > top_p
            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
            sorted_indices_to_remove[..., 0] = 0
            indices_to_remove = sorted_indices[sorted_indices_to_remove]
            logits[:, indices_to_remove] = -float("Inf")
            
            probs = torch.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            
            if next_token.item() == eot_id:
                break
                
            idx = torch.cat((idx, next_token), dim=1)
            
            # Early stop on new instruction header
            current_text = tokenizer.decode(idx[0].tolist())
            if "### Instruction:" in current_text.split("### Response:\n")[-1]:
                break
                
    full_text = tokenizer.decode(idx[0].tolist())
    answer = full_text.split("### Response:\n")[-1].split("### Instruction:")[0].strip()
    return answer.replace("<|endoftext|>", "")

print("Apex:", ask("What is the difference between a stack and a queue?"))

3. Launch Gradio Web Interface

python app.py

---

Complete Training Pipeline

                       Stage 1: Pre-Training
                       ---------------------
   [C4 English ~1.28B Tokens] ---> [Apex-64M Base Foundation Model]
                                       | (best_model.pt, val_loss ~4.02)
                                       v
                       Stage 2: Supervised Fine-Tuning
                       -------------------------------
   [Databricks Dolly 15k]      ---> [Prompt Loss Masking & 1-Token Shift]
                                       | (sft_model.pt, loss ~3.17)
                                       v
                       Stage 3: Deployment
                       -------------------
                               [Apex-64M Assistant]
                               - Multi-Turn Memory Buffer
                               - Toxic-BERT Safety Guardrails
                               - Interactive Gradio Web UI

For complete technical details, mathematical derivations, loss masking explanations, and parameter breakdowns, see ARCHITECTURE.md.

---

License

This project is open-source under the MIT License.