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

创建文本嵌入模型

• • 创建嵌入模型
  • • 数据
    • 模型
    • 损失函数
    • 评估器
    • 训练
  • MTEB（大规模文本嵌入基准测试）
  • • !!!!!!!!!离线下载
    • !!!!!!!!!!!离线导入

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

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

图片

可选\] - 安装所需包 如果您在 Google Colab（或任何其他云服务）上查看此笔记本，需要取消注释并运行以下代码块以安装本章的依赖项： 💡 注意：我们需要使用 GPU 来运行本笔记本中的示例。在 Google Colab 中，转到 运行时 \> 更改运行时类型 \> 硬件加速器 \> GPU \> GPU 类型 \> T4。 ```python # %%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条样本。 ```python 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=矛盾） ```python 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") ``` 查看数据集中第二条样本的内容，了解数据结构 ```python 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}` *** ** * ** *** #### 模型 离线下载模型 ```python 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模型作为基础嵌入模型 ```python from sentence_transformers import SentenceTransformer # Use a base model embedding_model = SentenceTransformer('/workspace/huggingface_cache/models/bert-base-uncased') ``` *** ** * ** *** #### 损失函数 导入损失函数模块，并定义SoftmaxLoss损失函数，需要指定句子嵌入维度和标签数量 ```python 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验证集，创建余弦相似度评估器来评估模型性能 ```python 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", ) ``` 这里离线下载和加载的步骤同上面 *** ** * ** *** #### 训练 导入训练参数模块，定义训练的超参数，包括输出目录、训练轮数、批次大小等 ```python 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 ```python pip install transformers==4.41 ``` 导入训练器模块，创建SentenceTransformerTrainer实例，传入模型、参数、数据集、损失函数和评估器 ```python from sentence_transformers.trainer import SentenceTransformerTrainer # 训练嵌入模型 trainer = SentenceTransformerTrainer( model=embedding_model, args=args, train_dataset=train_dataset, loss=train_loss, evaluator=evaluator ) trainer.train() ``` 输出  使用评估器评估训练好的模型性能，返回皮尔逊和斯皮尔曼相关系数等指标 ```python # 评估我们训练好的模型 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（大规模文本嵌入基准测试） #### !!!离线下载 ```python 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分类任务进行模型评估 ```python 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}}}