简介

Sentence Transformers 是一个广受认可的 Python 模块，用于训练或微调最先进的文本嵌入模型。在大型语言模型（LLM）领域，嵌入起着至关重要的作用，因为它能针对特定数据集显著提高相似性搜索等任务的性能。

最近，Hugging Face 发布了 3.0.0 版 Sentence Transformers，简化了训练、记录和评估过程。在本文中，我们将探讨如何使用我们的数据训练和微调句子转换器模型。

用于相似性搜索的嵌入

嵌入是将文本转换为固定大小的向量表示（浮点数）的过程，这种表示可以捕捉文本与其他词语之间的语义关系。如何将其用于相似性搜索？在相似性搜索中，我们将查询嵌入到向量数据库中。当用户提交查询时，我们需要在数据库中找到相似的查询。

首先，将所有文本数据转换为固定大小的矢量嵌入，并将其存储到矢量数据库中。接着，接受用户的查询，并将其转换为嵌入式数据。然后，在矢量数据库中找到与用户查询相似的搜索词或关键词，并检索出最接近的嵌入。这简单吗？是的，但要搜索最接近的嵌入，我们需要使用基于距离的算法，如余弦相似度、曼哈顿距离或欧氏距离。

什么是 SBERT？

SBERT （Sentence-BERT）是一种专门的句子转换器模型，用于高效的句子处理和比较。它采用连体网络架构，利用相同的 BERT 模型独立处理句子对。此外，SBERT 还在最终输出层利用均值池生成高质量的句子嵌入。

安装和设置

你可以使用在线笔记本，如 Google Colab。我还介绍了如何通过脚本执行训练代码。对于 Google Colab，请将运行环境设置为 T4 GPU 硬件。

!pip install -U "sentence-transformers[train]" accelerate datasets"sentence-transformers[train]" accelerate datasets

导入依赖项

import os
import json
import torch
import datasets
import pandas as pd
from torch.utils.data import DataLoader
from sentence_transformers import (
    SentenceTransformer, models,
    losses, util,
    InputExample, evaluation,
    SentenceTransformerTrainingArguments, SentenceTransformerTrainer
)
from accelerate import Accelerator
from datasets import load_dataset

在这篇文中，我使用了 Glue STS-B 数据和模型 sentence-transformers/all-MiniLM-L6-v2

data = load_dataset('sentence-transformers/stsb')'sentence-transformers/stsb')
train_data = data['train'].select(range(100))
val_data = data['validation'].select(range(100, 140))

在上面的代码块中，我选择了 100 个样本作为训练样本，40 个样本作为验证样本。这是因为 Colab 免费版的资源有限。你可以根据需要调整范围大小或导入整个数据集。

让我们看看训练数据中的随机样本数据

# Example data from 5th record (taking randomly to just display)
print("Sentence 1: ", train_data['sentence1'][5], "\nSentence 2: ", train_data['sentence2'][5], "\nScore: ", train_data['score'][5])

输出：

Sentence 1:  Some men are fighting. 
Sentence 2:  Two men are fighting. 
Score:  0.85

这将是我们的数据格式： 句子 1"、"句子 2 "和 "分数"。分数 "代表两个句子之间的接近或相似程度。在没有标签分数的情况下，只需相应地修改损失函数和评估器即可。

训练 SBERT

这是 SBERT 官方网站推荐的训练 SBERT 模型的方法。

训练代码：

def main():
   
    # Get number of GPUs working
    accelerator = Accelerator()
    print(f"Using GPUs: {accelerator.num_processes}")
    # Sentence Transformer BERT Model
    word_embedding_model = models.Transformer('sentence-transformers/all-MiniLM-L6-v2')
    # Applying pooling on final layer
    pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension())
    model = SentenceTransformer(modules=[word_embedding_model, pooling_model])
    # Define loss
    loss = losses.CoSENTLoss(model)
    # Define evaluator for evaluation
    evaluator = evaluation.EmbeddingSimilarityEvaluator(
        sentences1=val_data['sentence1'],
        sentences2=val_data['sentence2'],
        scores=val_data['score'],
        main_similarity=evaluation.SimilarityFunction.COSINE,
        name="sts-dev"
    )

    # Training arguments
    training_args = SentenceTransformerTrainingArguments(
        output_dir='./sbert-checkpoint', # Save checkpoints
        num_train_epochs=10,
        seed=33,
        per_device_train_batch_size=8,
        per_device_eval_batch_size=8,
        learning_rate=2e-5,
        fp16=True, # Loading model in mixed-precision
        warmup_ratio=0.1,
        evaluation_strategy="steps",
        eval_steps=2,
        save_total_limit=2,
        load_best_model_at_end=True,
        save_only_model=True,
        greater_is_better=True
    )

    # Train model
    trainer = SentenceTransformerTrainer(
        model=model,
        evaluator=evaluator,
        args=training_args,
        train_dataset=train_data,
        eval_dataset=val_data,
        loss=loss
    )
    trainer.train()
    # save the model
    model.save_pretrained("./sbert-model/")

现在，让我们逐步了解 main() 函数内部的功能：

• 定义 Accelerator() 以确定当前机器上可用的 GPU 数量。
• 从 HuggingFace 存储库中加载句子转换器模型，并使用均值池提取单词嵌入维度。在 SBERT 模型之后添加均值池化层作为输出。
• 定义损失函数（如 CoSENTLoss()），以便根据浮动相似性得分计算模型的损失。根据数据和标签，从 SBERT 的选项中选择合适的损失函数。
• 使用 Sentence Transformers 提供的 Evaluator() 类计算训练过程中的评估损失并获取特定指标。根据你的数据和用例选择合适的评估器，如 EmbeddingSimilarityEvaluator()。
• 使用从转换器 TrainingArguments 间接继承的 SentenceTransformerTrainer 类指定训练参数，如用于存储检查点的输出目录、每个设备（CPU/GPU）的批量大小、训练历元数、学习率、用于模型加载的 float16 精度、评估步骤等。
• 使用 SentenceTransformerTrainer 类训练模型，方法是定义训练数据和验证数据，可选择包含评估器，指定训练参数，以及定义损失函数。调用 train() 方法启动训练。执行训练并进一步保存模型。

训练 SBERT 的各种方法

定义完 main() 函数后，只需调用它即可启动模型训练过程。有几种方法可以做到这一点：

对于单 GPU：

• 如果使用 T4 GPU 在 Google Colab 免费版中运行代码，则只需创建一个新单元并调用函数：main()
• 如果在 Python 脚本中运行代码，则只需在终端运行 python 命令：python main.py。

对于多 GPU：

HuggingFace 变换器支持分布式数据并行（DDP）训练，可在多个 GPU 或多台机器上执行分布式并行训练。

• 如果你在包含多 GPU 的 colab 或任何笔记本中运行代码，那么：

from accelerator import notebook_launcher
notebook_launcher(main, num_processes=2)

将上述代码单独运行，就能在多 GPU 环境中运行。

• 对于 Python 脚本：

accelerate launch –multi-gpu –num_processes=2 main.py

这些是运行脚本或笔记本进行 SBERT 培训的一些常用方法。

测试模型

训练模型后，我们可以重新加载模型并进行推理测试。例如，如果我们有一个产品名称列表，而用户输入了搜索词，那么我们的目标就是识别出最相似的产品名称并给出分数。

在使用相似度分数作为标签对句子相似度数据进行嵌入模型训练后，现在就可以改进嵌入模型了。

以下是我们用于嵌入数据的产品名称示例列表：

# List of products
products = [
    "Apple iPhone 15 (256GB) | Silver",
    "Nike Air Max 2024 | Blue/White",
    "Samsung Galaxy S24 Ultra (512GB) | Phantom Black",
    "Sony PlayStation 5 Console | Digital Edition",
    "Dell XPS 13 Laptop | Intel i7, 16GB RAM, 512GB SSD",
    "Fitbit Charge 6 | Midnight Blue",
    "Bose QuietComfort 45 Headphones | Triple Black",
    "Canon EOS R6 Camera | 20.1 MP Mirrorless",
    "Microsoft Surface Pro 9 | Intel i5, 8GB RAM, 256GB SSD",
    "Adidas Ultraboost 21 Running Shoes | Core Black",
    "Amazon Kindle Paperwhite | 32GB, Waterproof",
    "LG OLED65C1PUB 65\" 4K Smart TV",
    "Garmin Forerunner 955 Smartwatch | Slate Grey",
    "Google Nest Thermostat | Charcoal",
    "KitchenAid Stand Mixer | 5-Quart, Empire Red",
    "Dyson V11 Torque Drive Cordless Vacuum",
    "JBL Charge 5 Portable Bluetooth Speaker | Squad",
    "Panasonic Lumix GH5 Camera | 20.3 MP, 4K Video",
    "Apple MacBook Pro 14\" | M1 Pro, 16GB RAM, 1TB SSD",
    "Under Armour HeatGear Compression Shirt | Black/Red"
]

接下来，加载经过微调的 SBERT 模型，并将产品名称转换为向量嵌入：

# Load fine-tuned model
model = SentenceTransformer('./sbert-model')

为了将产品名称转换为嵌入式，我们将利用 GPU 将其转换为张量。你可以使用以下代码来完成：

product_data = model.encode(products, convert_to_tensor=True).to("cuda")True).to("cuda")

通过将嵌入转换为 CUDA，我们可以利用 GPU 计算支持（dtype=torch.float32）；否则，如果选择 CPU，则默认为（dtype=float32）。

这个 product_data 将作为我们的矢量数据库，现在存储在内存中。或者，也可以使用 Qdrant、Pinecone、Chroma 等矢量数据库。

最后，创建一个函数，接受来自终端的用户查询或用户输入，并返回排名靠前的产品及其余弦相似度得分。

def search():
    query = input("Enter Query:\n")
    query_embeddings = model.encode([query], convert_to_tensor=True).to("cuda")
    hits = util.semantic_search(query_embeddings, product_data,
                                score_function=util.cos_sim)
   
    for i in range(5):
        best_search_term_id, best_search_term_core = hits[0][i]['corpus_id'], hits[0][i]['score']
        print("\nTop result: ", products[best_search_term_id])
        print("Score: ", best_search_term_core)

试运行：

可以看出，我们的模型表现优异，得分令人满意。为了进一步提高结果的相关性，可以考虑添加 0.5 的阈值比率。

结论

使用 SentenceTransformer 3.0.0 可以轻而易举地训练或微调嵌入模型。新版本通过 DDP 方法支持多 GPU 使用，并通过 Weights & Biases 引入了日志和实验功能。通过将我们的代码封装在单个主函数中并使用单个命令执行，开发人员可以大大简化工作流程。