XtremeDistil-L6-H384 Legal Document Classifier
This model is a fine-tuned version of microsoft/xtremedistil-l6-h384-uncased for multi-class legal document classification.
It classifies European Union legal documents into 29 categories using the document text as input.
Model Details
| Property | Value |
|---|---|
| Base Model | microsoft/xtremedistil-l6-h384-uncased |
| Architecture | BERT Encoder |
| Task | Multi-class Text Classification |
| Number of Classes | 29 |
| Framework | Hugging Face Transformers |
| Fine-tuning | Full Fine-tuning |
Intended Use
The model is designed for automatic categorization of legal documents.
Example applications include:
- Legal document organization
- Regulatory document classification
- Document routing
- Legal information retrieval pipelines
- Research on legal NLP
Dataset
The model was trained on a corpus of European Union legal documents consisting of 29 document categories.
Examples include:
- Regulations
- Directives
- Decisions
- Recommendations
- Opinions
- Communications
- Other legislative and administrative legal acts
The dataset was split using an 80/20 train-validation split with a fixed random seed.
Training Configuration
Hyperparameters
| Hyperparameter | Value |
|---|---|
| Epochs | 5 |
| Learning Rate | 5e-5 |
| Batch Size | 16 |
| Gradient Accumulation | 1 |
| Mixed Precision | BF16 |
| Optimizer | AdamW |
| Evaluation Strategy | Every Epoch |
| Checkpoint Strategy | Every Epoch |
| Early Stopping | Patience = 2 |
| Best Model Selection | Validation Accuracy |
Tokenization
The original tokenizer from the base model was used.
- Maximum sequence length: 512 tokens
- Dynamic padding using
DataCollatorWithPadding - Documents longer than 512 tokens are truncated
Evaluation
| Metric | Score |
|---|---|
| Validation Accuracy | 68.59% |
| Micro F1 | 68.59% |
| Validation Loss | 2.1866 |
Usage
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification
import torch
model_name = "IndranilB/xtremedistil-l6-h384-uncased-legal-classification-hpt"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
text = "Your legal document goes here."
inputs = tokenizer(
text,
truncation=True,
return_tensors="pt",
)
with torch.no_grad():
outputs = model(**inputs)
prediction = outputs.logits.argmax(dim=-1).item()
print(prediction)
Training Pipeline
The model was trained using the Hugging Face Trainer API with:
- Cross-Entropy Loss
- Automatic Mixed Precision (BF16)
- Weights & Biases experiment tracking
- Early stopping
- Best checkpoint selection based on validation accuracy
Limitations
- Maximum supported sequence length is 512 tokens.
- Longer legal documents are truncated before inference.
- Performance may decrease when critical information occurs beyond the first 512 tokens.
Future improvements include:
- Long-context transformer models
- Hierarchical document classification
- Chunk-based inference
- Parameter-Efficient Fine-Tuning (PEFT)
Acknowledgements
This work builds upon:
- Microsoft Research's XtremeDistil
- Hugging Face Transformers
- PyTorch
- Weights & Biases