HealthInsurance-FAQ-35M
A 35.65M-parameter GPT (decoder-only transformer) trained from scratch on a synthetic US health-insurance FAQ dataset (Medicare A/B/C/D, Medigap, ACA marketplace, annuities, IUL, life insurance, greetings, out-of-scope refusals). Trained on a single Apple M5 laptop GPU in about 65 minutes. Code, data, logs and evaluation are all included.
The architecture is copied from sraivante/tinystories-35m (base).
No pretrained weights were used. Training starts from random normal(0, 0.02) init and sees only the FAQ data.
| Folder | Checkpoint | Use it for |
|---|---|---|
final/ |
epoch 20 | Best answers (QA F1 0.379; ~97% exact recall of trained answers) |
best/ |
epoch 4, lowest val loss | Lowest perplexity on unseen text (val bits/char 1.006), weaker answers |
⚠️ General educational information only, not insurance advice. Answers can be wrong, especially for questions outside the training set. The data was generated with an LLM and has not been reviewed by a licensed or compliance professional. Have one approve it before any customer-facing use (CMS marketing rules apply to MA/PDP content).
Quick start
pip install torch tiktoken tokenizers huggingface_hub safetensors numpy
git clone https://huggingface.co/sraivante/healthinsurance-faq-35m && cd healthinsurance-faq-35m
python code/inference.py "What is Medigap?"
python code/inference.py --variant best "who are you"
python code/inference.py # interactive
Prompt format: User: {question}\nAssistant:. The answer ends with <|endoftext|>. Single-turn only.
Model
| Architecture | GPT-2 style: 9 layers, 6 heads, 384 dim, MLP 1536, pre-LayerNorm, GELU, learned absolute positions, tied input/output embeddings, no biases, dropout 0 |
| Parameters | 35,649,792 (19,316,736 of them are the tied token embedding) |
| Context | 1024 tokens |
| Tokenizer | GPT-2 BPE via tiktoken, vocab padded to 50,304. <|endoftext|> (50256) is BOS/EOS |
| Weights | safetensors, fp32. lm_head.weight is not stored (tied to wte.weight); load with strict=False (see code/inference.py) |
Training
- Data (sraivante/healthinsurance-faq, also in
data/): 1,795 train / 94 val Q&A pairs, plus 10 paraphrases per question. That gives 19,745 train documents (1.45M GPT-2 tokens) and 1,034 val documents. Paraphrases stay with their question's split, so no val question or rewording of one leaks into training. - Document-aware packing: each document is
<|endoftext|> User: … Assistant: … <|endoftext|>. Documents are packed into 1024-token rows (~99% fill).- Attention mask: same document AND causal. Positions restart at 0 for each document.
- Padding gets label −100 and attends only to itself.
- The model learns to emit EOS at the end of each answer.
- Optimizer: AdamW (β 0.9/0.95, wd 0.1). LR 1e-3 → 1e-4 cosine with 50 warmup steps. Grad clip 1.0.
- Batch and budget: effective batch 32 × 1024 tokens (micro-batch 4 × 8 accumulation). 20 epochs = 900 steps. fp32 on Apple MPS, ~7.4k tok/s.
- Init: normal(0, 0.02), residual projections 0.02/√(2·9). No phase-1 pre-pass.
Metrics per evaluation (results/healthinsurance-faq-35m/)
Val QA F1 is SQuAD token F1 of the greedy answer vs the reference answer, on the 94 val questions. Rep4 and Dist-2 are computed on generated tokens.
| Epoch | Step | Val PPL | Train PPL | Gap (val−train loss) | Val QA F1 | Val F1≥0.5 | Rep4 ↓ | Dist-2 ↑ |
|---|---|---|---|---|---|---|---|---|
| 2 | 90 | 39.65 | 25.68 | +0.435 | 0.127 | 0.011 | 0.291 | 0.091 |
4 (best/) |
180 | 25.85 | 4.65 | +1.716 | 0.228 | 0.043 | 0.070 | 0.373 |
| 6 | 270 | 39.85 | 1.95 | +3.018 | 0.271 | 0.117 | 0.010 | 0.575 |
| 8 | 360 | 51.71 | 1.61 | +3.469 | 0.294 | 0.170 | 0.005 | 0.617 |
| 10 | 450 | 64.38 | 1.42 | +3.813 | 0.343 | 0.234 | 0.004 | 0.640 |
| 12 | 540 | 76.81 | 1.30 | +4.079 | 0.350 | 0.191 | 0.002 | 0.653 |
| 14 | 630 | 96.86 | 1.23 | +4.363 | 0.375 | 0.277 | 0.001 | 0.653 |
| 16 | 720 | 109.29 | 1.20 | +4.512 | 0.370 | 0.266 | 0.006 | 0.643 |
| 18 | 810 | 117.96 | 1.18 | +4.601 | 0.379 | 0.287 | 0.003 | 0.642 |
20 (final/) |
900 | 124.17 | 1.18 | +4.658 | 0.379 | 0.266 | 0.001 | 0.674 |
Val perplexity bottoms out at epoch 4 and then rises as the model memorizes the training answers. The val answers never
appear in training, and ~2/3 of the training tokens are answers repeated across paraphrases. Answer quality (QA F1,
Dist-2) keeps improving anyway, which is why final/ is the recommended checkpoint for Q&A.
Does it answer correctly? (every question, final/)
comparison/all_questions_report.md and all_questions_answers.jsonl cover every train and val question, in its
original wording and one paraphrase each.
| Split | Wording | N | Exact match | Correct (F1≥0.8) | Mostly correct (F1≥0.5) | Mean F1 |
|---|---|---|---|---|---|---|
| train | original | 1795 | 97.4% | 97.7% | 99.4% | 0.990 |
| train | paraphrase | 1795 | 96.5% | 96.7% | 99.1% | 0.986 |
| val (unseen) | original | 94 | 0.0% | 1.1% | 26.6% | 0.379 |
| val (unseen) | paraphrase | 94 | 0.0% | 1.1% | 21.3% | 0.358 |
It works as an accurate answer bank for in-dataset questions, including reworded ones. On unseen questions it gives a fluent answer that is close to right about a quarter of the time.
Comparison: same data, different 35M design
The same data, packing, batch, epochs and evaluation were also used to train a custom 35.02M design from scratch:
21 layers × 352 dim, RMSNorm, RoPE, an own 10.7k byte-level BPE, pi-digit init and a phase-1 "unique-token cover"
pre-pass. Logs are in results/comparison-baseline-21l-rope/. Its weights are not included.
| Model | Checkpoint | Val bits/char ↓ | Val QA F1 ↑ | Val F1≥0.5 ↑ | Rep4 ↓ | Dist-2 ↑ | Train-question exact match |
|---|---|---|---|---|---|---|---|
| 21L RoPE, pi init | best (ep 6) | 1.062 | 0.211 | 0.021 | 0.029 | 0.333 | — |
| 21L RoPE, pi init | final (ep 20) | 1.205 | 0.288 | 0.138 | 0.002 | 0.607 | 28.1% |
| this model | best (ep 4) | 1.006 | 0.228 | 0.043 | 0.070 | 0.373 | — |
| this model | final (ep 20) | 1.491 | 0.379 | 0.266 | 0.001 | 0.674 | 97.4% |
Bits/char is tokenizer-independent, so it can be compared across the two models (their token perplexities cannot). The shallow-and-wide GPT-2 design wins on every quality metric and trains ~2× faster. These are single runs per design, and several factors differ at once (depth, tokenizer, positions, init, dropout, LR). Two notes on the custom model's pi-digit init:
- The literal
0.003 × digitinit, which makes every weight positive, stalled at val PPL ~530. - The zero-mean
0.003 × (digit − 4.5)init was needed to train at all. This matches the reference repo's finding that init scale matters, not its source.
Reproduce
pip install torch numpy tokenizers tiktoken gmpy2 pandas
python -u code/run.py all --variant tinystories # prepare (CPU) -> train (GPU) -> metrics table -> zip
python -u code/run.py all # the 21L RoPE / pi-init comparison model
python -u code/compare.py && python -u code/eval_all.py
python code/chat_ui.py # local streaming chat UI on http://localhost:8001
Training checkpoints every 100 steps and every epoch, and resumes automatically (run.py train). code/tests/ checks
that the packed loss equals the per-document loss (no cross-document attention) and that resume works.
Files
| Path | Contents |
|---|---|
final/, best/ |
model.safetensors + config.json |
code/ |
training library (autoreg_lib.py), runner, comparison/eval scripts, inference, chat UI |
data/ |
train.jsonl, val.jsonl ({"topic","q","a"}), paraphrases/*.jsonl |
results/ |
per-epoch metrics (md/csv/json), prep summaries and training logs for both designs |
comparison/ |
model comparison and all-questions correctness reports with every generated answer |
License
MIT, for code, weights and data.