AnglE - 基于BERT和LLM模型提升句子嵌入效果的方法

AnglE 项目介绍

项目概述

AnglE 是由 Mixedbread 赞助的一个项目，专注于优化文本嵌入表示。通过这个项目，用户可以轻松训练出优秀的基于 BERT 或大型语言模型（LLM）的句子嵌入。同时，AnglE 也提供了一个通用句子嵌入推理框架，可以对多种基于 Transformer 的句子嵌入进行推理。详细的使用文档可以访问：AnglE 文档。

特点

损失函数：

AnglE 损失函数
对比损失函数
CoSENT 损失函数
Espresso 损失函数

模型架构：

基于 BERT 的模型（如 BERT、RoBERTa、ELECTRA、ALBERT 等）
基于 LLM 的模型（如 LLaMA、Mistral、Qwen 等）
双向 LLM 模型（如 LLaMA、OpenELMo 等）

训练能力：

支持单 GPU 训练
支持多 GPU 训练

成就

AnglE 项目的论文已被 ACL 2024 的主会议接受。
使用 AnglE 项目训练的混合面包嵌入模型在多个语义文本相似性任务中取得了领先的成绩。

预训练模型

项目中提供了一些官方预训练模型，适用于不同的应用场景，例如通用的英语句子嵌入模型和医学相似性等特定领域模型。

快速开始

安装

要使用 AnglE，只需运行以下命令：

python -m pip install -U angle-emb

推理示例

AnglE 提供了便捷的 API 接口，用户可以通过简单的代码进行推理。以下是基于 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()
qv = angle.encode({'text': '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)

for dv in doc_vecs:
    print(cosine_similarity(qv[0], dv))

无提示词的推理：

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([
    '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))

自定义训练

用户可以根据需要准备数据集，自行训练适合自己需求的模型。为了加速推理过程，推荐使用 Mixedbread 的 batched 库。

python -m pip install batched

关注与联系

对于任何问题或建议，欢迎通过电子邮件联系项目团队：xmlee97@gmail.com。项目使用 MIT 许可证开放源代码，更多详情请参阅各个预训练模型的许可证说明。