slm-125m-base/README.md

---
license: apache-2.0
language:
- en
library_name: transformers
pipeline_tag: text-generation
tags:
- llama
- legal
- finance
- small-language-model
- pretrained-from-scratch
- vizuara
datasets:
- HFforLegal/case-law
- PleIAs/SEC
- HuggingFaceFW/fineweb-edu
metrics:
- perplexity
---

# SLM-125M-base

A **125.8M-parameter Llama-style language model pretrained from random weights** on a
legal + financial corpus. It is a **base completer** (next-token prediction only), not
an instruction-tuned chatbot — give it the start of a sentence and it continues in the
legal/financial register.

Built end-to-end on [Modal](https://modal.com) for the **Vizuara "SLM from scratch"
workshop**, replicating the reference pipeline at
[Vizuara-AI-Lab/slm-125m-from-scratch](https://github.com/Vizuara-AI-Lab/slm-125m-from-scratch).

- 🕹️ **Live demo:** https://slm-125m-site.vercel.app
- ⚡ **Inference API:** `https://dharsourav03--slm-125m-inference-web.modal.run` (`/generate`)

| | |
|---|---|
| **Parameters** | 125,847,552 (~125.8M, tied embeddings) |
| **Validation perplexity** | **10.87** (held-out 1% split, ~22.0M tokens) |
| **Train tokens** | 2.18B (1 epoch) |
| **Vocab / context** | 16,384 byte-level BPE / 1,024 tokens |

## Intended use

Prompt it with the **opening of a sentence** and it completes the thought in a legal or
financial style. It is a demonstration of the full from-scratch pipeline (data →
tokenizer → pretraining), not a production model.

**Not intended for:** factual question answering, chat, instruction following, or any
high-stakes legal/financial decision-making. At 125M parameters it holds very little
factual knowledge and will confidently produce plausible-sounding but incorrect content.

## Model architecture

Maps 1:1 to `transformers.LlamaConfig`.

| Component | Value |
|---|---|
| Architecture | Llama-style decoder-only transformer |
| Layers | 12 |
| Hidden size | 768 |
| Attention heads | 12 (head dim 64), MHA (kv-heads = 12) |
| MLP | SwiGLU, inner dim 3072 |
| Normalization | RMSNorm (pre-norm), eps 1e-5 |
| Positional encoding | RoPE (theta 10000) |
| Context length | 1024 tokens |
| Vocabulary | 16,384 (byte-level BPE), tied embeddings |
| Attention bias | none |

## Training data

Streamed from HuggingFace, cleaned, deduplicated, and decontaminated. Realized mix
(by real tokens):

| Source | Content | Share |
|---|---|---|
| [HFforLegal/case-law](https://huggingface.co/datasets/HFforLegal/case-law) | US court opinions | ~33% |
| [PleIAs/SEC](https://huggingface.co/datasets/PleIAs/SEC) | SEC filings (10-K, etc.) | ~39% |
| [HuggingFaceFW/fineweb-edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) | Educational web text | ~28% |

**Pipeline:** stream → deterministic rule-based cleaning (line filters, boilerplate
strip, repetition/language/OCR gates) → exact + MinHash near-dedup → 13-gram
decontamination against CaseHOLD / LexGLUE eval sets → 16K byte-level BPE tokenizer →
pack into 1024-token windows (99/1 train/val split). Final corpus: **2.18B train + 22.0M
val tokens**.

## Training procedure

| Hyperparameter | Value |
|---|---|
| Objective | Next-token cross-entropy (causal LM) |
| Epochs | 1 (2.18B tokens seen) |
| Optimizer | AdamW, betas (0.9, 0.95), weight decay 0.1 |
| Learning rate | 6e-4 → 6e-5, cosine decay |
| Warmup | 200M tokens |
| Gradient clip | 1.0 |
| Global batch | 524,288 tokens/step (micro-batch 32 × seq 1024, grad-accum 2) |
| Precision | bf16 |
| Hardware | 8×H100 (single-node DDP), ~24 min |

> The reference workshop trains ~5 epochs (→ val perplexity 8.50). This checkpoint is a
> **single-epoch** run (val perplexity 10.87); more epochs over the same fixed corpus
> lower perplexity further.

## Evaluation

Held-out validation perplexity is `exp(mean token-level cross-entropy)` over the 1%
val split (21,505 windows, ~22.0M tokens):

**Validation perplexity: 10.87** (val loss 2.3856).

## How to use

The model uses custom special tokens; prepending `<|bos|>` matches how it was served and
gives the best continuations.

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "Bias-variance-tradeoff/slm-125m-base"
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(repo, torch_dtype=torch.float32).eval()

prompt = "The plaintiff alleges that the defendant"
bos = tok.convert_tokens_to_ids("<|bos|>")
eos = tok.convert_tokens_to_ids("<|eos|>")
ids = torch.tensor([[bos] + tok.encode(prompt, add_special_tokens=False)])

out = model.generate(
    ids, max_new_tokens=90, min_new_tokens=40, do_sample=True,
    temperature=0.8, top_k=50, top_p=0.95, repetition_penalty=1.3,
    eos_token_id=eos, pad_token_id=eos,
)
print(tok.decode(out[0][ids.shape[1]:], skip_special_tokens=True))
```

## Limitations and bias

- **Base model, not aligned** — no instruction tuning, no RLHF, no safety filtering.
- **Not a knowledge base** — a 125M model stores only a few tens of MB of usable
  knowledge; it fabricates citations, statutes, and figures. Do not rely on any factual
  claim it makes.
- **Domain-skewed** — trained ~72% on US legal + SEC text, so it defaults to that
  register and reflects the biases of those corpora (US-centric law, corporate finance).
- **English only.**

## Citation / attribution

Trained from scratch as part of the **Vizuara AI Labs "SLM from scratch" workshop**,
following the pipeline in
[Vizuara-AI-Lab/slm-125m-from-scratch](https://github.com/Vizuara-AI-Lab/slm-125m-from-scratch).