AnglE:开创句子嵌入的新纪元
在自然语言处理(NLP)领域,高质量的句子嵌入对于各种下游任务至关重要。近日,一个名为AnglE的开源框架引起了研究人员的广泛关注。AnglE是一个用于训练和推理强大句子嵌入的框架,它基于创新的角度优化损失函数,能够轻松训练出最先进的BERT和LLM(大型语言模型)句子嵌入模型。本文将详细介绍AnglE的特性、使用方法及其在NLP领域取得的重要成就。
AnglE的核心特性
AnglE框架具有以下几个突出的特性:
-
多种损失函数: AnglE支持多种先进的损失函数,包括:
- AnglE loss: 基于角度优化的创新损失函数
- Contrastive loss: 对比损失
- CoSENT loss: 余弦相似度损失
- Espresso loss: 之前称为2DMSE的损失函数
-
灵活的backbone支持: AnglE可以使用多种预训练模型作为backbone,包括:
- BERT系列模型(BERT、RoBERTa、ELECTRA、ALBERT等)
- 大型语言模型(LLaMA、Mistral、Qwen等)
- 双向LLM模型(基于BiLLM项目)
-
高效训练: AnglE支持单GPU和多GPU训练,可以充分利用硬件资源加速模型训练过程。
-
易用性: AnglE提供了简洁的API,只需几行代码即可完成模型的训练和推理。
-
性能卓越: 基于AnglE训练的模型在多个权威基准测试中取得了最先进(SOTA)的成绩。
AnglE的安装与快速入门
要开始使用AnglE,首先需要通过pip安装:
python -m pip install -U angle_emb
安装完成后,我们可以快速上手使用AnglE进行句子嵌入的推理。以下是一个使用预训练BERT模型进行推理的示例:
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()
# 编码句子
doc_vecs = angle.encode([
'The weather is great!',
'The weather is very good!',
'I am going to bed'
])
# 计算相似度
for i, dv1 in enumerate(doc_vecs):
for dv2 in doc_vecs[i+1:]:
print(cosine_similarity(dv1, dv2))
这个简单的例子展示了如何使用AnglE加载预训练模型、编码句子并计算句子之间的相似度。
AnglE的预训练模型
AnglE项目提供了多个高质量的预训练模型,适用于不同的场景:
-
通用英语模型:
- WhereIsAI/UAE-Large-V1: 512 tokens, cls pooling, 适用于一般英语场景
-
代码相似度模型:
- WhereIsAI/UAE-Code-Large-V1: 512 tokens, cls pooling, 专门用于代码相似度任务
-
医学相似度模型:
- WhereIsAI/pubmed-angle-base-en: 512 tokens, cls pooling
- WhereIsAI/pubmed-angle-large-en: 512 tokens, cls pooling
-
基于LLM的模型:
- SeanLee97/angle-llama-13b-nli: 基于LLaMA-2-13b, 4096 tokens
- SeanLee97/angle-llama-7b-nli-v2: 基于LLaMA-2-7b, 4096 tokens
这些预训练模型涵盖了多个领域和规模,用户可以根据具体需求选择合适的模型。
AnglE的自定义训练
除了使用预训练模型,AnglE还支持用户使用自己的数据进行自定义训练。以下是一个简单的训练示例:
from datasets import load_dataset
from angle_emb import AnglE, AngleDataTokenizer
# 加载预训练模型作为起点
angle = AnglE.from_pretrained('SeanLee97/angle-bert-base-uncased-nli-en-v1', max_length=128, pooling_strategy='cls').cuda()
# 加载数据集
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"])
# 数据预处理
train_ds = ds['train'].shuffle().map(AngleDataTokenizer(angle.tokenizer, angle.max_length), num_proc=8)
valid_ds = ds['validation'].map(AngleDataTokenizer(angle.tokenizer, angle.max_length), num_proc=8)
# 开始训练
angle.fit(
train_ds=train_ds,
valid_ds=valid_ds,
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': 20.0,
'angle_w': 1.0,
'cosine_tau': 20,
'ibn_tau': 20,
'angle_tau': 20
},
fp16=True,
logging_steps=100
)
# 评估模型
corrcoef = angle.evaluate(ds['test'])
print('Spearman\'s corrcoef:', corrcoef)
这个例子展示了如何使用AnglE加载预训练模型、准备数据集、进行训练和评估。用户可以根据自己的需求调整训练参数和损失函数权重。
AnglE的突破性成就
AnglE自发布以来,在多个权威基准测试中取得了令人瞩目的成绩:
-
2024年3月8日,使用AnglE训练的mixedbread-ai/mxbai-embed-large-v1模型在MTEB排行榜上以64.68的平均分获得第一名。
-
2023年12月4日,AnglE团队发布的WhereIsAI/UAE-Large-V1模型在MTEB排行榜上以64.64的平均分获得当时的第一名。
-
2023年12月,AnglE在STS Benchmark语义文本相似度任务上取得了最先进的性能。
这些成就充分证明了AnglE在句子嵌入领域的强大实力和广阔前景。
AnglE的应用场景
AnglE可以应用于多种NLP任务,包括但不限于:
-
语义相似度计算: 计算句子或文档之间的语义相似度,可用于信息检索、文本聚类等任务。
-
文本分类: 利用高质量的句子嵌入提升文本分类的准确率。
-
信息检索: 在大规模文本库中快速检索相关文档。
-
问答系统: 提高问答系统的语义理解能力,更准确地匹配问题和答案。
-
文本摘要: 利用句子嵌入选择最具代表性的句子生成摘要。
-
机器翻译: 改善机器翻译系统的语义理解,提高翻译质量。
-
情感分析: 捕捉句子的细微语义差异,提高情感分析的准确性。
AnglE的未来展望
作为一个开源项目,AnglE正在持续发展和改进。未来,AnglE团队计划在以下方面进行更多探索:
-
多语言支持: 扩展AnglE以支持更多语言的句子嵌入训练。
-
跨模态学习: 探索将AnglE应用于图像-文本等跨模态学习任务。
-
效率优化: 进一步提高训练和推理的效率,使AnglE能够处理更大规模的数据和模型。
-
领域适应: 开发更多针对特定领域(如法律、医疗、金融等)的预训练模型。
-
与大型语言模型的结合: 探索如何将AnglE与最新的大型语言模型结合,进一步提升性能。
结语
AnglE作为一个强大的句子嵌入训练和推理框架,为NLP研究人员和工程师提供了一个便捷的工具。它不仅在多个基准测试中取得了最先进的成绩,还提供了灵活的API和丰富的预训练模型,使得高质量的句子嵌入变得触手可及。随着AnglE的不断发展和完善,我们有理由相信它将在推动NLP技术进步方面发挥越来越重要的作用。
无论您是NLP研究人员、学生还是工程师,AnglE都值得您深入探索和尝试。欢迎访问AnglE的GitHub仓库了解更多信息,并为这个开源项目做出贡献。让我们共同期待AnglE在未来为NLP领域带来更多突破性的进展。
