metadata
tags:
- distilbert
- legal
- text-classification
- routing
- ml-intern
license: apache-2.0
datasets:
- omersx/business-legal-disputes
metrics:
- accuracy
- f1
model-index:
- name: legal-router-distilbert-v2
results:
- task:
type: text-classification
name: Legal Agent Routing
dataset:
name: omersx/business-legal-disputes
type: legal-disputes
split: test
metrics:
- type: route_accuracy
value: 0.9025
- type: route_macro_f1
value: 0.8561
- type: attorney_f1
value: 0.3158
- type: attorney_recall
value: 0.8182
- type: source_f1
value: 0.9259
- type: missed_escalation_rate
value: 0.1818
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.