EN | 简体中文

AnglE 📐

由 Mixedbread 赞助

如需更详细的使用说明，请参阅 📘 文档： https://angle.readthedocs.io/en/latest/index.html

📢 使用 AnglE 训练/推理强大的句子嵌入。 该库源自论文：AnglE: 角度优化的文本嵌入。它只需几行代码即可训练最先进的基于 BERT/LLM 的句子嵌入。AnglE 也是一个通用的句子嵌入推理框架，支持多种基于 Transformer 的句子嵌入的推理。

✨ 特性

损失函数：

• 📐 AnglE 损失（ACL24）
• ⚖ 对比损失
• 📏 CoSENT 损失
• ☕️ Espresso 损失（ICLR 2025，又名 2DMSE，详情见：README_ESE）

骨干模型：

• 基于 BERT 的模型（BERT、RoBERTa、ModernBERT 等）
• 基于 LLM 的模型（LLaMA、Mistral、Qwen 等）
• 双向 LLM 基础模型（LLaMA、Mistral、Qwen、OpenELMo 等，详情请参阅：https://github.com/WhereIsAI/BiLLM）

训练：

• 单 GPU 训练
• 多 GPU 训练

未来还将添加更多功能。

🏆 成就

📅 2024年5月16日 | 论文“AnglE: 角度优化的文本嵌入”被 ACL 2024 主会场接收。

📅 2024年3月13日 | 论文“BeLLM: 面向句子嵌入的反向依赖增强大型语言模型”被 NAACL 2024 主会场接收。

📅 2024年3月8日 | 🍞 mixedbread 的嵌入 (mixedbread-ai/mxbai-embed-large-v1) 在 MTEB Leaderboard 上以平均分 64.68 获得 SOTA！该模型使用 AnglE 训练而成。恭喜 mixedbread！

📅 2023年12月4日 | 我们的通用句子嵌入 WhereIsAI/UAE-Large-V1 在 MTEB Leaderboard 上以平均分 64.64 获得 SOTA！该模型同样使用 AnglE 训练而成。

📅 2023年12月 | AnglE 在 STS 语义文本相似度基准测试中达到 SOTA 表现！

🤗 官方预训练模型

基于 BERT 的模型：

🤗 HF	最大标记数	池化策略	场景
WhereIsAI/UAE-Large-V1	512	cls	英语，通用
WhereIsAI/UAE-Code-Large-V1	512	cls	代码相似性
WhereIsAI/pubmed-angle-base-en	512	cls	医学相似性
WhereIsAI/pubmed-angle-large-en	512	cls	医学相似性

基于 LLM 的模型：

🤗 HF (lora 权重)	骨干模型	最大标记数	提示词	池化策略	场景
SeanLee97/angle-llama-13b-nli	NousResearch/Llama-2-13b-hf	4096	Prompts.A	最后一个 token	英语，相似性测量
SeanLee97/angle-llama-7b-nli-v2	NousResearch/Llama-2-7b-hf	4096	Prompts.A	最后一个 token	英语，相似性测量

💡 您可以在 HuggingFace Collection 中找到更多使用 AnglE 训练的第三方嵌入。

🚀 快速入门

⬇️ 安装

使用 uv

uv pip install -U angle-emb

或 pip

pip install -U angle-emb

---

🔍 推理

1️⃣ 基于 BERT 的模型

选项 A：使用提示词（用于检索任务）

使用包含 {text} 占位符的提示词。可通过 Prompts.list_prompts() 查看可用提示词。

from angle_emb import AnglE, Prompts
from angle_emb.utils import cosine_similarity

# 加载模型
angle = AnglE.from_pretrained('WhereIsAI/UAE-Large-V1', pooling_strategy='cls').cuda()

# 使用提示词编码查询，不使用提示词编码文档
qv = angle.encode(['天气怎么样？'], to_numpy=True, prompt=Prompts.C)
doc_vecs = angle.encode([
    '天气真好！',
    '今天下雨了。',
    '我要去睡觉了'
], to_numpy=True)

# 计算相似度
for dv in doc_vecs:
    print(cosine_similarity(qv[0], dv))

选项 B：不使用提示词（用于相似性任务）

from angle_emb import AnglE
from angle_emb.utils import cosine_similarity

# 加载模型
angle = AnglE.from_pretrained('WhereIsAI/UAE-Large-V1', pooling_strategy='cls').cuda()

# 编码文档
doc_vecs = angle.encode([
    '天气真好！',
    '天气非常好！',
    '我要去睡觉了'
])

# 计算两两相似度
for i, dv1 in enumerate(doc_vecs):
    for dv2 in doc_vecs[i+1:]:
        print(cosine_similarity(dv1, dv2))

---

2️⃣ 基于 LLM 的模型

对于基于 LoRA 的模型，需同时指定骨干模型和 LoRA 权重。对于 LLM 模型，务必设置 is_llm=True。

import torch
from angle_emb import AnglE, Prompts
from angle_emb.utils import cosine_similarity

# 加载带有 LoRA 权重的 LLM
angle = AnglE.from_pretrained(
    'NousResearch/Llama-2-7b-hf',
    pretrained_lora_path='SeanLee97/angle-llama-7b-nli-v2',
    pooling_strategy='last',
    is_llm=True,
    torch_dtype=torch.float16
).cuda()

# 使用提示词编码
doc_vecs = angle.encode([
    '天气真好！',
    '天气非常好！',
    '我要去睡觉了'
], prompt=Prompts.A)

# 计算相似度
for i, dv1 in enumerate(doc_vecs):
    for dv2 in doc_vecs[i+1:]:
        print(cosine_similarity(dv1, dv2))

---

3️⃣ 基于 BiLLM 的模型

通过设置 apply_billm=True 并指定模型类，启用双向 LLM。

import os
import torch
from angle_emb import AnglE
from angle_emb.utils import cosine_similarity

# 设置 BiLLM 环境变量
os.environ['BiLLM_START_INDEX'] = '31'

# 加载 BiLLM 模型
angle = AnglE.from_pretrained(
    'NousResearch/Llama-2-7b-hf',
    pretrained_lora_path='SeanLee97/bellm-llama-7b-nli',
    pooling_strategy='last',
    is_llm=True,
    apply_billm=True,
    billm_model_class='LlamaForCausalLM',
    torch_dtype=torch.float16
).cuda()

# 使用自定义提示进行编码
doc_vecs = angle.encode([
    '天气真好！',
    '天气非常好！',
    '我要去睡觉了'
], prompt='句子 {text} 的代表性词语是：')

# 计算相似度
for i, dv1 in enumerate(doc_vecs):
    for dv2 in doc_vecs[i+1:]:
        print(cosine_similarity(dv1, dv2))

---

4️⃣ 浓缩版/套娃模型

通过截断层和嵌入维度，实现灵活的模型压缩。

from angle_emb import AnglE
from angle_emb.utils import cosine_similarity

# 加载模型
angle = AnglE.from_pretrained('mixedbread-ai/mxbai-embed-2d-large-v1', pooling_strategy='cls').cuda()

# 截断到特定层
angle = angle.truncate_layer(layer_index=22)

# 使用截断后的嵌入尺寸进行编码
doc_vecs = angle.encode([
    '天气真好！',
    '天气非常好！',
    '我要去睡觉了'
], embedding_size=768)

# 计算相似度
for i, dv1 in enumerate(doc_vecs):
    for dv2 in doc_vecs[i+1:]:
        print(cosine_similarity(dv1, dv2))

---

5️⃣ 第三方模型

使用 AnglE 可以加载任何基于 Transformer 的模型（例如 sentence-transformers、BAAI/bge 等）。

from angle_emb import AnglE

# 加载第三方模型
model = AnglE.from_pretrained('mixedbread-ai/mxbai-embed-large-v1', pooling_strategy='cls').cuda()

# 编码文本
vec = model.encode('hello world', to_numpy=True)
print(vec)

---

⚡ 批量推理

使用 batched 库加速推理（推荐用于大规模处理）。

uv pip install batched

import batched
from angle_emb import AnglE

# 加载模型
model = AnglE.from_pretrained("WhereIsAI/UAE-Large-V1", pooling_strategy='cls').cuda()

# 启用动态批处理
model.encode = batched.dynamically(model.encode, batch_size=64)

# 编码大批量数据
vecs = model.encode([
    '天气真好！',
    '天气非常好！',
    '我要去睡觉了'
] * 50)

🕸️ 自定义训练

💡 完整详情请参阅 官方训练文档。

---

🗂️ 第一步：准备数据集

AnglE 支持三种数据集格式。请根据任务选择：

格式	列	描述	使用场景
格式 A	text1, text2, label	成对文本及相似度分数（0–1）	相似度评分
格式 B	query, positive	查询-文档对	检索任务，无需硬负样本
格式 C	query, positive, negative	查询及其正负样本	对比学习

注意事项：

• 所有格式均使用 HuggingFace 的 datasets.Dataset
• text1、text2、query、positive 和 negative 可以是字符串或字符串列表（列表时会随机采样）

---

🚂 第二步：训练方法

方案 A：命令行训练（推荐）

单 GPU：

CUDA_VISIBLE_DEVICES=0 angle-trainer --help

多 GPU 与 FSDP：

CUDA_VISIBLE_DEVICES=0,1,2,3 WANDB_MODE=disabled accelerate launch \
  --multi_gpu \
  --num_processes 4 \
  --main_process_port 2345 \
  --config_file examples/FSDP/fsdp_config.yaml \
  -m angle_emb.angle_trainer \
  --gradient_checkpointing 1 \
  --use_reentrant 0 \
  ...

多 GPU（标准方式）：

CUDA_VISIBLE_DEVICES=0,1,2,3 WANDB_MODE=disabled accelerate launch \
  --multi_gpu \
  --num_processes 4 \
  --main_process_port 2345 \
  -m angle_emb.angle_trainer \
  --model_name_or_path YOUR_MODEL \
  --train_name_or_path YOUR_DATASET \
  ...

📁 更多示例：examples/Training

---

方案 B：Python API 训练

from datasets import load_dataset
from angle_emb import AnglE

# 第一步：加载预训练模型
angle = AnglE.from_pretrained(
    'SeanLee97/angle-bert-base-uncased-nli-en-v1',
    max_length=128,
    pooling_strategy='cls'
).cuda()

# 第二步：准备数据集（格式 A 示例）
ds = load_dataset('mteb/stsbenchmark-sts')
ds = ds.map(lambda obj: {
    "text1": str(obj["sentence1"]),
    "text2": str(obj['sentence2']),
    "label": obj['score']
})
ds = ds.select_columns(["text1", "text2", "label"])

# 第三步：训练模型
angle.fit(
    train_ds=ds['train'].shuffle(),
    valid_ds=ds['validation'],
    output_dir='ckpts/sts-b',
    batch_size=32,
    epochs=5,
    learning_rate=2e-5,
    save_steps=100,
    eval_steps=1000,
    warmup_steps=0,
    gradient_accumulation_steps=1,
    loss_kwargs={
        'cosine_w': 1.0,
        'ibn_w': 1.0,
        'angle_w': 0.02,
        'cosine_tau': 20,
        'ibn_tau': 20,
        'angle_tau': 20
    },
    fp16=True,
    logging_steps=100
)

# 第四步：评估
corrcoef = angle.evaluate(ds['test'])
print('Spearman 相关系数：', corrcoef)

---

⚙️ 高级配置

训练特殊模型

模型类型	CLI 标志	描述
LLM	--is_llm 1 + LoRA 参数	必须手动启用 LLM 模式
BiLLM	--apply_billm 1 --billm_model_class LlamaForCausalLM	双向 LLM（指南）
浓缩版（ESE）	--apply_ese 1 --ese_kl_temperature 1.0 --ese_compression_size 256	套娃风格嵌入

应用提示词

格式	标志	适用对象
格式 A	--text_prompt "text: {text}"	同时应用于 text1 和 text2
格式 B/C	--query_prompt "query: {text}"	仅适用于 query 字段
格式 B/C	--doc_prompt "document: {text}"	仅适用于 positive 和 negative 字段

列映射（旧版兼容性）

无需修改即可适配旧数据集：

# 命令行
--column_rename_mapping "text:query"

# Python
列重命名映射={"text": "query"}

模型转换

将训练好的模型转换为 sentence-transformers 格式：

python scripts/convert_to_sentence_transformers.py --help

---

💡 微调技巧

📖 完整文档

格式	建议
格式 A	增加 cosine_w 或减小 ibn_w
格式 B	仅调整 ibn_w 和 ibn_tau
格式 C	设置 cosine_w=0, angle_w=0.02, 并配置 cln_w + ibn_w

防止灾难性遗忘：

• 设置 teacher_name_or_path 进行知识蒸馏
• 自蒸馏时使用相同的模型路径
• ⚠️ 确保教师模型和学生模型使用 相同的分词器

---

🔄 与 sentence-transformers 集成

任务	状态	备注
训练	⚠️ 部分	SentenceTransformers 提供了 AnglE 损失，但为了获得最佳效果，请使用官方的 angle_emb
推理	✅ 全面	可以通过 examples/convert_to_sentence_transformers.py 转换训练好的模型

🫡 引用

如果您使用了我们的代码和预训练模型，请通过以下方式引用我们的工作以支持我们：

@article{li2023angle,
  title={AnglE优化文本嵌入},
  author={李贤明和李静},
  journal={arXiv 预印本 arXiv:2309.12871},
  year={2023}
}

📜 更改日志

📅	描述
2025年1月	v0.6.0 - 重大重构 🎉： • 移除 AngleDataTokenizer - 无需预先对数据集进行分词！ • 移除 DatasetFormats 类 - 使用字符串字面量（'A', 'B', 'C'） • 移除 LLM 模型的自动检测 - 手动设置 is_llm • 将 --prompt_template 重命名为 --text_prompt（仅适用于格式 A） • 新增 --query_prompt 和 --doc_prompt 用于格式 B/C • 新增 --column_rename_mapping，以便在不修改旧数据集的情况下适配 • 更新数据格式：格式 B/C 现在使用 query、positive、negative 字段 • 支持格式 B/C 中基于列表的采样 • 更新示例，使用 accelerate launch • 请参阅 MIGRATION_GUIDE.md 获取升级说明
2024年5月21日	支持 Espresso 句子嵌入
2024年2月7日	支持仅使用正样本对进行训练（格式 C：query, positive）
2023年12月4日	发布通用英语句子嵌入模型：WhereIsAI/UAE-Large-V1
2023年11月2日	发布英文预训练模型：SeanLee97/angle-llama-13b-nli
2023年10月28日	发布两款中文预训练模型：SeanLee97/angle-roberta-wwm-base-zhnli-v1 和 SeanLee97/angle-llama-7b-zhnli-v1；新增中文 README.md

📧 联系方式

如果您有任何问题或建议，请随时通过电子邮件与我们联系：xmlee97@gmail.com

© 许可证

本项目采用 MIT 许可证。关于预训练模型的许可信息，请参阅相应模型的许可证文件。

AnglE 快速上手指南

AnglE 是一个用于训练和推理高性能句子嵌入（Sentence Embeddings）的开源库，支持 BERT 和 LLM 架构，实现了 ACL 2024 论文提出的角度优化损失函数。

1. 环境准备

• 系统要求：Linux / macOS / Windows (推荐 Linux)
• Python 版本：3.8+
• 硬件要求：
  • 推理：CPU 或 GPU (CUDA 兼容)
  • 训练：推荐 NVIDIA GPU (支持单卡及多卡 FSDP 训练)
• 前置依赖：PyTorch, Transformers, Accelerate

2. 安装步骤

推荐使用 uv 进行快速安装，也可使用标准的 pip。

方式一：使用 uv (推荐)

uv pip install -U angle-emb

方式二：使用 pip

pip install -U angle-emb

提示：如需加速下载，可在命令后添加国内镜像源，例如：-i https://pypi.tuna.tsinghua.edu.cn/simple

3. 基本使用

3.1 模型推理 (Inference)

AnglE 支持加载 HuggingFace 上的预训练模型进行向量化编码。

场景 A：基于 BERT 的模型 (通用检索/相似度)

from angle_emb import AnglE, Prompts
from angle_emb.utils import cosine_similarity

# 1. 加载模型 (自动下载)
# pooling_strategy 根据模型类型选择 'cls' 或 'last'
angle = AnglE.from_pretrained('WhereIsAI/UAE-Large-V1', pooling_strategy='cls').cuda()

# 2. 编码文本
# 检索任务建议给 Query 添加 Prompt (如 Prompts.C)，文档通常不需要
query_vec = angle.encode(['what is the weather?'], to_numpy=True, prompt=Prompts.C)
doc_vecs = angle.encode([
    'The weather is great!',
    'it is rainy today.',
    'i am going to bed'
], to_numpy=True)

# 3. 计算余弦相似度
for dv in doc_vecs:
    score = cosine_similarity(query_vec[0], dv)
    print(f"Similarity: {score}")

场景 B：基于 LLM 的模型 (LoRA 微调)

加载 LLM 基座模型及对应的 LoRA 权重时，必须设置 is_llm=True。

import torch
from angle_emb import AnglE, Prompts

# 加载 LLM 基座 + LoRA 权重
angle = AnglE.from_pretrained(
    'NousResearch/Llama-2-7b-hf',
    pretrained_lora_path='SeanLee97/angle-llama-7b-nli-v2',
    pooling_strategy='last',
    is_llm=True,
    torch_dtype=torch.float16
).cuda()

# 编码 (LLM 模型通常强依赖 Prompt)
vectors = angle.encode([
    'The weather is great!',
    'The weather is very good!'
], prompt=Prompts.A)

3.2 批量推理优化

处理大规模数据时，结合 batched 库可显著提升速度。

uv pip install batched

import batched
from angle_emb import AnglE

model = AnglE.from_pretrained("WhereIsAI/UAE-Large-V1", pooling_strategy='cls').cuda()

# 启用动态批处理
model.encode = batched.dynamically(model.encode, batch_size=64)

# 执行大批量编码
vecs = model.encode(['text sample'] * 1000)

3.3 自定义训练 (Custom Training)

AnglE 支持三种数据集格式，可通过 CLI 或 Python API 启动训练。

数据集格式示例 (Format C - 对比学习): 需包含 query, positive, negative 三列。

启动训练 (单卡 CLI 模式):

# 查看帮助信息
CUDA_VISIBLE_DEVICES=0 angle-trainer --help

# 示例启动命令 (需替换实际路径)
CUDA_VISIBLE_DEVICES=0 angle-trainer \
  --model_name_or_path bert-base-uncased \
  --train_name_or_path path/to/your/dataset \
  --output_dir ./saved_models \
  --num_train_epochs 3 \
  --per_device_train_batch_size 32 \
  --learning_rate 2e-5

注意：完整训练参数配置及多卡 FSDP 训练方案，请参考官方文档：https://angle.readthedocs.io/en/latest/notes/training.html