Fine-tuning
a 7 Billion
Parameter
Legal AI

What is LegalBench
and why does it exist?

Before a single line of training code, we needed to understand the data. This is always step one — in RAG pipelines, in fine-tuning, in anything. You look at the raw material first.

LegalBench is a Stanford dataset of 162 legal reasoning tasks. We only needed 4. The first step was just opening it up and looking at what's inside — like exploring a new database before writing any queries.

from datasets import load_dataset

# Wrong path first (nguyen-brat/legalbench) → DatasetNotFoundError
# Correct path is the Stanford author's account:
ds = load_dataset("nguha/legalbench", "contract_nli_confidentiality_of_agreement")

print(ds['test'][0])
# Output:
# {'index': 0, 'answer': 'Yes',
#  'text': "Confidential Information: means all confidential
#           information disclosed by a Party...",
#  'document_name': 'Aspiegel_NDA_template.pdf'}

The data had text and answer: "Yes". Clean. Simple. But that format alone can't be used for training — we needed to convert it into something called DPO format. More on that next.

What is DPO and why
do we need it?

Raw data just says "the answer is Yes." But to actually teach a model, you need to show it two answers and say: this one is good, that one is bad. Learn the difference.

DPO stands for Direct Preference Optimisation. Every training sample needs exactly 3 things:

In our case, since all tasks were Yes/No, we made it simple: chosen = correct answer, rejected = opposite answer. The conversion function handled all 4 task types differently based on their structure.

def convert_to_dpo_format(sample, task_name):
    # contract_qa has a 'question' field, others don't
    if task_name == "contract_qa":
        prompt = f"""You are a legal expert reviewing contracts.

Clause: {sample['text']}

Question: {sample['question']}
Answer only Yes or No."""
    else:
        prompt = f"""You are a legal expert reviewing contracts.

Clause: {sample['text']}

Does this clause satisfy the legal condition?
Answer only Yes or No."""

    correct = sample["answer"].strip()
    wrong = "No" if correct == "Yes" else "Yes"

    return {
        "prompt": prompt,
        "chosen": correct,    # ✅ correct answer
        "rejected": wrong,    # ❌ wrong answer
        "task": task_name,
    }

After running this across all 4 tasks: 1560 samples total, shuffled and split 80/10/10 into train/val/test. Phase 1 done.

What is a Tokenizer?
Models don't speak English.

This came up right before we loaded the model. You know how the OpenAI API charges you by "tokens" not characters? Here's what a token actually is.

Type any text below to see exactly how it gets tokenised — how many tokens it becomes and what numbers they map to:

What is LoRA?
Training 0.13% of a model.

When Madhav asked about r and lora_alpha, this is what we covered. It's the most important concept in this whole project.

lora_config = LoraConfig(
    r=16,               # adapter size — 16 is sweet spot for T4/A10G
    lora_alpha=32,      # always set to 2× your r value
    lora_dropout=0.05,  # small regularisation
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
    # these are the attention layers — where reasoning happens
    task_type=TaskType.CAUSAL_LM,
    bias="none",
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# trainable params: 10,092,544 || all params: 7,625,709,056 || trainable%: 0.1323
# ↑ only 0.13% of the model actually trains. Rest is frozen.

Training config —
every parameter explained.

These are all the config values that confused us. Every single one explained the same way — with an example, not a definition.

Kaggle died.
Twice.

This is the part nobody writes about in tutorials. Things broke. Here's exactly what happened and how we fixed each one.

After 2 crashes, we compared GPU providers properly:

Training complete.
What the numbers mean.

After 49 minutes on Modal's A10G, training finished. Here's the actual output — and what each number actually tells us.

Despite the rewards issue, the infrastructure is complete. Data pipeline works, model loads, LoRA attaches, training runs on Modal. The next step is fixing the data quality — then re-running training and comparing against GPT-4o mini and Claude Haiku.