Tiny Shakespeare Chat

A small GPT written and trained from scratch (no pretrained weights, no tokenizer library). It is a character-level, nanoGPT-style decoder with 10,751,232 parameters. It was first pretrained on the Tiny Shakespeare corpus to predict the next character, then chat-tuned on 7,096 consecutive dialogue-line pairs, so it answers a message with the "next line" of an imaginary play. Its replies are play-style verse, not facts.

Model

Usage

Python

from huggingface_hub import snapshot_download
import sys
path = snapshot_download("shalev396/tiny-shakespeare-chat")
sys.path.insert(0, path)
import model
predictor = model.load(path, device="cpu")   # or "cuda"
print(predictor.predict("How fares the king?"))
print(predictor.predict("And the queen?", history=[["How fares the king?", "He is well, my lord."]],
                        max_new_tokens=150, temperature=0.7, top_k=40, seed=0))
for piece in predictor.stream("Speak to me of love"):   # streaming, one character at a time
    print(piece, end="", flush=True)

history accepts [[user, bot], ...] pairs or [{"role": ..., "content": ...}, ...] messages. Sampling is random; pass seed for a repeatable reply. Requirements: torch, huggingface_hub, safetensors.

Space API (free): the Space exposes /predict (message, history as a JSON string, temperature, max_new_tokens in; [reply, seconds, device] out). curl and @gradio/client examples are in its README.

Inference Endpoint: deploy this repo from its page (Deploy -> Inference Endpoints, CPU or GPU). handler.py loads the model once (on cuda when the endpoint has a GPU).

curl $ENDPOINT_URL -H "Authorization: Bearer $HF_TOKEN" -H "Content-Type: application/json" \
  -d '{"inputs": {"message": "And the queen?", "history": [["How fares the king?", "He is well, my lord."]]},
       "parameters": {"max_new_tokens": 150, "temperature": 0.7, "top_k": 40, "seed": 1}}'
# [{"generated_text": "..."}]

A plain string ({"inputs": "How fares the king?"}) works too.

Training

• Data: Tiny Shakespeare (karpathy/char-rnn input.txt, 1,115,394 characters of Shakespeare plays).
• Stage A, pretrain: next-character prediction on the raw play, first 90 % train / last 10 % val. 5,000 iterations, batch 64 x 256 characters, AdamW (betas 0.9/0.95, weight decay 0.1 on matrices only), 100 warmup iterations, cosine LR 3e-4 -> 3e-5, gradient clipping 1.0, fp16 autocast on CUDA.
• Stage B, chat-tune: the play is parsed into consecutive (speaker, utterance) turns. Every pair of neighbouring turns becomes <|user|> {line} <|end|> <|bot|> {next line} <|end|>, plus a variant whose reply starts with SPEAKER:. The samples are shuffled (seed 42) and joined, and the last 5 % of that stream is val. 2,000 iterations from the stage A weights, same optimiser, 50 warmup iterations, cosine LR 1e-4 -> 1e-5.
• This checkpoint: converted into model.GPT from the earlier full run (legacy project 13_tiny_shakespeare_chat_pt, one CUDA GPU, 541 s for both stages, 2026-08-01). Every tensor loads with strict=True, the logits match the legacy code exactly (max |diff| 0.0) and seeded generations are identical. All metrics below were re-computed locally on CPU with this repo's training/src/engine.evaluate.

Full code: training/ · Colab.

Experiments

Every variant scored with the same deterministic evaluation (non-overlapping 256-character windows, every validation character scored once) on both validation splits (from metrics.json -> comparison). Cross-entropy in nats per character, lower is better. The deployed model is in bold.

• Pretraining alone takes the loss on unseen text from 2.48 (bigram) to 1.50 nats per character (2.16 bits per character).
• Chat-tuning drops the chat val loss from 1.56 to 1.05. Part of that is the model learning the chat format (markers, speaker names, where a reply ends). Part is the leak described in Limitations, so 1.05 is optimistic.

Training curves of the deployed run (the run's own estimates on 40 random 64 x 256 batches per split). Stage A's val loss flattens around 1.50 after about 3,000 iterations while train keeps falling (1.02 at the end), a mild overfit. Stage B's val loss is still creeping down at 2,000 iterations (1.054).

Evaluation

The deployed model on the chat val split (215,040 characters scored) and the text val split (111,360 characters). bpc = bits per character (loss / ln 2), ppl = per-character perplexity (e^loss).

Loss by position in the context window: the first characters of a window are hard to predict (about 2.4 nats with no context), and the loss levels off after roughly 30-50 characters. The model mostly uses about a line of context.

Limitations

• Toy model. 10.75M parameters trained on 1.1 MB of text. It imitates the style (speaker names, verse, archaic words) but invents words, loses the thread within a few lines and does not understand the message. Nothing it says is factual.
• Chat val leaks. The chat samples come from the whole play, and both variants of each pair are shuffled together before the split, so 95.8 % of chat-val samples have their twin (same dialogue pair, other variant) in chat-train. The chat-tuned model has also seen the stage A validation text inside chat pairs. So chat val loss is optimistic, and text val loss after stage B is not a held-out number. The clean held-out figure is stage A's 1.495 on text val. The split was kept so that the code reproduces these weights.
• 68 symbols. Characters outside the corpus alphabet (digits other than 3, brackets, double quotes, accents, emoji) are silently dropped from the message.
• 256-character context. Only the last 3 turns are kept, and long prompts are cropped from the left.
• Random sampling. The same message gives a different reply each time unless seed is set. A reply can be empty (the model may emit <|end|> at once) or cut off at max_new_tokens.
• The plays contain archaic, violent or offensive language, and the model can reproduce it.