LegalBERT Multi-Label Posture Classification

This repository implementeda fine-tuned LegalBERT model for multi-label posture classification of US judicial opinions.

📋 Table of Contents

• Project Overview
• Environment Setup
• Project Structure
• Data Requirements
• Quick Start
• Detailed Usage
• Model Architecture
• Output Artifacts

🎯 Project Overview

This solution addresses the Thomson Reuters Data Science Challenge by:

1. Question 1: Data exploration and basic statistics reporting
2. Question 2: Building an exploratory multi-label classification model
3. Question 3: Providing recommendations for next steps

The model fine-tunes nlpaueb/legal-bert-base-uncased using LoRA (Low-Rank Adaptation) for parameter-efficient training on the top 10 most frequent procedural postures plus an "Other" category.

🔧 Environment Setup

Prerequisites

• Python: 3.11+ (tested on 3.11 & 3.12)
• Operating System: macOS (Apple Silicon/Intel), Linux, Windows
• Memory: Minimum 8GB RAM, recommended 16GB+
• Storage: At least 5GB free space

Step 1: Create Virtual Environment

# Using conda (recommended)
conda create -n tr-challenge python=3.11
conda activate tr-challenge

Step 2: Install Dependencies

pip install -r requirements.txt

📁 Project Structure

ThomsonReutersChallenge/
├── README.md                           # This file
├── requirements.txt                    # Python dependencies
├── run.sh                             # Quick execution script
├── .gitignore                         # Git ignore rules
├── data/                              # Data directory
│   └── TRDataChallenge2023.txt       # Main dataset (JSONL format)
│
├── config/                            # Configuration files
│   └── postures.json                 # Top 10 postures with kappa values
│
├── src/                               # Source code
│   ├── __init__.py                   # Package initialization
│   ├── main.py                       # Main execution script
│   ├── data_processor.py             # Data loading and preprocessing
│   ├── model.py                      # LegalBERT model with LoRA
│   └── trainer.py                    # Training and evaluation logic
│
├── outputs/                           # Generated outputs
│   ├── *.png                        # Data exploration plots
│   └── legalbert_multi/             # Model artifacts
│       ├── tokenizer/               # Tokenizer files
│       ├── pytorch_model.bin        # Trained model weights
│       ├── thresholds.json          # Optimized classification thresholds
│       ├── metrics.json             # Evaluation metrics
│       └── plots/                   # Training and evaluation plots
│
└── TR_Data_Challenge_Exploratory_Analysis.ipynb  # Jupyter notebook with exploratory analysis on data (Question 1)

📊 Data Requirements

Expected Data Files

1. data/TRDataChallenge2023.txt: Main dataset in JSONL format

   • Each line contains a JSON object with documentId, postures, and sections
   • File size: ~500MB
   • Format: UTF-8 encoded

2. config/postures.json: Configuration file with top 10 postures and their kappa values

   • Already included in the repository
   • Contains inter-annotator agreement scores

Data Format Example

{
  "documentId": "doc_12345",
  "postures": ["On Appeal (Civil)", "Motion to Dismiss"],
  "sections": [
    {
      "headtext": "Introduction",
      "paragraphs": ["This case involves...", "The plaintiff argues..."]
    }
  ]
}

🚀 Quick Start

Jupyter Notebook Analysis (For Challenge Question 1)

# Start Jupyter
jupyter notebook

# Open TR_Data_Challenge_Exploratory_Analysis.ipynb
# Run all cells for complete analysis

Model Training Option 1: Use the Run Script (For Challenge Question 2)

# Make script executable
chmod +x run.sh

# Run the complete pipeline
./run.sh

Model Training Option 2: Manual Execution (For Challenge Question 2)

# Train from scratch
python src/main.py \
    --data_path ./data/TRDataChallenge2023.txt \
    --postures_path ./config/postures.json \
    --output_dir ./outputs/my_model

# Evaluate existing model
python src/main.py \
    --do_eval \
    --model_dir ./outputs/legalbert_multi \
    --data_path ./data/TRDataChallenge2023.txt \
    --postures_path ./config/postures.json

🔍 Detailed Usage

Training Configuration

The model supports extensive customization through command-line arguments:

python src/main.py \
    --do_train \
    --data_path data/TRDataChallenge2023.txt \
    --postures_path config/postures.json \
    --output_dir outputs/custom_model \
    --model_name nlpaueb/legal-bert-base-uncased \
    --max_chunk_length 512 \
    --stride 128 \
    --lora_r 8 \
    --lora_alpha 32 \
    --batch_size 2 \
    --grad_accum_steps 16 \
    --learning_rate 3e-4 \
    --num_epochs 10 \
    --early_stop_patience 3 \
    --random_seed 42

Key Parameters

• --max_chunk_length: Maximum tokens per chunk (default: 512)
• --stride: Overlap between chunks (default: 128)
• --lora_r: LoRA rank parameter (default: 8)
• --lora_alpha: LoRA alpha parameter (default: 32)
• --batch_size: Per-device batch size (default: 4 for Mac M4 and 32 for NVIDIA RTX)
• --grad_accum_steps: Gradient accumulation steps (default: 32 for Mac M4 and 4 for NVIDIA RTX)

Continuing Training

python src/main.py \
    --do_train \
    --continue_training \
    --model_dir ./outputs/legalbert_multi \
    --output_dir ./outputs/legalbert_multi_continued \
    --num_epochs 3

🏗️ Model Architecture

Base Model

• Foundation: nlpaueb/legal-bert-base-uncased
• Domain: Legal text (pre-trained on legal documents)
• Size: 110M parameters

Efficient Fine-tuning

• Method: LoRA (Low-Rank Adaptation)
• Trainable Parameters: ~2% of total model parameters
• Target Modules: Query, Key, Value projection layers
• LoRA Config: r=8, α=32, dropout=0.1

Text Processing Pipeline

1. Document Chunking: Split long documents into ≤512 token chunks with 128-token stride
2. Section Formatting: [HEAD] {headtext} [SEP] {paragraph_text}
3. Attention Pooling: Learnable attention mechanism for document-level aggregation
4. Classification Head: 768 → 256 → 11 dimensions with Sigmoid activation

Training Strategy

• Loss Function: Binary Cross-Entropy with Logits
• Class Weighting: Kappa-based positive sample weighting
• Optimization: AdamW (lr=3e-4, weight_decay=0.01)
• Mixed Precision: bf16 on supported hardware
• Early Stopping: Based on validation macro-F1 score

📈 Output Artifacts

Model Files

• outputs/legalbert_multi/pytorch_model.bin: Trained model weights
• outputs/legalbert_multi/tokenizer/: Tokenizer configuration
• outputs/legalbert_multi/thresholds.json: Optimized classification thresholds

Evaluation Results

• outputs/legalbert_multi/metrics.json: Comprehensive evaluation metrics
• outputs/legalbert_multi/plots/: Training curves and performance visualizations

Data Analysis

• outputs/*.png: Exploratory data analysis plots
• Dataset statistics and posture distribution visualizations

📝 Reproducibility Notes

Tested Environments

• ✅ macOS 13+ (Apple Silicon M4)
• ✅ WSL Ubuntu 22.04+ (CUDA 12.4+)

Fixed Random Seeds

All experiments use random_seed=42 for reproducibility across:

• Data splitting (70/15/15 train/val/test)
• Model initialization
• Training procedures

Version Pinning

All dependencies are version-pinned in requirements.txt to ensure consistent results.

Expected Results

With the default configuration, you should achieve:

• Macro-F1: ~0.65-0.75 on validation set
• Micro-F1: ~0.70-0.80 on validation set
• Training Time: 30-120 minutes on modern GPU