AnglE:一个强大的句子嵌入训练和推理框架

AnglE:开创句子嵌入的新纪元

在自然语言处理(NLP)领域,高质量的句子嵌入对于各种下游任务至关重要。近日,一个名为AnglE的开源框架引起了研究人员的广泛关注。AnglE是一个用于训练和推理强大句子嵌入的框架,它基于创新的角度优化损失函数,能够轻松训练出最先进的BERT和LLM(大型语言模型)句子嵌入模型。本文将详细介绍AnglE的特性、使用方法及其在NLP领域取得的重要成就。

AnglE的核心特性

AnglE框架具有以下几个突出的特性:

1. 多种损失函数: AnglE支持多种先进的损失函数,包括:

   • AnglE loss: 基于角度优化的创新损失函数
   • Contrastive loss: 对比损失
   • CoSENT loss: 余弦相似度损失
   • Espresso loss: 之前称为2DMSE的损失函数

2. 灵活的backbone支持: AnglE可以使用多种预训练模型作为backbone,包括:

   • BERT系列模型(BERT、RoBERTa、ELECTRA、ALBERT等)
   • 大型语言模型(LLaMA、Mistral、Qwen等)
   • 双向LLM模型(基于BiLLM项目)

3. 高效训练: AnglE支持单GPU和多GPU训练,可以充分利用硬件资源加速模型训练过程。

4. 易用性: AnglE提供了简洁的API,只需几行代码即可完成模型的训练和推理。

5. 性能卓越: 基于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项目提供了多个高质量的预训练模型,适用于不同的场景:

1. 通用英语模型:

   • WhereIsAI/UAE-Large-V1: 512 tokens, cls pooling, 适用于一般英语场景

2. 代码相似度模型:

   • WhereIsAI/UAE-Code-Large-V1: 512 tokens, cls pooling, 专门用于代码相似度任务

3. 医学相似度模型:

   • WhereIsAI/pubmed-angle-base-en: 512 tokens, cls pooling
   • WhereIsAI/pubmed-angle-large-en: 512 tokens, cls pooling

4. 基于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自发布以来,在多个权威基准测试中取得了令人瞩目的成绩:

1. 2024年3月8日,使用AnglE训练的mixedbread-ai/mxbai-embed-large-v1模型在MTEB排行榜上以64.68的平均分获得第一名。

2. 2023年12月4日,AnglE团队发布的WhereIsAI/UAE-Large-V1模型在MTEB排行榜上以64.64的平均分获得当时的第一名。

3. 2023年12月,AnglE在STS Benchmark语义文本相似度任务上取得了最先进的性能。

这些成就充分证明了AnglE在句子嵌入领域的强大实力和广阔前景。

AnglE的应用场景

AnglE可以应用于多种NLP任务,包括但不限于:

1. 语义相似度计算: 计算句子或文档之间的语义相似度,可用于信息检索、文本聚类等任务。

2. 文本分类: 利用高质量的句子嵌入提升文本分类的准确率。

3. 信息检索: 在大规模文本库中快速检索相关文档。

4. 问答系统: 提高问答系统的语义理解能力,更准确地匹配问题和答案。

5. 文本摘要: 利用句子嵌入选择最具代表性的句子生成摘要。

6. 机器翻译: 改善机器翻译系统的语义理解,提高翻译质量。

7. 情感分析: 捕捉句子的细微语义差异,提高情感分析的准确性。

AnglE的未来展望

作为一个开源项目,AnglE正在持续发展和改进。未来,AnglE团队计划在以下方面进行更多探索:

1. 多语言支持: 扩展AnglE以支持更多语言的句子嵌入训练。

2. 跨模态学习: 探索将AnglE应用于图像-文本等跨模态学习任务。

3. 效率优化: 进一步提高训练和推理的效率,使AnglE能够处理更大规模的数据和模型。

4. 领域适应: 开发更多针对特定领域(如法律、医疗、金融等)的预训练模型。

5. 与大型语言模型的结合: 探索如何将AnglE与最新的大型语言模型结合,进一步提升性能。

结语

AnglE作为一个强大的句子嵌入训练和推理框架,为NLP研究人员和工程师提供了一个便捷的工具。它不仅在多个基准测试中取得了最先进的成绩,还提供了灵活的API和丰富的预训练模型,使得高质量的句子嵌入变得触手可及。随着AnglE的不断发展和完善,我们有理由相信它将在推动NLP技术进步方面发挥越来越重要的作用。

无论您是NLP研究人员、学生还是工程师,AnglE都值得您深入探索和尝试。欢迎访问AnglE的GitHub仓库了解更多信息,并为这个开源项目做出贡献。让我们共同期待AnglE在未来为NLP领域带来更多突破性的进展。