legal-slm-336m

A 332.4M-parameter decoder-only language model pretrained from scratch on Indian legal text. This directory is the complete release: weights, config and tokenizer, loadable with no conversion step and no trust_remote_code.

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("legalslm/legal-slm-336m")
tok = AutoTokenizer.from_pretrained("legalslm/legal-slm-336m")

ids = tok("Section 103 of the Bharatiya Nyaya Sanhita", return_tensors="pt")
print(tok.decode(model.generate(**ids, max_new_tokens=40)[0]))

pad_token is set to <|endoftext|> (id 0), the same as eos_token, so batched scoring pads rather than raising. config.json declares vocab_size 32768, which matches the tokenizer exactly.

Not instruction-tuned. This is a base model: it continues text, it does not follow instructions.

Try it

python run_model.py                                   # generation from a default prompt
python run_model.py --from-json scoring_tests.json    # reproduce a stored result
python run_model.py --from-json scoring_tests.json --section bns:140
python run_model.py --prompt "Section 105 of the Bharatiya Nyaya Sanhita, 2023 —"

run_model.py ships in this folder and imports nothing beyond torch and transformers. It shows both modes — greedy generation, and log-probability scoring of candidate continuations, which is the mode the bits-per-byte result measures.

scoring_tests.json is the whole study behind the scoring claims below: 192 BNS sections x 3 probes, 576 records, each with the prompt, the candidates, the exact string fed to the model, the expected answer and the score it got. So the claims here can be recomputed rather than taken on trust. Read its conclusion field before quoting a number from it — the model discriminates statutory register and punishment category, but does NOT verify that a clause belongs to the section it cites, and the shuffled_control records are what establish that.

Architecture, and why it loads as Llama

332.4M parameters (298.9M non-embedding)
d_model 1024 · 24 layers · 16 query heads / 4 KV heads · context 1024
RMSNorm · RoPE · SwiGLU (hidden 3200) · GQA · tied embeddings
vocab 32,768 — custom byte-level BPE trained on this corpus

The architecture is Llama's, so the released weights are mapped onto LlamaForCausalLM rather than shipped with custom modelling code. The conversion was verified numerically, not by "it loaded": logits from the original implementation and the converted model agree to max |diff| 0.00000 with 100% argmax agreement in fp32, and all 218 stored tensors are bit-identical to the mapped weights. A conversion that silently transposes the fused QKV or uses the wrong RoPE convention loads perfectly and answers nonsense, so equality was measured.

Headline result

Bits per byte on held-out text, every model scored at context 1024 with identical documents and byte accounting. Lower is better. BPB rather than perplexity because perplexity is per-token and a custom 32K vocabulary would flatter this model against any other tokenizer.

30.9% better than the best competitor on Indian legal text and 25.2% better on statutes, against a model 49% larger trained on 18T tokens. Size-matched against gpt2-medium: 36.8% and 33.6% better. Same parameter count, different corpus.

And the cost: this is the worst model in the table on general text, beaten even by GPT-2 at 124M. Web text was 18.4% of the corpus. A domain model that beat everything on everything would indicate a contaminated evaluation rather than a better model, so the loss is reported as prominently as the wins.

Contamination is asymmetric and it favours the comparison models: this corpus was decontaminated against these exact evaluation sets (415,481 documents removed by 13-gram filtering, with an assertion that no held-out document reached a shard). None of the five comparison models carries such a guarantee, and pile-of-law and Indian Kanoon scrapes are plausible in several of their pretraining sets. These margins are a lower bound; theirs may include memorisation.

Training

14.0B tokens seen at this checkpoint · 78,667 steps · context 1024
Round 1: 12.198B tokens, 17h43m on one H200, 205,000 tok/s, MFU 44.4% median
Round 2: resumed with a corrected data position and a diluted anneal
AdamW, decoupled weight decay 0.1 (excluding norms and embeddings)
WSD schedule · bf16 autocast over fp32 master weights · gradient clipping 1.0

Corpus: 12.0B tokens from 21 sources, deduplicated, quality-filtered and decontaminated. Broad interleaved data for most of training, then an anneal concentrating statutes, Supreme Court judgments and instruction-style data.

This checkpoint was selected from a five-point decay series by a five-seed study on a downstream task — 71 legal fact patterns × 5 seeds × 3 candidates — not by taking the final step. Accuracy did not separate the candidates (pairwise McNemar p = 0.33 to 0.82); it was chosen for the lowest fabrication rate and much the lowest seed variance (sd 0.80 against 2.42 for the nearest rival, which had a seed collapse to 4/71).

Structural tokens

Every training document was prefixed with five reserved control tokens:

<|jur:in|><|doc:judgment|><|court:high|><|era:both|><|decade:2010|>  …text…  <|endoftext|>

38 reserved ids: 8 jurisdictions, 7 document types, 7 court levels, 4 statutory eras (ipc, bns, both, na), 10 decades, plus <|endoftext|> and <|pad|>.

They must be emitted literally to condition generation — writing "Bombay High Court" in a prompt does not activate <|court:high|>. Do NOT prepend the legal tags for out-of-domain prompts: measured on MMLU, prepending them costs 2.2 points (21.1% with tags against 23.3% without).

Limitations

• It has no section-conditional knowledge of punishments. This was measured directly rather than inferred. Scoring a fixed candidate set by log-probability, the model picks the same answer for a given question whether it is shown the real section or a different one's — paired test gained 0, lost 0, p = 1.0. It answers "2 years" for 32 of 40 different sections, exactly the majority-class rate. Asked whether a fine is an alternative, additional or the only punishment, it gives one answer for all 20 items. Do not use it to look up penalties.
• It will invent citations. At ~42 tokens per parameter it has not memorised case law. Generated case numbers, judge names and citations will be well-formed and frequently false. This checkpoint fabricates 4.7% of the section numbers it emits — the lowest of its series, not zero.
• In-context examples override it. A worked example in the prompt (for example "3 years") is copied as the answer on 90–100% of prompts that show one and 0% of those that do not, and this corrupts log-probability ranking too, not merely free generation. Few-shot formatting examples will be echoed.
• Not instruction-tuned. It continues text.
• Weakest of its peers on general text. Use a general model for general text.
• BNS coverage is thin, and BNSS/BSA thinner still. The Bharatiya Nyaya Sanhita reached the 12B corpus only through IPC->BNS correspondence notes appended to judgments, plus a mapping table carrying 343 of 358 BNS sections. The BNS, BNSS and BSA bare Acts are not in that corpus at all; they were added in a ~11M-token overlay shard for the second training slot, at 12x. BNSS (531 sections) and BSA (170) together are 0.015% of the 14.22B tokens the model saw. It has seen them; it cannot be said to have learned them.
• Post-2023 law is under-represented. The corpus is predominantly IPC-era.

For anything requiring correct statutory text, pair it with retrieval. Its strength is modelling the language of Indian law — where it beats every comparison model — not recalling specific provisions.

License and provenance

The training corpus contains text from gated sources (NyayaAnumana, IL-TUR) whose terms do not permit redistribution, so the corpus is not released. The weights are released for evaluation.