slm-125m-qa

ankitw497/slm-125m-qa

A full fine-tune of ankitw497/slm-125m-base (125.8M params, trained from scratch on legal/financial/web text) on a grounded QA dataset — 10,000 question-answer pairs synthesized with Gemini 3.1 Flash Lite from the same pretraining corpus, validated for grounding, and deduplicated.

Training

• Full fine-tune (no LoRA/adapters) — all 125.8M parameters updated
• 9,800 train / 200 val examples, chat format (<|user|>question<|assistant|>answer<|eos|>)
• Loss masked to the answer span only (question tokens contribute no gradient)
• 1x T4, AdamW, cosine LR schedule, peak lr 3e-5
• Checkpoint selection: best validation loss, not last step. Training ran 918 steps (3 epochs); val loss stopped improving at step 300 (well under 1 epoch) and rose afterward as the model overfit the small training set. This checkpoint is from step 300, not the final step.
• Best val loss: 2.1224 (perplexity 8.35), vs 22.9 baseline perplexity before fine-tuning

Known limitation

This is closed-book fine-tuning — the model was never given the source passage at training or inference time, only the question. At 125.8M parameters it does not have the capacity to reliably memorize facts from a 2.4B-token corpus, so while it has learned the domain's style well, treat specific facts, names, and figures in its answers as unverified. A retrieval-augmented (RAFT-style) setup, where the model is given the relevant source text at inference time rather than relying on memorization, would be the more reliable approach at this model size.

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("ankitw497/slm-125m-qa")
model = AutoModelForCausalLM.from_pretrained("ankitw497/slm-125m-qa")

prompt = "<|user|>{your question}<|assistant|>"
ids = tok(prompt, return_tensors="pt", return_token_type_ids=False)
out = model.generate(**ids, max_new_tokens=80, repetition_penalty=1.3, no_repeat_ngram_size=3)
print(tok.decode(out[0], skip_special_tokens=True))