Fine-Tune Mistral for Legal Tasks in Under 60 Minutes

# Create a clean environment
python -m venv legal-finetune
source legal-finetune/bin/activate

# Install dependencies
pip install transformers==4.40.0 \
            peft==0.10.0 \
            trl==0.8.6 \
            datasets==2.19.0 \
            bitsandbytes==0.43.1 \
            accelerate==0.29.3 \
            wandb  # optional, for tracking

# dataset_prep.py
import json

# Each example follows this structure
example = {
    "instruction": "Identify and explain the indemnification clause in the following contract excerpt.",
    "input": "Section 12. Indemnification. The Vendor shall defend, indemnify, and hold harmless the Client from any claims arising out of Vendor's breach of this Agreement...",
    "output": "This is a one-sided indemnification clause. The Vendor bears full indemnification obligations toward the Client. Key elements: (1) Vendor must actively defend claims, not just reimburse; (2) Coverage includes breach of agreement claims only, not tort claims; (3) No cap on indemnification amount is stated — this is a significant risk for the Vendor. Recommendation: negotiate a mutual indemnification clause or add a liability cap."
}

# Format for Mistral's chat template
def format_instruction(sample):
    # Mistral uses [INST] tags for chat-style fine-tuning
    return f"[INST] {sample['instruction']}\n\n{sample['input']} [/INST] {sample['output']}"

# Load your JSONL dataset
with open("legal_data.jsonl", "r") as f:
    data = [json.loads(line) for line in f]

formatted = [format_instruction(d) for d in data]
print(f"Dataset ready: {len(formatted)} examples")
print(f"Sample:\n{formatted[0][:300]}...")

# train.py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, TaskType

MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.3"

# 4-bit quantization config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",       # NormalFloat4 — best for LLMs
    bnb_4bit_compute_dtype=torch.bfloat16,  # Faster than float16 on Ampere GPUs
    bnb_4bit_use_double_quant=True,  # Reduces memory another ~0.4 bits/param
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
tokenizer.pad_token = tokenizer.eos_token  # Mistral has no dedicated pad token

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=bnb_config,
    device_map="auto",
)

print(f"Model loaded. Memory: {model.get_memory_footprint() / 1e9:.1f} GB")

# train.py (continued)
from trl import SFTTrainer
from transformers import TrainingArguments
from datasets import Dataset

# LoRA config — rank 16 is a good default for domain adaptation
lora_config = LoraConfig(
    r=16,                    # Rank: higher = more capacity, more VRAM
    lora_alpha=32,           # Scaling factor (usually 2x rank)
    target_modules=[         # Which layers to adapt
        "q_proj", "k_proj",
        "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj"
    ],
    lora_dropout=0.05,
    bias="none",
    task_type=TaskType.CAUSAL_LM,
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# Expected: ~0.6% of params are trainable — keeps memory low

# Load formatted dataset
with open("legal_data.jsonl") as f:
    raw = [json.loads(l) for l in f]

dataset = Dataset.from_list([{"text": format_instruction(d)} for d in raw])
dataset = dataset.train_test_split(test_size=0.1)

# Training arguments
args = TrainingArguments(
    output_dir="./mistral-legal-lora",
    num_train_epochs=3,              # 3 epochs is usually enough for 500-1k examples
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,   # Effective batch size = 16
    learning_rate=2e-4,              # Standard for LoRA fine-tuning
    fp16=False,
    bf16=True,                       # Use bfloat16 on Ampere (A100, RTX 3090+)
    logging_steps=10,
    save_strategy="epoch",
    evaluation_strategy="epoch",
    load_best_model_at_end=True,
    report_to="none",                # Set to "wandb" if you want tracking
)

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    dataset_text_field="text",
    max_seq_length=2048,
    args=args,
)

trainer.train()
trainer.save_model("./mistral-legal-final")
print("Training complete.")

# merge_and_export.py
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

BASE_MODEL = "mistralai/Mistral-7B-Instruct-v0.3"
LORA_WEIGHTS = "./mistral-legal-final"

# Load base in float16 for merging (no quantization)
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float16,
    device_map="cpu",  # Merge on CPU to avoid OOM
)

model = PeftModel.from_pretrained(base_model, LORA_WEIGHTS)
merged = model.merge_and_unload()  # Bakes LoRA weights into base model

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
merged.save_pretrained("./mistral-legal-merged")
tokenizer.save_pretrained("./mistral-legal-merged")

print("Merged model saved to ./mistral-legal-merged")
print(f"Size: ~14GB in float16")

# test_inference.py
from transformers import pipeline

pipe = pipeline(
    "text-generation",
    model="./mistral-legal-merged",
    torch_dtype="auto",
    device_map="auto",
)

prompt = "[INST] Does the following clause create a unilateral or mutual obligation? Explain.\n\nThe Client shall reimburse Vendor for all reasonable out-of-pocket expenses incurred in connection with the Services. [/INST]"

result = pipe(prompt, max_new_tokens=300, do_sample=False)
print(result[0]["generated_text"])