legal-risk-classifier-lora

Legal Risk Classifier — Mistral 7B + LoRA

A LoRA fine-tuned adapter for Mistral 7B Instruct v0.3 that classifies commercial contract clauses by type and risk level, and generates a plain-English explanation — all in a single structured JSON output.

Model Description

This model takes a contract clause as input and returns:

• Clause type — one of 41 legal categories (e.g. Indemnification, Governing Law, Limitation of Liability)
• Risk level — High, Medium, or Low
• Summary — a one-sentence plain-English explanation
• Reason — why the clause carries that risk level

It was fine-tuned using LoRA (Low-Rank Adaptation) on top of the frozen Mistral 7B Instruct v0.3 base model, trained on the CUAD (Contract Understanding Atticus Dataset) legal clause dataset.

Why a decoder model instead of an encoder (e.g. BERT)?

Encoder-only models output a single label. This task requires three simultaneous outputs — classification, risk assessment, and natural-language explanation — in one coherent response. A decoder-based generative model produces all three in a single forward pass. Mistral 7B was additionally pre-trained on a broad corpus that includes legal text, giving it strong zero-shot legal reasoning before any fine-tuning was applied.

Example Output

{
  "clause_type": "Indemnification",
  "risk_level": "High",
  "summary": "You are responsible for all losses the other party suffers.",
  "reason": "Broad unlimited indemnification — no liability cap specified."
}

Training Details

Training Data

CUAD — an expert-annotated legal dataset produced by the Atticus Project. 510 real commercial contracts annotated by practising lawyers and supervised law students across 41 clause categories.

• Raw examples: 84,000+ question–answer pairs
• After filtering to positive (non-empty) examples: ~4,000 genuine clause instances
• Train / test split: 95% / 5%, shuffled with seed 42 for reproducibility

Risk-level labels (High / Medium / Low) are not present in CUAD and were added as part of this project, based on legal exposure principles:

Unmapped clause types default to Medium as a conservative measure.

Training Procedure

• Base model: mistralai/Mistral-7B-Instruct-v0.3
• Method: LoRA via HuggingFace PEFT
• Hardware: NVIDIA DGX Spark (GB10, Blackwell architecture, 128GB unified memory)
• Precision: BF16 (full precision — no quantisation)
• LoRA rank (r): 32
• LoRA alpha: 64
• Target modules: q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj
• Epochs: 3
• Batch size: 4 (gradient accumulation 4 → effective batch size 16)
• Learning rate: 2e-4
• Trainable parameters: ~85M of 7.2B total (1.2%)
• Adapter size: ~150MB

Evaluation

Evaluated on 454 held-out test examples never seen during training (High: 78, Medium: 325, Low: 51).

Results — Before vs After Fine-Tuning

Comparison With Published Baselines

Key Findings

• Low Risk Recall improved by 330% relative (19.6% → 84.3%) — the most dramatic improvement. Before fine-tuning the base model rarely identified genuinely low-risk clauses, which in production creates alert fatigue. After fine-tuning it can confidently clear low-risk clauses without unnecessary human review.
• JSON reliability reached 100% with zero inference failures, critical for any production API integration.
• Zero-shot Mistral 7B already exceeded the best published zero-shot result in ContractEval 2025 (~52% F1), suggesting strong innate legal reasoning from pre-training alone.
• High Risk Recall remains the hardest category (67.95%), reflecting class imbalance in CUAD (78 High Risk vs 325 Medium Risk test examples). In production, all High Risk classifications should be routed to mandatory human review regardless of model confidence.

How to Use

from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

base_model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-Instruct-v0.3",
    torch_dtype="bfloat16",
    device_map="auto"
)
model = PeftModel.from_pretrained(base_model, "Govardhan12345/legal-risk-classifier-lora")
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")

prompt = """<s>[INST] You are a legal risk analyst specialising in commercial contracts.

Analyse the following contract clause:
The contractor shall indemnify and hold harmless the employer from all claims.

Respond ONLY with a valid JSON object:
{
  "clause_type": "...",
  "risk_level": "High or Medium or Low",
  "summary": "...",
  "reason": "..."
} [/INST]"""

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

Intended Use

This model is designed as a decision-support tool for identifying and summarising contract risk — it is not a substitute for legal advice. High Risk classifications and edge cases should always be reviewed by a qualified professional before any contractual decision is made.

Limitations

• Trained on CUAD, which consists of US commercial contracts — coverage of other jurisdictions or specialised contract types (e.g. UK construction subcontracts) may be limited without further domain-specific fine-tuning.
• High Risk Recall (67.95%) is lower than Medium and Low Risk due to class imbalance in the training data.
• Outputs should be treated as a first-pass risk flag, not a final legal determination.

Framework Versions

• PEFT 0.12.0
• Transformers
• BF16 training on NVIDIA DGX Spark (GB10, Blackwell)