narcolepticchicken/legal-router-distilbert-v2

Legal Router DistilBERT v2

A multi-head text classification model for legal-agent routing. Given a user's legal request, it predicts:

Output	Type	Classes
route	4-class	contract_law, commercial_law, tort_law, employment_law
attorney_review_required	binary	True / False
source_required	binary	True / False

Architecture

DistilBERT-base-uncased (67M params) with 3 independent classification heads sharing the same backbone. The attorney head uses a weighted cross-entropy loss (10.5× positive weight) to handle extreme class imbalance.

Training

• Data: 5,000 labeled samples from omersx/business-legal-disputes, synthesized with keyword heuristics
• Optimizer: AdamW, lr=2e-5, weight_decay=0.01, 5 epochs
• Precision: bf16
• Hardware: NVIDIA T4 (16GB)

Results (400-sample test set)

Metric	Classifier	Prompt Baseline (SmolLM2-360M)
Route Accuracy	90.25%	~40%
Route Macro F1	0.856	~0.25
Attorney Accuracy	90.25%	~85%
Attorney F1	0.316	~0.15
Attorney Recall	81.8%	~30%
Source Accuracy	86.75%	~70%
Source F1	0.926	~0.80
Missed Escalation	18.2%	~60%
Latency (per sample)	~5ms	~2000ms
Parse Failures	0%	~15%

Key findings:

• 400× faster than the prompt baseline (5ms vs 2000ms)
• 90% route accuracy vs ~40% for zero-shot prompting
• 81.8% attorney recall — catches 9 of 11 escalation cases
• 0% parse failures — deterministic structured output vs ~15% JSON parse failures with the LLM

Limitations

• Labels are synthesized via keyword heuristics, not human-annotated
• Extreme attorney imbalance (11/400 test positive) limits attorney F1
• English legal domain only

Usage

from transformers import DistilBertTokenizerFast
import torch

ROUTE_MAP = {0: 'contract_law', 1: 'commercial_law', 2: 'tort_law', 3: 'employment_law'}

tokenizer = DistilBertTokenizerFast.from_pretrained("narcolepticchicken/legal-router-distilbert-v2")
# See model repo for full DistilBertForLegalRoutingV2 class definition
from legal_router import DistilBertForLegalRoutingV2
model = DistilBertForLegalRoutingV2.from_pretrained("narcolepticchicken/legal-router-distilbert-v2")

text = "The plaintiff alleges breach of contract regarding the supply agreement..."
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
with torch.no_grad():
    outputs = model(**inputs)
route = ROUTE_MAP[torch.argmax(outputs.route_logits, dim=-1).item()]
attorney = bool(torch.argmax(outputs.attorney_logits, dim=-1).item())
source = bool(torch.argmax(outputs.source_logits, dim=-1).item())
print(f"Route: {route}, Attorney Review: {attorney}, Source Required: {source}")

Generated by ML Intern

This model was generated by ML Intern, an agent for ML research and development on the Hugging Face Hub.