Fine-Tuning a Language Model for Legal Question Answering March, 2024

The legal domain is filled with complex terminology, nuanced interpretations, and the need for precise answers. That makes it a good test for natural language processing (NLP), and generic language models often struggle with the subtleties of legal text. To satisfy my curiosity, in this post I walk through fine-tuning a large language model (LLM) for legal question answering. You can find the raw dataset and the complete notebook in legal-gemma repository.

The Toolkit

  • Google Colab: The experimentation playground for GPU resources.
  • LegalBench: A curated dataset of legal questions and answers.
  • HuggingFace: Its libraries streamline working with LLMs.
  • BitsAndBytes: For memory-efficient quantization.
  • PEFT and LoRA: Techniques to selectively fine-tune parts of the model.

The Dataset: LegalBench

Before diving into code, let's understand the structure of LegalBench. According to its authors, LegalBench is "a benchmark consisting of different legal reasoning tasks. Each task has an associated dataset, consisting of input-output pairs."

Step-by-Step Guide

Project Setup

First, head to Google Colab and start a new notebook. Next, make sure you have all the following installed on the notebook: 


!pip install --upgrade pandas datasets transformers bitsandbytes peft trl huggingface_hub

Load and Transform the Data

I do not use the LegalBench dataset in the originally provided format but instead transform it from a multi-task, multi-dataset form to a question/answering dataset where the answers could be Yes/No, multiple-choice, or open-ended. You can find the [long and ugly] transformation code here. Basically, it takes one row from the original dataset and prepares a context, a question, and an answer depending on the task type. I decided to move the data transformation code to a separate Python file to improve code readability.

So, in the implementation, I start by reading the provided dataset and transforming it using the data transformation that I built.


  from Data import transform_data
  import pandas as pd

  df = pd.read_json('raw_data_sample.json')

  data = [transform_data(row) for row in df.itertuples()]
  indexes, contexts, questions, answers = zip(*data)

  assert len(contexts) == len(questions) == len(answers) and len(contexts) > 0
  input_texts = [{"text": f"Answer the Question based on the given Context.\nContext: {c}\nQuestion: {q}\nAnswer: {a}"} for c, q, a in zip(contexts, questions, answers)]

Next, I use scikit-learn to split the data into two sets: train (95%), and test (5%).


  from sklearn.model_selection import train_test_split
  from datasets import Dataset

  train_texts, test_texts = train_test_split(input_texts, test_size=0.05, random_state=7)

  train_dataset = Dataset.from_list(train_texts)
  test_dataset = Dataset.from_list(test_texts)

Prepare for Loading the Language Model

Then, I import the necessary libraries for fine-tuning the gemma-2b model. I use HuggingFace for this purpose.


  import torch

  from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
  from transformers import BitsAndBytesConfig, set_seed

  from peft import LoraConfig, PeftModel, prepare_model_for_kbit_training, get_peft_model

  from trl import SFTTrainer

  set_seed(7)
  model_name = "google/gemma-2b-it"

Since loading gemma-2b requires accepting particular terms and conditions, one needs to accept these terms and log in to the HuggingFace hub with an access token. These terms and conditions can be accepted by logging into HuggingFace. Then log in using your access token.


  from huggingface_hub import login
  login(token='YOUR_TOKEN')

Fine-tuning the Model

The next step is loading the model and the corresponding tokenizer.

Fitting the LLM on my GPU!

Since fine-tuning gemma-2b required more memory than I had available (15GB), I needed to take advantage of quantization to reduce the memory usage. Basically, BitsAndBytes quantization uses lower-precision data types to enable loading larger models.

My configuration loads the linear layers of the model with 4-bit integer precision.


  bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
  )

Loading the Model

I load the model and its corresponding tokenizer using HuggingFace. Then, I pass it to prepare_model_for_kbit_training so that the BitsAndBytes quantization that I configured earlier is applied to the loaded model.


  tokenizer = AutoTokenizer.from_pretrained(model_name, truncation=True, truncation_side = "left")
  tokenizer.pad_token = tokenizer.eos_token
  tokenizer.padding_value = tokenizer.eos_token_id
  model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb_config)
  model = prepare_model_for_kbit_training(model)

The next configuration that I found necessary for successfully fine-tuning gemma-2b within my time and hardware constraints was reducing the number of trainable parameters. Recent work has shown that even training a very small percentage of pretrained parameters can be beneficial. That is where LoraConfig comes in. The LoRA technique freezes every parameter in the pretrained model and instead trains a pair of small matrices that it inserts into each layer of the Transformer.


  peft_config = LoraConfig(
    r=8,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=['q_proj', 'up_proj', 'v_proj', 'k_proj', 'o_proj', 'down_proj', 'gate_proj']
  )
  model = get_peft_model(model, peft_config)

Then, I set the training arguments. Due to the time constraints, I experimented with only 3 different learning rates to choose the best one, and kept the rest of the hyperparameters fixed. Since the fine-tuning data is relatively small, I decided to only do a small number of training steps.


  training_arguments = TrainingArguments(
    output_dir="./results",
    do_eval=True,
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
    gradient_accumulation_steps=4,
    optim="paged_adamw_8bit",
    save_steps=100,
    logging_steps=100,
    learning_rate=1e-5,
    eval_steps=100,
    num_train_epochs=1,
    warmup_ratio=0.02,
    lr_scheduler_type="linear",
    load_best_model_at_end=True,
    save_strategy="steps",
    evaluation_strategy="steps"
  )

Next, I set up the fine-tuning trainer so that it uses the training_arguments defined above.


  trainer = SFTTrainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    dataset_text_field="text",
    tokenizer=tokenizer,
    args=training_arguments,
    peft_config=peft_config,
    max_seq_length=512
  )

Measuring the Perplexity before Fine-tuning

Before fine-tuning, I will measure the perplexity of the original (non-fine-tuned) model on the legal test set. 


  import math

  eval_results = trainer.evaluate()
  print(f"Perplexity: {math.exp(eval_results['eval_loss']):.2f}")


  Perplexity: 67.86

Let's record the model's answers to 100 randomly chosen samples from the test set. I will use them to manually evaluate the models' responses.


  from random import choices

  subset_test_dataset = list(test_dataset.select(choices(range(len(test_dataset)), k=100)))

  with open('ground_truth_subset_test.txt', 'w') as f:
    for i, entry in enumerate(subset_test_dataset):
      f.write(str(i) + "\n")
      f.write(entry['text'])
      f.write('\n====================\n')

  with open('original_subset_test.txt', 'w') as f:
    for i, entry in enumerate(subset_test_dataset):
      text = entry["text"][:entry["text"].find('\nAnswer: ')] + '\nAnswer: '
      inputs = tokenizer(text, return_tensors="pt", add_special_tokens=True)
      outputs = model.generate(**inputs, max_new_tokens=20, pad_token_id=tokenizer.eos_token_id)
      result = tokenizer.decode(outputs[0], skip_special_tokens=True)
      f.write(str(i) + "\n")
      f.write(result)
      f.write('\n====================\n')

Fine-tuning

Now, we can finally start the fine-tuning using the configurations and hyperparameters set earlier. 


  torch.cuda.empty_cache()
  model.config.use_cache = False
  trainer.train()
  model.config.use_cache = True

Save the Model


  trainer.model.save_pretrained('legal_'+model_name)

Evaluating the Perplexity of the New Model

Let's evaluate the perplexity of the fine-tuned model.


  new_eval_results = trainer.evaluate()
  print(f"Perplexity: {math.exp(new_eval_results['eval_loss']):.2f}")
  assert new_eval_results['eval_loss'] < eval_results['eval_loss']
Output: 

	Perplexity: 4.56

As expected (and desired!), the perplexity of the fine-tuned model is lower than that of the original model.

Manual Evaluation

Let's record the model's answers to a subset of the test set.


  with open('fine-tuned_subset_test.txt', 'w') as f:
    for i, entry in enumerate(subset_test_dataset):
      text = entry["text"][:entry["text"].find('\nAnswer: ')] + '\nAnswer: '
      inputs = tokenizer(text, return_tensors="pt", add_special_tokens=True)
      outputs = model.generate(**inputs, max_new_tokens=20)
      result = tokenizer.decode(outputs[0], skip_special_tokens=True)
      f.write(str(i) + "\n")
      f.write(result)
      f.write('\n====================\n')

I used the ground truth, the answers generated by the original model, and the answers generated by the fine-tuned model to perform a manual evaluation on a subset of the test set. The original model answers the given question correctly in 36% of the test cases whereas the fine-tuned model correctly answers 62% of the test cases. You can find the model's behavior for individual samples in a separate file. What surprises me is that this much improvement came from so little effort.

Conclusion and Considerations

As we saw in this simple exercise, fine-tuning an LLM can significantly boost its performance on domain-specific tasks. Keep in mind that:
  1. The results are only as good as the dataset used.
  2. Legal language changes, so we will need to update the model over time.
  3. We should be mindful of the potential biases and fairness issues in the model's output.