SLM Code Optimizer

A code optimization transformer built entirely from scratch in PyTorch — no pretrained weights, no fine-tuning, no shortcuts.

---

What This Is

A Small Language Model (SLM) that takes suboptimal Python code, analyzes its time and space complexity, and rewrites it with a more efficient approach.

The goal was never to compete with GPT-4. The goal was to understand exactly how language models work — from byte-pair encoding to rotary positional embeddings to causal attention — by building every piece by hand.

Input:

def has_duplicates(arr):
    for i in range(len(arr)):
        for j in range(i + 1, len(arr)):
            if arr[i] == arr[j]:
                return True
    return False

Output:

Analysis: Original time complexity is O(n²) due to nested loops.
Space complexity is O(1). Optimized version uses a set for O(n)
average time complexity and O(n) space.

def has_duplicates(arr):
    seen = set()
    for item in arr:
        if item in seen:
            return True
        seen.add(item)
    return False

---

Architecture

Every component was written from scratch. Nothing was imported from HuggingFace, NLTK, or any ML library beyond PyTorch itself.

slm-project/
├── tokenizer/
│   └── bpe_tokenizer.py      # Byte Pair Encoding, built from scratch
├── data/
│   └── dataset.py            # PyTorch Dataset + DataLoader
├── model/
│   ├── rope.py               # Rotary Positional Encoding
│   ├── attention.py          # Multi-Head Self-Attention + causal mask
│   └── transformer.py        # Full transformer + SLMModel
├── training/
│   └── trainer.py            # Training loop with fp16, grad accum, cosine LR
├── inference/
│   └── generate.py           # Autoregressive generation, temp/top-k/top-p
├── static/
│   └── index.html            # Web UI
├── app.py                    # Flask server
├── config.py                 # All hyperparameters
├── train.py                  # Main training entry point
└── generate_dataset.py       # Parallelized dataset generation pipeline

Model Config

Parameter	Value
Architecture	Decoder-only Transformer
Parameters	33.9M
d_model	512
num_heads	8
num_layers	8
vocab_size	8,000
max_seq_len	512
Positional Encoding	RoPE
Normalization	RMSNorm (Pre-Norm)
Activation	SwiGLU

Training Config

Parameter	Value
Optimizer	AdamW (β₁=0.9, β₂=0.95)
Learning Rate	3e-4 with cosine decay
Warmup Steps	100
Batch Size	4 (effective 16 with grad accum)
Gradient Accumulation	4 steps
Mixed Precision	fp16
Epochs	15
Hardware	GTX 1650 Ti (4GB VRAM)
Training Time	~30 hours

---

Dataset

30,089 training pairs generated synthetically using the Gemini API, covering 35 problem categories across 15 anti-patterns.

Generation pipeline highlights:

• 5 parallel threads (one per API key) for ~5x throughput
• 3-model rotation per key: flash-lite → flash → pro
• 15 total API slots before any rate limiting stalls generation
• Auto-resume from checkpoint — safe to stop and restart anytime
• ~25,000 pairs generated in a single day at peak throughput

Problem categories include: array manipulation, dynamic programming, graph traversal, string operations, sliding window, monotonic stack, union-find, heap/priority queue, backtracking, memoization patterns, NumPy vectorization, concurrency patterns, memory optimization, and more.

---

Training Results

Run	Examples	Epochs	Final Loss
Test run	600	15	1.4771
Full run	30,089	15	0.1574

Loss curve (full run):

Epoch  1  →  ~3.2
Epoch  3  →  ~1.8
Epoch  6  →  ~0.8
Epoch 10  →  ~0.38
Epoch 13  →  ~0.23
Epoch 15  →  0.1574

---

Capability Profile

Evaluated on 7 test cases across increasing complexity:

Pattern	Result	Notes
Nested loops → hash set	✅ Perfect	Correct analysis + runnable code
String concatenation → join	✅ Perfect	Reduced to one idiomatic line
List → set intersection	⚠️ Inconsistent	Correct concept, occasional variable hallucination
Redundant computation in loop	❌ Missed	Subtler pattern, not well generalized
Recursive memoization	✅ Perfect	Correct closure pattern with nested function
Multi-pattern, complex function	❌ Breaks down	Beyond single-pattern reasoning capacity
300-line input	❌ Truncated	Exceeds 512-token sequence length limit

Effective use case: Single well-defined anti-patterns in functions up to ~50 lines at temperature 0.1–0.3.

---

Known Limitations

• 512-token sequence limit — inputs longer than ~50-80 lines get truncated, producing incoherent output
• Single-pattern reasoning — works well on one anti-pattern per function; degrades on multiple interacting inefficiencies
• Non-deterministic — temperature > 0.3 can produce variable name hallucinations on the same input
• No conversational ability — pure code transformation model, not an instruction follower

---

Running Locally

1. Install dependencies

pip install torch flask flask-cors google-genai

2. Generate dataset (skip if using provided data)

# Add your Gemini API keys to generate_dataset.py
python generate_dataset.py

3. Train

python train.py

4. Launch web UI

python app.py
# Open http://localhost:5000

---

Web Interface

Dark-themed side-by-side editor with:

• Input panel for suboptimal code
• Output panel for analysis + optimized code
• Temperature slider (recommended: 0.1–0.3 for code tasks)
• Token limit slider
• Ctrl+Enter shortcut
• Copy output button

---

What I Learned

Building this from scratch meant understanding every design decision:

• Why RoPE over learned embeddings — relative position awareness, generalizes better
• Why Pre-Norm over Post-Norm — more stable gradients in deep networks
• Why SwiGLU over ReLU — gating mechanism is more expressive, used by LLaMA/Mistral
• Why the loss is shifted by 1 — model at position i predicts token at i+1
• Why gradient accumulation — simulates larger batch sizes within 4GB VRAM
• Why the lock is never held during API calls — threading parallelism requires minimal critical sections

---

Tech Stack

PyTorch · Flask · Google Gemini API · Python threading · fp16 mixed precision

---

License

MIT