---
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

```python
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}")
```

<!-- ml-intern-provenance -->
## Generated by ML Intern

This model was generated by [ML Intern](https://github.com/huggingface/ml-intern), an agent for ML research and development on the Hugging Face Hub.