Train your own AI: Fine tune a large language model for sentence similarity

“Fine-tuning” means adapting an existing machine learning model for specific tasks or use cases. In this post I’m going to walk you through how you can fine tune a large language model for sentence similarity using some hand annotated test data. This example is in the psychology domain.

You need training data consisting of pairs of sentences, and a “ground truth” of how similar you want those sentences to be when you train your custom sentence similarity model.

We are going to train this model Harmony, a web-based tool which allows researchers in social sciences fields such as psychology to compare questionnaire items and understand the semantic similarity between those items as a percentage value. Harmony originally had an underlying large language model which is a standard multilingual LLM from HuggingFace, but this general model had some shortcomings, so we decided to fine-tune a psychology domain model.

This notebook is adapted from the scripts for the Harmony/DOXA AI competition (fine tune an LLM for the psychology domain), but it’s general and can be applied to other domains. Credit to Jeremy Lo Ying Ping for the training code.

Useful links

• Download the Jupyter notebook (you can also open it in Colab)
• Download the training data
• Cosine similarity function
• Harmony DOXA competition
• More information about sentence similarities

Getting started: installing and importing dependencies and downloading the training data

First I’m going to install and import all the libraries which we’re going to use: Pandas, and the HuggingFace libraries.

python
# !pip install pandas==2.2.2 transformers==4.43.1 sentence-transformers[train]==3.0.1

import pandas as pd

Now we can download the training data with wget.

!wget https://naturallanguageprocessing.com/harmony-matching-training-data.csv.zip

Now we’ve got our training data here, we can have a look inside:

df = pd.read_csv("harmony-matching-training-data.csv.zip")

df

Splitting into train and test sets

Now we’re going to split this into a train set and a test set.

from sklearn.model_selection import train_test_split

df_train, df_test = train_test_split(df)

df_train.reset_index(inplace=True)
df_test.reset_index(inplace=True)

df_train.drop(columns=["index"], inplace=True)
df_test.drop(columns=["index"], inplace=True)

len(df_train), len(df_test)

(1763, 588)

Converting from Pandas to HuggingFace dataset format

Now let’s import everything from Pandas dataframes into HuggingFace datasets which is what we will need for the training:

from datasets import Dataset

dataset_train = Dataset.from_pandas(df_train)

dataset_train

Dataset({
    features: ['sentence1', 'sentence2', 'score'],
    num_rows: 1763
})

dataset_test = Dataset.from_pandas(df_test)

Initialising the model that we will fine tune

We’re going to start from the sentence transformer model all-mpnet-base-v2 and fine tune it.

Although there were a number of mental health models on hugging on the HuggingFace Hub, we are going to start with a generalist model which produces embeddings in 768 dimensions.

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-mpnet-base-v2")

model

SentenceTransformer(
  (0): Transformer({'max_seq_length': 384, 'do_lower_case': False}) with Transformer model: MPNetModel 
  (1): Pooling({'word_embedding_dimension': 768, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
  (2): Normalize()
)

Let’s import the training libraries. We’re going to use the cosine similarity loss function.

from sentence_transformers import SentenceTransformerTrainer, SentenceTransformerTrainingArguments

from sentence_transformers.losses import CosineSimilarityLoss

loss = CosineSimilarityLoss(model)

Setting up the training

Let’s define our training hyperparameters:

trainer = SentenceTransformerTrainer(
    model = model,
    args = SentenceTransformerTrainingArguments(
        output_dir="checkpoints",
        num_train_epochs=3,
        per_device_eval_batch_size=16
    ),
    train_dataset=dataset_train,
    loss = loss
)

trainer

<sentence_transformers.trainer.SentenceTransformerTrainer at 0x7a081c90f810>

Running the training

The training took about 10 minutes on my laptop without a GPU.

trainer.train()

TrainOutput(global_step=663, training_loss=0.05342011286302569, metrics={'train_runtime': 553.4312, 'train_samples_per_second': 9.557, 'train_steps_per_second': 1.198, 'total_flos': 0.0, 'train_loss': 0.05342011286302569, 'epoch': 3.0})

model

SentenceTransformer(
  (0): Transformer({'max_seq_length': 384, 'do_lower_case': False}) with Transformer model: MPNetModel 
  (1): Pooling({'word_embedding_dimension': 768, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
  (2): Normalize()
)

Training has finished. Let’s encode a sentence with our model!

Now we can test the encoding by encoding a single sentence:

model.encode(["I feel nervous, anxious or afraid"])

array([[-9.93097015e-03, -7.88239613e-02,  6.90429413e-04,
        -5.47722541e-03, -7.27352547e-03,  5.51959611e-02,
         .... [abbreviated] ...
         2.22559460e-02, -4.00685295e-02, -5.54708019e-02,
         9.85449180e-03, -2.92167068e-02, -2.18849741e-02,
        -3.62964422e-02, -1.33359190e-02, -3.98904458e-02]], dtype=float32)

Embed the entire dataframe

We can take all of the two text columns in our dataframe and embed them:

sentence_1_embeddings = model.encode(df_test.sentence1)

sentence_2_embeddings = model.encode(df_test.sentence2)

sentence_1_embeddings.shape

(588, 768)

Calculate the cosine similarity

Now we’re going to use a cosine similarity function. I’ve just copy pasting this from the Harmony matching function.

from numpy import dot, matmul, ndarray, matrix
from numpy.linalg import norm
import numpy as np
def cosine_similarity(vec1: ndarray, vec2: ndarray) -> ndarray:
    dp = dot(vec1, vec2.T)
    m1 = matrix(norm(vec1, axis=1))
    m2 = matrix(norm(vec2.T, axis=0))

    return np.asarray(dp / matmul(m1.T, m2))
    

similarity_matrix = cosine_similarity(sentence_1_embeddings, sentence_2_embeddings)

similarity_matrix.shape

(588, 588)

df_test["y_pred"] = [similarity_matrix[i,i] for i in range(len(similarity_matrix))]

df_test

df_test["residual"] = df_test.y_pred - df_test.score

df_test

np.mean(df_test.residual * df_test.residual)

np.float64(0.06922512204080025)

np.mean(np.abs(df_test.residual))

np.float64(0.20813965150724928)