Tirth2109/legal-slm

Legal SLM (400M)

A highly specialized, 400-million parameter Small Language Model (SLM) trained from scratch to understand and reason over legal frameworks, specifically optimized for the Indian Penal Code (IPC) and international legal corpora.

Unlike standard statistical language models, this architecture integrates Knowledge Graph (KG) Contrastive Alignment. By physically anchoring legal entities (Subjects, Predicates, Objects) in its vector space during pre-training, the model demonstrates high hallucination resistance and structural reasoning capabilities natively suited for Retrieval-Augmented Generation (RAG).

Model Details

• Architecture: Custom Llama-style Causal Decoder (LegalSLMForCausalLM)
• Parameters: ~400M
• Layers: 24
• Hidden Size: 1024
• Attention Heads: 16
• Context Window: 1024 tokens
• Vocab Size: 48,000
• Activations: SwiGLU FFNs with RMSNorm

Usage

Because this model uses a custom architecture, you must pass trust_remote_code=True to allow Hugging Face to load the proprietary modeling_legalslm.py and configuration_legalslm.py files included in this repository.

Installation

pip install transformers torch

Loading the Model

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

repo_id = "your-username/legal-slm"  # Replace with your HF repo ID

tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(repo_id, trust_remote_code=True)

# Ensure padding token is set for batched generation
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

Prompting Structure

The model was aligned using a blended Supervised Fine-Tuning (SFT) approach. It supports both open-ended conversational queries and strict context-grounded RAG tasks.

1. Context-Grounded RAG (Recommended)

prompt = """<|user|>
INSTRUCTION: You are a legal assistant. If the Context below is tagged [CONTEXT: RELEVANT], answer the question STRICTLY using those facts. If it is tagged [CONTEXT: NOT FOUND], say plainly that you don't have information on this. Do not invent information in either case.

Context:
[CONTEXT: RELEVANT]
<Section 420> <defines> <Cheating and delivery dishonestly inducing of property>

Question: What legal offense does Section 420 define?
<|assistant|>
"""

inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=50, temperature=0.1)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

2. Open Conversational Query

prompt = "<|user|>\nWhat are the provisions for theft under the IPC?\n<|assistant|>\n"
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=100, temperature=0.6)

Training Pipeline

The model was developed in two primary phases:

Phase 1: Dual-Objective Pre-Training

• Dataset: 9.4 Billion tokens (75% IPC/Indian Legal, 25% Global Legal).
• Technique: Standard Next-Token Prediction combined with a Contrastive Projection Head. Synthetic Knowledge Graph triples (e.g., <subj> Section 420 </subj> <pred> penalty is </pred> <obj> ... </obj>) were injected every ~1,500 tokens. The contrastive head structurally mapped these entities closer together in the vector space, establishing robust relational memory.
• Hardware: NVIDIA V100 GPU.

Phase 2: Blended Supervised Fine-Tuning (SFT)

• Dataset: A blended mixture of Techmaestro369/indian-legal-texts-finetuning (Conversational QA) and synthetic GraphRAG contextual prompts.
• Technique: Low-learning-rate alignment to prevent catastrophic forgetting of conversational abilities while enforcing strict adherence to provided <Context> blocks.

Limitations & Ethical Considerations

• Not Legal Counsel: This model is an experimental AI research artifact. It does not provide certified legal advice. Always consult a qualified attorney for real-world legal matters.
• Jurisdictional Bias: The weights are heavily skewed toward the Indian Penal Code (IPC) and related Indian frameworks. While it possesses baseline syntactic understanding of global law, it may confidently misapply Indian legal principles to foreign jurisdictions unless explicitly grounded via RAG context.
• Early Training State: As a 400M parameter model trained within strict compute constraints, generative fluency (grammar, repetition) may occasionally degrade on highly complex, multi-hop reasoning prompts.