完整教程：AI大模型入门到实战系列（十四）创建文本嵌入模型

创建文本嵌入模型

探索训练和微调嵌入模型的方法。
图片图片图片在 Colab 中打开

本笔记来自 Jay Alammar 和 Maarten Grootendorst 所著《动手学大语言模型》一书的第十章。

图片
[可选] - 安装所需包
如果您在 Google Colab（或任何其他云服务）上查看此笔记本，需要取消注释并运行以下代码块以安装本章的依赖项：

注意：我们需要使用 GPU 来运行本笔记本中的示例。在 Google Colab 中，转到运行时 > 更改运行时类型 > 硬件加速器 > GPU > GPU 类型 > T4。

# %%capture
# !pip install -q accelerate>=0.27.2 peft>=0.9.0 bitsandbytes>=0.43.0 transformers>=4.38.2 trl>=0.7.11 sentencepiece>=0.1.99
# !pip install -q sentence-transformers>=3.0.0 mteb>=1.1.2 datasets>=2.18.0

创建嵌入模型

数据

离线下载数据glue，从GLUE基准数据集中加载MNLI（多类型自然语言推理）任务的训练集，然后从中选取前50,000条样本。

from datasets import load_dataset
import os
# 设置代理（如果需要）
os.environ['HTTP_PROXY'] = 'http://127.0.0.1:10809'
os.environ['HTTPS_PROXY'] = 'http://127.0.0.1:10809'
# 设置缓存目录
os.environ['HF_HOME'] = 'D:/huggingface_cache'
# 方法2：下载前50,000条MNLI训练集（如你最初想要的）
train_dataset = load_dataset("glue", "mnli", split="train", cache_dir="D:/huggingface_cache/datasets").select(range(50_000))
print(f"50k子集大小: {len(train_dataset)}")  # 50,000条
# 查看数据
print(train_dataset[0])  # 查看第一条数据

从datasets库中离线加载MNLI数据集，该数据集包含前提、假设和标签（0=蕴含，1=中立，2=矛盾）

from datasets import load_dataset
# Load MNLI dataset from GLUE
# 0 = entailment, 1 = neutral, 2 = contradiction
train_dataset = load_dataset('/workspace/huggingface_cache/datasets/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}

模型

离线下载模型

from huggingface_hub import snapshot_download
import os
os.environ['HTTP_PROXY'] = 'http://127.0.0.1:10809'
os.environ['HTTPS_PROXY'] = 'http://127.0.0.1:10809'
# 设置模型路径
model_id = "bert-base-uncased"
local_dir = "D:/huggingface_cache/models/bert-base-uncased"
# 确保目录存在
os.makedirs(local_dir, exist_ok=True)
# 下载模型（对应 hf download 命令）
print(f"开始下载模型: {model_id}")
print(f"保存到: {local_dir}")
try:
snapshot_download(
repo_id=model_id,
local_dir=local_dir,
local_dir_use_symlinks=False,  # 不使用符号链接
force_download=True,           # 强制重新下载
resume_download=False          # 不续传
# allow_patterns="*.gguf"   这里要注释掉，不然下载的是空文件
)
print("✅ 模型下载完成！")
except Exception as e:
print(f"❌ 下载失败: {e}")

从sentence-transformers库导入SentenceTransformer类，并离线加载bert-base-uncased模型作为基础嵌入模型

from sentence_transformers import SentenceTransformer
# Use a base model
embedding_model = SentenceTransformer('/workspace/huggingface_cache/models/bert-base-uncased')

损失函数

导入损失函数模块，并定义SoftmaxLoss损失函数，需要指定句子嵌入维度和标签数量

from sentence_transformers import losses
# 定义损失函数。在soft-max损失中，我们还需要显式设置标签数量。
train_loss = losses.SoftmaxLoss(
model=embedding_model,
sentence_embedding_dimension=embedding_model.get_sentence_embedding_dimension(),
num_labels=3
)

评估器

导入评估器模块，加载STS-B验证集，创建余弦相似度评估器来评估模型性能

from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator
# Create an embedding similarity evaluator for stsb
val_sts = load_dataset('/workspace/huggingface_cache/datasets/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.training_args import SentenceTransformerTrainingArguments
# 定义训练参数
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,
)

这里可能报错

TypeError: compute_loss() got an unexpected keyword argument ‘num_items_in_batch’

解决方法是将transformers的版本回退到4.41

pip install transformers==4.41

导入训练器模块，创建SentenceTransformerTrainer实例，传入模型、参数、数据集、损失函数和评估器

from sentence_transformers.trainer import SentenceTransformerTrainer
# 训练嵌入模型
trainer = SentenceTransformerTrainer(
model=embedding_model,
args=args,
train_dataset=train_dataset,
loss=train_loss,
evaluator=evaluator
)
trainer.train()

输出
在这里插入图片描述

使用评估器评估训练好的模型性能，返回皮尔逊和斯皮尔曼相关系数等指标

# 评估我们训练好的模型
evaluator(embedding_model)

{ 'pearson_cosine': 0.432, # 余弦相似度的皮尔逊相关系数 'spearman_cosine': 0.509, # 余弦相似度的斯皮尔曼相关系数 ✓ 'pearson_manhattan': 0.475, # 曼哈顿距离的皮尔逊相关系数 'spearman_manhattan': 0.502, # 曼哈顿距离的斯皮尔曼相关系数 'pearson_euclidean': 0.461, # 欧氏距离的皮尔逊相关系数 'spearman_euclidean': 0.497, # 欧氏距离的斯皮尔曼相关系数 'pearson_dot': 0.418, # 点积的皮尔逊相关系数 'spearman_dot': 0.440, # 点积的斯皮尔曼相关系数 'pearson_max': 0.475, # 最好的皮尔逊系数（曼哈顿距离） 'spearman_max': 0.509 # 最好的斯皮尔曼系数（余弦相似度） }

MTEB（大规模文本嵌入基准测试）

!!!离线下载

from datasets import load_dataset, DatasetDict
import os
# 设置代理（如果需要）
os.environ['HTTP_PROXY'] = 'http://127.0.0.1:10809'
os.environ['HTTPS_PROXY'] = 'http://127.0.0.1:10809'
# 1. 定义你想要的本地保存路径
local_save_path = "D:/huggingface_cache/"
# 2. 从官方路径下载数据集（使用 streaming=False 确保完整下载）
print(f"正在从Hub下载Banking77数据集到: {local_save_path}")
try:
# 直接下载整个数据集
dataset = load_dataset("PolyAI/banking77", cache_dir=local_save_path)
print("数据集下载完成！")
# 3. 检查并保存为本地格式（确保是DatasetDict格式）
if isinstance(dataset, DatasetDict):
print("正在将数据集保存为本地格式...")
dataset.save_to_disk(local_save_path)
print(f"数据集已成功保存至: {local_save_path}")
else:
# 如果返回的是单个Dataset，也将其包装为DatasetDict保存
print("检测到单个数据集，正在转换为DatasetDict并保存...")
dataset_dict = DatasetDict({"train": dataset}) # 假设它为训练集
dataset_dict.save_to_disk(local_save_path)
print(f"数据集已成功保存至: {local_save_path}")
except Exception as e:
print(f"下载或保存过程中出现错误: {e}")
# 可以考虑更详细的错误处理

!!!离线导入

导入MTEB基准测试模块，选择Banking77分类任务进行模型评估

from mteb.tasks import Banking77Classification
from datasets import load_from_disk
import os
# 设置离线模式
# os.environ.update({
#     'HF_DATASETS_OFFLINE': '1',
#     'HF_HUB_OFFLINE': '1',
#     'TRANSFORMERS_OFFLINE': '1',
# })
# 猴子补丁：替换数据集加载函数
import datasets
original_load_dataset = datasets.load_dataset
def offline_load_dataset(path, *args, **kwargs):
"""自定义加载函数，只从本地加载"""
if path == "mteb/banking77":
print(f"使用本地数据集代替 {path}")
local_path = "/workspace/huggingface_cache/datasets/banking77"
return load_from_disk(local_path)
else:
# 对于其他数据集，抛出错误或尝试本地加载
raise ConnectionError(f"离线模式：无法加载数据集 {path}")
# 临时替换函数
datasets.load_dataset = offline_load_dataset
try:
# 创建任务
task = Banking77Classification()
# 直接评估
print("开始评估...")
evaluation = MTEB(tasks=[task])
results = evaluation.run(embedding_model)
#     results = task.evaluate(embedding_model)
print("评估结果:")
print(results)
finally:
# 恢复原函数
datasets.load_dataset = original_load_dataset

安装mteb 可能报错

TypeError: unsupported operand type(s) for |: ‘type’ and ‘NoneType’
这个错误是因为 Python 3.8 不支持 | 运算符（联合类型）语法。在 Python 3.10 及以上版本中，type | None 是有效的语法（等效于 Optional[type]），但在 Python 3.8 中会报错

解决方法是降低mteb的版本
pip install mteb==1.1.2

输出

开始评估…
──────────────────────────── Selected task─────────────────────────────────
Classification
- Banking77Classification, s2s
使用本地数据集代替 mteb/banking77
评估结果:
{‘Banking77Classification’: {‘mteb_version’: ‘1.1.2’, ‘dataset_revision’: ‘0fd18e25b25c072e09e0d92ab615fda904d66300’, ‘mteb_dataset_name’: ‘Banking77Classification’, ‘test’: {‘accuracy’: 0.5085064935064935, ‘f1’: 0.5067556813716814, ‘accuracy_stderr’: 0.010948637349356522, ‘f1_stderr’: 0.010753042483624088, ‘main_score’: 0.5085064935064935, ‘evaluation_time’: 20.13}}}

posted @ 2026-01-24 18:38 clnchanpin 阅读(0) 评论(0) 收藏举报

刷新页面返回顶部