Chapter 10 - Creating Text Embedding Models

Creating an Embedding Model

Data

from datasets import load_dataset

# Load MNLI dataset from GLUE
# 0 = entailment, 1 = neutral, 2 = contradiction
train_dataset = load_dataset("glue", "mnli", split="train").select(range(50_000))
train_dataset = train_dataset.remove_columns("idx")
train_dataset[2]
{'premise': 'One of our number will carry out your instructions minutely.',
 'hypothesis': 'A member of my team will execute your orders with immense precision.',
 'label': 0}

Model

from sentence_transformers import SentenceTransformer

# Use a base model
embedding_model = SentenceTransformer('bert-base-uncased')
WARNING:sentence_transformers.SentenceTransformer:No sentence-transformers model found with name bert-base-uncased. Creating a new one with mean pooling.
/usr/local/lib/python3.10/dist-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.
  warnings.warn(

Loss Function

from sentence_transformers import losses

# Define the loss function. In soft-max loss, we will also need to explicitly set the number of labels.
train_loss = losses.SoftmaxLoss(
    model=embedding_model,
    sentence_embedding_dimension=embedding_model.get_sentence_embedding_dimension(),
    num_labels=3
)

Evaluation

from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator

# Create an embedding similarity evaluator for stsb
val_sts = load_dataset('glue', 'stsb', split='validation')
evaluator = EmbeddingSimilarityEvaluator(
    sentences1=val_sts["sentence1"],
    sentences2=val_sts["sentence2"],
    scores=[score/5 for score in val_sts["label"]],
    main_similarity="cosine",
)

Training

from sentence_transformers.training_args import SentenceTransformerTrainingArguments

# Define the training arguments
args = SentenceTransformerTrainingArguments(
    output_dir="base_embedding_model",
    num_train_epochs=1,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    warmup_steps=100,
    fp16=True,
    eval_steps=100,
    logging_steps=100,
)
from sentence_transformers.trainer import SentenceTransformerTrainer

# Train embedding model
trainer = SentenceTransformerTrainer(
    model=embedding_model,
    args=args,
    train_dataset=train_dataset,
    loss=train_loss,
    evaluator=evaluator
)
trainer.train()
[1563/1563 06:10, Epoch 1/1]
Step Training Loss
100 1.080700
200 0.959400
300 0.916200
400 0.870200
500 0.849100
600 0.854200
700 0.835200
800 0.825200
900 0.818100
1000 0.800300
1100 0.781600
1200 0.777100
1300 0.786600
1400 0.767900
1500 0.797100 

TrainOutput(global_step=1563, training_loss=0.8453957184872716, metrics={'train_runtime': 372.5713, 'train_samples_per_second': 134.202, 'train_steps_per_second': 4.195, 'total_flos': 0.0, 'train_loss': 0.8453957184872716, 'epoch': 1.0})
# Evaluate our trained model
evaluator(embedding_model)
{'pearson_cosine': 0.3710938716460552,
 'spearman_cosine': 0.45148122260403883,
 'pearson_manhattan': 0.4037396904694362,
 'spearman_manhattan': 0.4396893995197567,
 'pearson_euclidean': 0.390788259199341,
 'spearman_euclidean': 0.43444104358464286,
 'pearson_dot': 0.3392927926047231,
 'spearman_dot': 0.3530708415227247,
 'pearson_max': 0.4037396904694362,
 'spearman_max': 0.45148122260403883}

MTEB

from mteb import MTEB

# Choose evaluation task
evaluation = MTEB(tasks=["Banking77Classification"])

# Calculate results
results = evaluation.run(embedding_model)
results
───────────────────────────────────────────────── Selected tasks  ─────────────────────────────────────────────────

Classification

    - Banking77Classification, s2s







/usr/local/lib/python3.10/dist-packages/joblib/externals/loky/backend/fork_exec.py:38: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
  pid = os.fork()






{'Banking77Classification': {'mteb_version': '1.1.2',
  'dataset_revision': '0fd18e25b25c072e09e0d92ab615fda904d66300',
  'mteb_dataset_name': 'Banking77Classification',
  'test': {'accuracy': 0.46022727272727276,
   'f1': 0.45802738001849663,
   'accuracy_stderr': 0.009556987908238961,
   'f1_stderr': 0.01072225943077292,
   'main_score': 0.46022727272727276,
   'evaluation_time': 29.63}}}






VRAM Clean-up - You will need to run the code below to partially empty the VRAM (GPU RAM). If that does not work, it is advised to restart the notebook instead. You can also run !nivia-smi to check current usage.




# # Empty and delete trainer/model
# trainer.accelerator.clear()
# del trainer, embedding_model

# # Garbage collection and empty cache
# import gc
# import torch

# gc.collect()
# torch.cuda.empty_cache()






import gc
import torch

gc.collect()
torch.cuda.empty_cache()








Loss Fuctions


⚠️ VRAM Clean-up * Restart the notebook in order to clean-up memory if you move on to the next training example.




Cosine Similarity Loss




from datasets import Dataset, load_dataset

# Load MNLI dataset from GLUE
# 0 = entailment, 1 = neutral, 2 = contradiction
train_dataset = load_dataset("glue", "mnli", split="train").select(range(50_000))
train_dataset = train_dataset.remove_columns("idx")

# (neutral/contradiction)=0 and (entailment)=1
mapping = {2: 0, 1: 0, 0:1}
train_dataset = Dataset.from_dict({
    "sentence1": train_dataset["premise"],
    "sentence2": train_dataset["hypothesis"],
    "label": [float(mapping[label]) for label in train_dataset["label"]]
})






from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator

# Create an embedding similarity evaluator for stsb
val_sts = load_dataset('glue', 'stsb', split='validation')
evaluator = EmbeddingSimilarityEvaluator(
    sentences1=val_sts["sentence1"],
    sentences2=val_sts["sentence2"],
    scores=[score/5 for score in val_sts["label"]],
    main_similarity="cosine"
)






from sentence_transformers import losses, SentenceTransformer
from sentence_transformers.trainer import SentenceTransformerTrainer
from sentence_transformers.training_args import SentenceTransformerTrainingArguments

# Define model
embedding_model = SentenceTransformer('bert-base-uncased')

# Loss function
train_loss = losses.CosineSimilarityLoss(model=embedding_model)

# Define the training arguments
args = SentenceTransformerTrainingArguments(
    output_dir="cosineloss_embedding_model",
    num_train_epochs=1,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    warmup_steps=100,
    fp16=True,
    eval_steps=100,
    logging_steps=100,
)

# Train model
trainer = SentenceTransformerTrainer(
    model=embedding_model,
    args=args,
    train_dataset=train_dataset,
    loss=train_loss,
    evaluator=evaluator
)
trainer.train()




WARNING:sentence_transformers.SentenceTransformer:No sentence-transformers model found with name bert-base-uncased. Creating a new one with mean pooling.
/usr/local/lib/python3.10/dist-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.
  warnings.warn(






    

      
      
      [1563/1563 06:04, Epoch 1/1]
    

    





Step
Training Loss






100
0.231900




200
0.168900




300
0.170900




400
0.157800




500
0.152900




600
0.156100




700
0.149300




800
0.154500




900
0.150900




1000
0.145600




1100
0.147800




1200
0.145600




1300
0.145100




1400
0.142000




1500
0.141600



























TrainOutput(global_step=1563, training_loss=0.15676780793427353, metrics={'train_runtime': 364.4779, 'train_samples_per_second': 137.183, 'train_steps_per_second': 4.288, 'total_flos': 0.0, 'train_loss': 0.15676780793427353, 'epoch': 1.0})








# Evaluate our trained model
evaluator(embedding_model)




{'pearson_cosine': 0.7222320710908391,
 'spearman_cosine': 0.725059765496038,
 'pearson_manhattan': 0.7338172618636865,
 'spearman_manhattan': 0.7323465534428775,
 'pearson_euclidean': 0.7332726423686017,
 'spearman_euclidean': 0.7316943270141215,
 'pearson_dot': 0.6603672299249149,
 'spearman_dot': 0.6624301208511642,
 'pearson_max': 0.7338172618636865,
 'spearman_max': 0.7323465534428775}






⚠️ VRAM Clean-up * Restart the notebook in order to clean-up memory if you move on to the next training example.




import gc
import torch

gc.collect()
torch.cuda.empty_cache()








Multiple Negatives Ranking Loss




import random
from tqdm import tqdm
from datasets import Dataset, load_dataset

# # Load MNLI dataset from GLUE
mnli = load_dataset("glue", "mnli", split="train").select(range(50_000))
mnli = mnli.remove_columns("idx")
mnli = mnli.filter(lambda x: True if x['label'] == 0 else False)

# Prepare data and add a soft negative
train_dataset = {"anchor": [], "positive": [], "negative": []}
soft_negatives = mnli["hypothesis"]
random.shuffle(soft_negatives)
for row, soft_negative in tqdm(zip(mnli, soft_negatives)):
    train_dataset["anchor"].append(row["premise"])
    train_dataset["positive"].append(row["hypothesis"])
    train_dataset["negative"].append(soft_negative)
train_dataset = Dataset.from_dict(train_dataset)
len(train_dataset)




16875it [00:01, 14110.96it/s]






16875








from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator

# Create an embedding similarity evaluator for stsb
val_sts = load_dataset('glue', 'stsb', split='validation')
evaluator = EmbeddingSimilarityEvaluator(
    sentences1=val_sts["sentence1"],
    sentences2=val_sts["sentence2"],
    scores=[score/5 for score in val_sts["label"]],
    main_similarity="cosine"
)






from sentence_transformers import losses, SentenceTransformer
from sentence_transformers.trainer import SentenceTransformerTrainer
from sentence_transformers.training_args import SentenceTransformerTrainingArguments

# Define model
embedding_model = SentenceTransformer('bert-base-uncased')

# Loss function
train_loss = losses.MultipleNegativesRankingLoss(model=embedding_model)

# Define the training arguments
args = SentenceTransformerTrainingArguments(
    output_dir="mnrloss_embedding_model",
    num_train_epochs=1,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    warmup_steps=100,
    fp16=True,
    eval_steps=100,
    logging_steps=100,
)

# Train model
trainer = SentenceTransformerTrainer(
    model=embedding_model,
    args=args,
    train_dataset=train_dataset,
    loss=train_loss,
    evaluator=evaluator
)
trainer.train()




WARNING:sentence_transformers.SentenceTransformer:No sentence-transformers model found with name bert-base-uncased. Creating a new one with mean pooling.
/usr/local/lib/python3.10/dist-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.
  warnings.warn(






    

      
      
      [528/528 03:00, Epoch 1/1]
    

    





Step
Training Loss






100
0.345200




200
0.105500




300
0.079000




400
0.062200




500
0.069000

















TrainOutput(global_step=528, training_loss=0.12795479415041028, metrics={'train_runtime': 180.866, 'train_samples_per_second': 93.301, 'train_steps_per_second': 2.919, 'total_flos': 0.0, 'train_loss': 0.12795479415041028, 'epoch': 1.0})








# Evaluate our trained model
evaluator(embedding_model)




{'pearson_cosine': 0.8070727434643791,
 'spearman_cosine': 0.8106193672462586,
 'pearson_manhattan': 0.8213132116968124,
 'spearman_manhattan': 0.8164551132664518,
 'pearson_euclidean': 0.820988086354926,
 'spearman_euclidean': 0.8160139830687847,
 'pearson_dot': 0.7429357515240518,
 'spearman_dot': 0.7316164586329814,
 'pearson_max': 0.8213132116968124,
 'spearman_max': 0.8164551132664518}












Fine-tuning


⚠️ VRAM Clean-up * Restart the notebook in order to clean-up memory if you move on to the next training example.




import gc
import torch

gc.collect()
torch.cuda.empty_cache()






Supervised




from datasets import load_dataset
from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator

# Load MNLI dataset from GLUE
# 0 = entailment, 1 = neutral, 2 = contradiction
train_dataset = load_dataset("glue", "mnli", split="train").select(range(50_000))
train_dataset = train_dataset.remove_columns("idx")

# Create an embedding similarity evaluator for stsb
val_sts = load_dataset('glue', 'stsb', split='validation')
evaluator = EmbeddingSimilarityEvaluator(
    sentences1=val_sts["sentence1"],
    sentences2=val_sts["sentence2"],
    scores=[score/5 for score in val_sts["label"]],
    main_similarity="cosine"
)






from sentence_transformers import losses, SentenceTransformer
from sentence_transformers.trainer import SentenceTransformerTrainer
from sentence_transformers.training_args import SentenceTransformerTrainingArguments

# Define model
embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

# Loss function
train_loss = losses.MultipleNegativesRankingLoss(model=embedding_model)

# Define the training arguments
args = SentenceTransformerTrainingArguments(
    output_dir="finetuned_embedding_model",
    num_train_epochs=1,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    warmup_steps=100,
    fp16=True,
    eval_steps=100,
    logging_steps=100,
)

# Train model
trainer = SentenceTransformerTrainer(
    model=embedding_model,
    args=args,
    train_dataset=train_dataset,
    loss=train_loss,
    evaluator=evaluator
)
trainer.train()
























/usr/local/lib/python3.10/dist-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.
  warnings.warn(









































    

      
      
      [1563/1563 01:57, Epoch 1/1]
    

    





Step
Training Loss






100
0.155500




200
0.110000




300
0.118600




400
0.115300




500
0.110700




600
0.101000




700
0.113100




800
0.099800




900
0.109600




1000
0.105800




1100
0.094900




1200
0.106400




1300
0.105300




1400
0.105200




1500
0.106600



























TrainOutput(global_step=1563, training_loss=0.10982195932897176, metrics={'train_runtime': 117.3739, 'train_samples_per_second': 425.989, 'train_steps_per_second': 13.316, 'total_flos': 0.0, 'train_loss': 0.10982195932897176, 'epoch': 1.0})








# Evaluate our trained model
evaluator(embedding_model)




{'pearson_cosine': 0.8489503881223601,
 'spearman_cosine': 0.8484667083117318,
 'pearson_manhattan': 0.8503843871673679,
 'spearman_manhattan': 0.8475679105384369,
 'pearson_euclidean': 0.8513072191805562,
 'spearman_euclidean': 0.8484667083117318,
 'pearson_dot': 0.8489503890256918,
 'spearman_dot': 0.8484667083117318,
 'pearson_max': 0.8513072191805562,
 'spearman_max': 0.8484667083117318}








# Evaluate the pre-trained model
original_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
evaluator(original_model)




/usr/local/lib/python3.10/dist-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.
  warnings.warn(






{'pearson_cosine': 0.8696194608752055,
 'spearman_cosine': 0.8671637433378804,
 'pearson_manhattan': 0.8670399009851635,
 'spearman_manhattan': 0.8663946139224048,
 'pearson_euclidean': 0.867871599362501,
 'spearman_euclidean': 0.8671643653432983,
 'pearson_dot': 0.8696194616795601,
 'spearman_dot': 0.8671631197908374,
 'pearson_max': 0.8696194616795601,
 'spearman_max': 0.8671643653432983}






⚠️ VRAM Clean-up * Restart the notebook in order to clean-up memory if you move on to the next training example.




import gc
import torch

gc.collect()
torch.cuda.empty_cache()








Augmented SBERT


Step 1: Fine-tune a cross-encoder




import pandas as pd
from tqdm import tqdm
from datasets import load_dataset, Dataset
from sentence_transformers import InputExample
from sentence_transformers.datasets import NoDuplicatesDataLoader

# Prepare a small set of 10000 documents for the cross-encoder
dataset = load_dataset("glue", "mnli", split="train").select(range(10_000))
mapping = {2: 0, 1: 0, 0:1}

# Data Loader
gold_examples = [
    InputExample(texts=[row["premise"], row["hypothesis"]], label=mapping[row["label"]])
    for row in tqdm(dataset)
]
gold_dataloader = NoDuplicatesDataLoader(gold_examples, batch_size=32)

# Pandas DataFrame for easier data handling
gold = pd.DataFrame(
    {
    'sentence1': dataset['premise'],
    'sentence2': dataset['hypothesis'],
    'label': [mapping[label] for label in dataset['label']]
    }
)




100%|██████████| 10000/10000 [00:00<00:00, 25870.92it/s]








from sentence_transformers.cross_encoder import CrossEncoder

# Train a cross-encoder on the gold dataset
cross_encoder = CrossEncoder('bert-base-uncased', num_labels=2)
cross_encoder.fit(
    train_dataloader=gold_dataloader,
    epochs=1,
    show_progress_bar=True,
    warmup_steps=100,
    use_amp=False
)




/usr/local/lib/python3.10/dist-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.
  warnings.warn(
Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['classifier.bias', 'classifier.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
















Step 2: Create new sentence pairs




# Prepare the silver dataset by predicting labels with the cross-encoder
silver = load_dataset("glue", "mnli", split="train").select(range(10_000, 50_000))
pairs = list(zip(silver['premise'], silver['hypothesis']))




Step 3: Label new sentence pairs with the fine-tuned cross-encoder (silver dataset)




import numpy as np

# Label the sentence pairs using our fine-tuned cross-encoder
output = cross_encoder.predict(pairs, apply_softmax=True, show_progress_bar=True)
silver = pd.DataFrame(
    {
        "sentence1": silver["premise"],
        "sentence2": silver["hypothesis"],
        "label": np.argmax(output, axis=1)
    }
)




Step 4: Train a bi-encoder (SBERT) on the extended dataset (gold + silver dataset)




# Combine gold + silver
data = pd.concat([gold, silver], ignore_index=True, axis=0)
data = data.drop_duplicates(subset=['sentence1', 'sentence2'], keep="first")
train_dataset = Dataset.from_pandas(data, preserve_index=False)






from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator

# Create an embedding similarity evaluator for stsb
val_sts = load_dataset('glue', 'stsb', split='validation')
evaluator = EmbeddingSimilarityEvaluator(
    sentences1=val_sts["sentence1"],
    sentences2=val_sts["sentence2"],
    scores=[score/5 for score in val_sts["label"]],
    main_similarity="cosine"
)






from sentence_transformers import losses, SentenceTransformer
from sentence_transformers.trainer import SentenceTransformerTrainer
from sentence_transformers.training_args import SentenceTransformerTrainingArguments

# Define model
embedding_model = SentenceTransformer('bert-base-uncased')

# Loss function
train_loss = losses.CosineSimilarityLoss(model=embedding_model)

# Define the training arguments
args = SentenceTransformerTrainingArguments(
    output_dir="augmented_embedding_model",
    num_train_epochs=1,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    warmup_steps=100,
    fp16=True,
    eval_steps=100,
    logging_steps=100,
)

# Train model
trainer = SentenceTransformerTrainer(
    model=embedding_model,
    args=args,
    train_dataset=train_dataset,
    loss=train_loss,
    evaluator=evaluator
)
trainer.train()






# Evaluate our trained model
evaluator(embedding_model)






trainer.accelerator.clear()




Step 5: Evaluate without silver dataset




# Combine gold + silver
data = pd.concat([gold], ignore_index=True, axis=0)
data = data.drop_duplicates(subset=['sentence1', 'sentence2'], keep="first")
train_dataset = Dataset.from_pandas(data, preserve_index=False)

# Define model
embedding_model = SentenceTransformer('bert-base-uncased')

# Loss function
train_loss = losses.CosineSimilarityLoss(model=embedding_model)

# Define the training arguments
args = SentenceTransformerTrainingArguments(
    output_dir="gold_only_embedding_model",
    num_train_epochs=1,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    warmup_steps=100,
    fp16=True,
    eval_steps=100,
    logging_steps=100,
)

# Train model
trainer = SentenceTransformerTrainer(
    model=embedding_model,
    args=args,
    train_dataset=train_dataset,
    loss=train_loss,
    evaluator=evaluator
)
trainer.train()






# Evaluate our trained model
evaluator(embedding_model)




Compared to using both the silver and gold datasets, using only the gold dataset reduces the performance of the model!


⚠️ VRAM Clean-up * Restart the notebook in order to clean-up memory if you move on to the next training example.




import gc
import torch

gc.collect()
torch.cuda.empty_cache()








Unsupervised Learning




Tranformer-based Denoising AutoEncoder (TSDAE)




# Download additional tokenizer
import nltk
nltk.download('punkt')






from tqdm import tqdm
from datasets import Dataset, load_dataset
from sentence_transformers.datasets import DenoisingAutoEncoderDataset

# Create a flat list of sentences
mnli = load_dataset("glue", "mnli", split="train").select(range(25_000))
flat_sentences = mnli["premise"] + mnli["hypothesis"]

# Add noise to our input data
damaged_data = DenoisingAutoEncoderDataset(list(set(flat_sentences)))

# Create dataset
train_dataset = {"damaged_sentence": [], "original_sentence": []}
for data in tqdm(damaged_data):
    train_dataset["damaged_sentence"].append(data.texts[0])
    train_dataset["original_sentence"].append(data.texts[1])
train_dataset = Dataset.from_dict(train_dataset)






train_dataset[0]






# # Choose a different deletion ratio
# flat_sentences = list(set(flat_sentences))
# damaged_data = DenoisingAutoEncoderDataset(
#     flat_sentences,
#     noise_fn=lambda s: DenoisingAutoEncoderDataset.delete(s, del_ratio=0.6)
# )






from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator

# Create an embedding similarity evaluator for stsb
val_sts = load_dataset('glue', 'stsb', split='validation')
evaluator = EmbeddingSimilarityEvaluator(
    sentences1=val_sts["sentence1"],
    sentences2=val_sts["sentence2"],
    scores=[score/5 for score in val_sts["label"]],
    main_similarity="cosine"
)






from sentence_transformers import models, SentenceTransformer

# Create your embedding model
word_embedding_model = models.Transformer('bert-base-uncased')
pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), 'cls')
embedding_model = SentenceTransformer(modules=[word_embedding_model, pooling_model])






from sentence_transformers import losses

# Use the denoising auto-encoder loss
train_loss = losses.DenoisingAutoEncoderLoss(
    embedding_model, tie_encoder_decoder=True
)
train_loss.decoder = train_loss.decoder.to("cuda")






from sentence_transformers.trainer import SentenceTransformerTrainer
from sentence_transformers.training_args import SentenceTransformerTrainingArguments

# Define the training arguments
args = SentenceTransformerTrainingArguments(
    output_dir="tsdae_embedding_model",
    num_train_epochs=1,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    warmup_steps=100,
    fp16=True,
    eval_steps=100,
    logging_steps=100,
)

# Train model
trainer = SentenceTransformerTrainer(
    model=embedding_model,
    args=args,
    train_dataset=train_dataset,
    loss=train_loss,
    evaluator=evaluator
)
trainer.train()






# Evaluate our trained model
evaluator(embedding_model)






import gc
import torch

gc.collect()
torch.cuda.empty_cache()












 Back to top