Athena — 4.87M Shakespeare Dialogue Model
Athena is an experimental 4,871,168-parameter causal Transformer, built in PyTorch and trained from random initialization on dialogue extracted from William Shakespeare's complete works. Its goal is to explore learning conversational turn-taking from a small, self-prepared dataset. This release contains the 3,250-step checkpoint.
Athena has learned recognizable English and Shakespeare-style patterns, but its responses remain short, repetitive, and often unrelated to the prompt. It is an early research and learning prototype, not a validated general conversational assistant.
Architecture
| Property | Value |
|---|---|
| Framework | PyTorch, custom nn.Module |
| Unique trainable parameters | 4,871,168 |
| Transformer layers | 6 |
| Hidden width | 256 |
| Attention heads | 8 |
| Feed-forward width | 1,024 |
| Attention | Causal self-attention |
| Normalization | Pre-norm blocks and final LayerNorm |
| Activation | GELU |
| Dropout | 0.1 during training |
| Position encoding | Learned absolute embeddings |
| Context length | 256 tokens (UTF-8 bytes plus special tokens) |
| Vocabulary | 260 tokens |
| Input/output embeddings | Tied weights |
Tokenization maps UTF-8 bytes directly to IDs 0–255. Special IDs are PAD=256, BOS=257, EOS=258, and TURN=259. No learned tokenizer or external vocabulary is required. The context window is 256 bytes/tokens, not 256 words. Multi-byte characters use more than one token.
Training data
Source: The Complete Works of William Shakespeare, Project Gutenberg ebook #100, provided as a PDF. The edition metadata reports an update date of October 29, 2024.
The cleaned corpus contains 31,348 speaker turns in 753 sequences from 38 plays. Cleaning removed detected stage directions from 4,813 turns, excluded standalone prologue/epilogue/chorus sections and non-dialogue works, and quarantined nine ambiguous turns. Original Shakespearean wording is retained. Some subtle stage directions may remain because cleaning used automated rules and sample review.
Speaker identities and source references are preserved in the preparation records. Training input omits character names and uses TURN tokens between utterances. Training teaches continuation and turn boundaries; it does not assign persistent user/assistant identities.
The source scenes were deterministically split using seed 42: 598 training scene groups, 74 validation groups, and 74 test groups. Fragments of a source scene remain in the same split. Splits are approximately 80/10/10 by scene group, not by tokens; entire plays are not held out.
Cleaned dataset SHA-256:
ebdaef9f5d4898964b02618e2727ad49ae04ab71e4b102d58f2b51e9918b7889
Training procedure
- Initialization: random weights; no pretrained model used.
- Objective: next-token cross entropy, with padding targets ignored.
- Optimizer: AdamW.
- Configuration in supplied code: learning rate 0.0003, weight decay 0.01, gradient clipping at norm 1.0.
- Batch size: 16; sequence length: 256.
- Random training windows sampled within scene sequences, weighted by sequence length.
- Float32 training in the supplied implementation; no mixed precision configured.
- Checkpoint step: 3,250 optimizer updates.
- Training workflow: Google Colab notebook. Actual GPU model and total training duration were not recorded in the supplied checkpoint.
Hyperparameters above describe the supplied training code. The release checkpoint preserves optimizer settings and architecture; historical command-line overrides and a complete training log were not supplied. Optimizer state is omitted from this inference release. Use the original backup ZIP to resume training.
Evaluation
Both checkpoints were evaluated on the same 128 sampled windows of 256 tokens from held-out test scenes: 32,768 target tokens total, sampling seed 991, token-weighted cross entropy. Sampling may revisit or overlap text. This is a sampled evaluation, not the full test corpus.
| Checkpoint | Test loss (nats/token) ↓ | Byte-token perplexity ↓ | Next-token accuracy ↑ |
|---|---|---|---|
| 250 steps | 2.4059 | 11.0889 | 29.53% |
| 3,250 steps | 1.4584 | 4.2992 | 55.00% |
Test loss decreased by approximately 39.4%. These metrics measure byte/token prediction; they do not measure factual accuracy or the percentage of questions answered correctly. No conversational benchmark or human preference study has been completed.
With temperature 0.8, top-k 40, generation seed 123, and a maximum of 160 generated bytes, the 3,250-step checkpoint produced:
| Prompt | Response |
|---|---|
| Hello, how are you? | But there? |
| What is your name? | But there? |
| Why are you sad? | By you. |
Samples illustrate current limitations rather than reliable conversational skill. Generation may differ across software versions and devices. Detailed evaluation output is in comparison_results.json.
Usage
Download all files from this repository into one folder, install PyTorch, then run:
python -m pip install -r requirements.txt
python athena.py chat --checkpoint model.pt
Type /reset to clear the conversation or /quit to exit. The custom code runs on CUDA when available and otherwise uses CPU. This small model can run on a CPU.
For direct loading:
import torch
from athena import Athena
checkpoint = torch.load("model.pt", map_location="cpu", weights_only=True)
model = Athena(**checkpoint["config"])
model.load_state_dict(checkpoint["model"])
model.eval()
This is a custom PyTorch checkpoint, not a Transformers AutoModelForCausalLM package. The model card's text-generation task tag describes its purpose; automatic Hugging Face hosted inference/widget support is not configured. Use the included code to run it.
Intended use and limitations
Intended for learning, small language-model experiments, Shakespeare-style text generation, and studying the effects of additional training. The short context restricts conversation memory. Athena can produce malformed words, invalid UTF-8 sequences decoded with replacement characters, incoherent text, repetitions, and empty turns. It does not reliably follow instructions or ground responses in facts.
The corpus contains historical language and social biases, including discriminatory themes, violence, and sexual material. Outputs may reflect these. No safety alignment or moderation training was performed. Do not rely on its output for consequential decisions.
Release files and licensing
README.md: this model card.model.pt: inference checkpoint (model weights, configuration, step, and dataset hash).athena.py: model implementation and command-line chat; also includes original preparation/training functions.requirements.txt: PyTorch dependency.comparison_results.json: measured checkpoint comparison.cleaning_report.json: corpus cleaning summary.
A license for the newly created model weights and code has not yet been selected, so this card does not assert one. The source edition includes Project Gutenberg terms; availability and rights can vary by jurisdiction. Source: https://www.gutenberg.org/ebooks/100.