MiniLM

A small, from-scratch language modeling project trained on Shakespeare text. MiniLM implements two different approaches to text generation — a classical N-gram model and a hand-built Transformer (MiniGPT) — sharing a common BPE tokenizer, so you can train, compare, and interact with both through a local web UI.

Features

• Custom BPE tokenizer trained on the project's text corpus
• N-gram language model (configurable order, default 4-gram) with statistical text generation and backoff smoothing
• MiniGPT Transformer — a from-scratch GPT-style architecture (attention, embeddings, block-based generation) trained with a standard next-token objective
• Comparison tooling (compare.py) to train/evaluate both models side by side
• Local web app (app.py) — a single-file Flask GUI to generate text from either model, or both at once, with adjustable prompt, max tokens, and temperature

Project Structure

MiniLM/
├── checkpoints/
│   └── best_model.pt          # Trained Transformer weights + config
├── data/
│   └── shakespear.txt         # Training corpus
├── docs/
│   └── Mini_LM_Project_Plan.docx
├── ngram/
│   ├── ngram_model.py         # N-gram model implementation (backoff smoothing)
│   ├── ngram_model.json       # Trained/cached N-gram model
│   └── test_ngram.py
├── notebooks/                 # Exploratory / scratch notebooks
├── tokenizer/
│   ├── bpe.py                 # BPE tokenizer implementation
│   ├── vocab.json             # Trained vocabulary (tokenizer state)
│   └── test_bpe.py            # Also the tokenizer's training entry point
├── transformer/
│   ├── model.py                # MiniGPT model definition
│   ├── attention.py            # Attention mechanism
│   └── train.py                # Transformer training script
├── app.py                      # Flask web UI for text generation
├── compare.py                  # Train/compare N-gram vs Transformer
├── requirements.txt
└── README.md

Requirements

• Python 3.10+
• PyTorch (CUDA build recommended if you have a compatible GPU — the app auto-detects and uses cuda if available, otherwise falls back to cpu)
• Flask

Install everything with:

pip install -r requirements.txt

Setup & Usage

All commands below are run from the project root (MiniLM/), with the virtual environment activated (see Issue: ModuleNotFoundError below if you hit import errors).

1. Train the tokenizer

Trains the BPE tokenizer on a subset of the corpus and saves the vocabulary to tokenizer/vocab.json. (Despite the filename, test_bpe.py is the actual training entry point — it both trains and runs roundtrip validation.)

.venv\Scripts\python.exe tokenizer\test_bpe.py

2. Train the Transformer

Trains the MiniGPT model and saves the best checkpoint to checkpoints/best_model.pt. Takes roughly 10–20 minutes on an RTX 3050-class GPU (early stopping may finish it sooner).

.venv\Scripts\python.exe transformer\train.py

3. Train the N-gram model & compare

Builds a statistical N-gram model (default order 4) from data/shakespear.txt, caches it to ngram/ngram_model.json, and prints a side-by-side perplexity/generation comparison against the trained Transformer.

.venv\Scripts\python.exe compare.py

4. Launch the web UI

Once a tokenizer and at least one model are trained:

.venv\Scripts\python.exe app.py

Then open http://127.0.0.1:5000 in your browser.

Important: if you retrain the tokenizer, N-gram model, or Transformer while app.py is already running, you must restart the Flask process. Models are loaded once into memory at startup, not per-request — the running server has no way to know the underlying files (or code) changed. See Issue: stale server after retraining below.

Using the Web UI

The interface lets you:

• Enter a Prompt (defaults to ROMEO:, matching the Shakespeare training data)
• Choose a Model: Transformer, N-gram, or Both (side-by-side comparison)
• Adjust Max new tokens (20–400)
• Adjust Temperature (0.1–1.5) — lower values produce more predictable text, higher values produce more varied/random text

Click Generate to produce output. If a model hasn't been trained yet, the app will show an error telling you which script to run first.

Note: generation is stochastic — the same prompt, model, and temperature will produce a different output on every click, since there's no fixed random seed. This is expected, not a bug.

Notes

• The app loads all models once at startup (not per request) since loading is the slow part — expect a short delay before the server reports Ready on device: <cuda|cpu>.
• app.py runs Flask's built-in development server (debug=False). This is fine for local experimentation but not intended for production deployment — use a proper WSGI server (e.g. Gunicorn or Waitress) if you need to serve this beyond your own machine.
• Model outputs are a product of a small vocabulary and modest training corpus, so expect stylistically convincing but not fully coherent Shakespearean-style text.

---

Issues Encountered & Fixes

A running log of real problems hit during development — what caused them, how they were fixed, and why that specific fix was chosen over alternatives.

1. Flask missing from requirements.txt

Symptom: app.py imports Flask, but requirements.txt only listed torch, numpy, matplotlib, tokenizers, and notebook. A fresh pip install -r requirements.txt would leave app.py unable to run (ModuleNotFoundError: No module named 'flask').

Cause: Flask had been installed ad-hoc during development (per the original docstring instruction pip install flask) but never added to the tracked requirements file.

Fix: Added Flask==3.1.3 (current stable release) to requirements.txt, pinned to a specific version like every other dependency, so installs stay reproducible across machines instead of silently picking up whatever the latest Flask happens to be at install time.

2. N-gram model generating pure noise on longer outputs

Symptom: With model=ngram, output degraded into unreadable character soup with no spaces or word structure at all (e.g. cityeveryWhyCOMspiritundwordstribuneshas...), especially at higher max_tokens.

Cause: get_next_token_probs() used flat Laplace (add-alpha) smoothing:

probs[tok_id] = (count + alpha) / (context_count + alpha * vocab_size)

When the model hit a 3-token context it had never seen during training (very common — a 4-gram context space is enormous relative to a single Shakespeare corpus), context_count and every count were 0, collapsing the formula to alpha / (alpha * vocab_size) = 1 / vocab_size. Every token became exactly equally likely — pure uniform random sampling over the entire vocabulary, completely ignoring language structure.

Fix: Rewrote ngram_model.py to train and store counts at every order from 1 (unigram) up to n, and implemented backoff: if the full n-gram context is unseen, fall back to a shorter (n−1)-gram context, then shorter again, down to the unigram distribution — using real observed statistics at whichever level actually has data, instead of jumping straight to uniform noise.

Why backoff over alternatives: Increasing alpha would only soften the uniform collapse, not eliminate it. Training a lower-order model from the start (e.g. dropping to a fixed trigram) would throw away useful long-context signal on the (many) contexts that were seen. Backoff keeps the best available signal at every context length without discarding the high-order model's precision when it does have data. The tradeoff: this changed the save/load JSON format (per-order counts instead of one flat dict), so any cached ngram_model.json had to be deleted and retrained.

3. N-gram output still running words together after the backoff fix

Symptom: After fix #2, output had real words and grammar fragments, but still occasionally jammed words together with no space (e.g. manbetter, notperate).

Cause: A separate, deeper issue in the tokenizer, not the N-gram model. The original pre-tokenization regex (\w+|[^\w\s]+|\s+) split whitespace into its own standalone chunk, so a space was always just one token among ~1000, structurally indistinguishable from any word-fragment token. At low backoff orders, the model had no signal telling it "a boundary must go here," so it sometimes picked two word-tokens back to back with no space token between them.

Fix: Changed the pre-tokenization pattern in bpe.py to GPT-2-style leading-space fusion: " ?\w+| ?[^\w\s]+|\s+". A single leading space now attaches to the token chunk that follows it ("the king" → ["the", " king"] instead of ["the", " ", "king"]), so BPE merges bake the space directly into the resulting subword tokens. A token like " king" either appears whole or not at all — there's no way to "skip the space but keep the word."

Why this over a decode-time patch: A simpler-looking fix would be to post-process generated text and insert spaces heuristically after the fact. That was rejected because it would paper over the model's actual behavior rather than fix it — the model would still be choosing to run words together internally; a string-cleanup pass just hides it, and would break down on cases a heuristic didn't anticipate. Fixing it at the tokenizer level makes correct spacing a structural guarantee instead of a best-effort patch.

Cost of this decision: Changing the tokenizer changes what every token ID means, so it forced a full retraining cascade:

1. Deleted tokenizer/vocab.json, ngram/ngram_model.json, checkpoints/best_model.pt
2. Retrained the tokenizer (tokenizer/test_bpe.py)
3. Retrained the Transformer (transformer/train.py) — the expensive step, ~10–20 min
4. Retrained the N-gram model (compare.py)

Result: The Transformer's output now has consistently correct spacing and grammar. The N-gram model still occasionally misplaces boundaries — this is treated as an inherent limitation of a 4-gram count-based model, not a remaining bug: with a context of only 3 prior tokens and a ~1000-token vocab, most contexts are unique/unseen, so the model has no reliable way to learn "don't emit a word-start token here." The Transformer doesn't share this problem because it attends to the full preceding sequence rather than a fixed short window.

Note: perplexity numbers are not comparable before vs. after this fix. The new tokenizer produces different (longer, fused) tokens, so each next-token prediction is a different, harder classification problem. A higher perplexity number after this change does not mean worse output — the actual generated text is qualitatively better.

4. ModuleNotFoundError: torch / flask when running scripts

Symptom: python compare.py (or any script) failed with ModuleNotFoundError, despite the packages being installed.

Cause: Running plain python invokes the system/global Python install, not the project's .venv, where the actual dependencies live.

Fix: Either activate the virtual environment first —

.venv\Scripts\Activate.ps1
python compare.py

— or call the venv's Python directly without activating:

.venv\Scripts\python.exe compare.py

(If PowerShell blocks activation with an execution-policy error, run Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass first, then retry — this only affects the current terminal session, not the system permanently.)

5. Stale Flask server after retraining a model

Symptom: After fixing and retraining the N-gram model, the web UI still produced the old broken (pure-noise) output.

Cause: app.py loads the tokenizer and both models once, at process startup — not per-request. A Flask server that had been running since before the code/model changes was still holding the old class definitions and old cached weights/counts in memory. Python does not hot-reload already-imported modules or already-loaded model state.

Fix: Stop the running server (Ctrl+C) and restart it (.venv\Scripts\python.exe app.py) any time the tokenizer, N-gram model, or Transformer checkpoint changes on disk.

6. git push rejected: "Updates were rejected because the remote contains work that you do not have locally"

Symptom: A local commit was made successfully, but git push failed.

Cause: The remote (GitHub) had commits the local repo didn't — e.g. from editing repository metadata (description/topics) directly on GitHub's web UI.

Fix:

git pull origin main
git push

git pull merges the remote changes into the local branch (opening an editor for a merge commit message if needed); once merged, the push proceeds normally. If the same lines had been changed both locally and remotely, this would instead produce a merge conflict requiring manual resolution — not encountered in this case, since only unrelated metadata differed.