You can deploy this model and start earning money today!

narcolepticchicken/
legal-router-v1

Name: legal-router-v1
Availability: InStock

Views

No views yet

Legal Agent Router v1

A compact multi-task classifier that acts as the first-pass control layer for a legal AI agent. Built on SmolLM2-360M-Instruct with a custom multi-head architecture.

Model Overview

Attribute	Value
Base Model	`HuggingFaceTB/SmolLM2-360M-Instruct`
Architecture	Shared encoder + 6 classification heads
Parameters	~360M
Max Seq Length	512
Precision	bfloat16

Routing Decisions (6 parallel tasks)

Request Type (11 classes): contract_draft, contract_review, legal_research, case_analysis, litigation_strategy, regulatory_compliance, ip_question, employment_law, client_consultation, general_legal_info, out_of_scope
Attorney Review Required (binary)
Source Documents Required (binary)
Direct Answer Permitted (binary)
Tool/Lane (9 classes): contract_analyzer, case_law_retriever, statute_retriever, regulatory_checker, ip_database, drafting_engine, general_qa, human_attorney, safety_guard
Safety Flags (6 multi-label): unsafe, out_of_scope, confidentiality_risk, jurisdiction_mismatch, conflict_of_interest, urgent_deadline

Training Recipe

Dataset: Synthetic 2,200 examples (1,777 train / 423 test) with normal, hard, hidden, and adversarial splits
LR: 2e-5 (cosine schedule, 10% warmup)
Epochs: 3
Batch size: 32
Loss: Sum of CrossEntropyLoss (5 tasks) + BCEWithLogitsLoss (multi-label safety flags)

Literature Basis

Lawma (ICLR 2025): Multi-task legal classification beats commercial LLMs
SafeRoute (ACL 2025): Confidence-based routing to large models
GQR-Bench (ECAI 2025): Binary CE + thresholding for OOD/abstention

Usage

python

from transformers import AutoTokenizer
import torch
from train import LegalRouterModel

tokenizer = AutoTokenizer.from_pretrained("narcolepticchicken/legal-router-v1", trust_remote_code=True)
model = LegalRouterModel.from_pretrained("narcolepticchicken/legal-router-v1")

text = "Draft an NDA for my startup meeting with investors next week."
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
with torch.no_grad():
    outputs = model(**inputs)
    request_type = torch.argmax(outputs["logits"]["request_type"], dim=-1)
    attorney_review = torch.argmax(outputs["logits"]["attorney_review"], dim=-1)
    # ... etc

narcolepticchicken/legal-router-v1