--- 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).