AnglE 学习资料汇总 - 强大的句子嵌入训练和推理框架

AnglE

AnglE简介

AnglE是一个用于训练和推理强大句子嵌入的开源框架。它源自论文《AnglE: Angle-optimized Text Embeddings》，允许用户通过几行代码就能训练出最先进的基于BERT或LLM的句子嵌入。AnglE还是一个通用的句子嵌入推理框架，可以推理各种基于transformer的句子嵌入。

AnglE Logo

主要特性

AnglE具有以下主要特性:

支持多种损失函数:AnglE损失、对比损失、CoSENT损失等
支持多种主干网络:基于BERT的模型、基于LLM的模型等
支持单GPU和多GPU训练
提供了预训练模型
易于使用的API

安装

使用pip安装AnglE:

python -m pip install -U angle_emb

快速使用

推理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))

推理LLM模型

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

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([
    {'text': 'The weather is great!'},
    {'text': 'The weather is very good!'},
    {'text': 'i am going to bed'}
], prompt=Prompts.A)

for i, dv1 in enumerate(doc_vecs):
    for dv2 in doc_vecs[i+1:]:
        print(cosine_similarity(dv1, dv2))

训练自定义模型

AnglE提供了便捷的API用于训练自定义模型:

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
)

学习资源

官方文档 - 详细的使用说明和API文档
GitHub仓库 - 源代码和更多示例
Hugging Face模型库 - 预训练模型
论文 - 技术细节和原理

AnglE架构

总结

AnglE是一个强大而灵活的句子嵌入框架,适用于各种NLP任务。通过本文介绍的学习资源,读者可以快速上手AnglE,并将其应用到自己的项目中。无论是使用预训练模型还是训练自定义模型,AnglE都提供了简单易用的API。希望这些资源能帮助你充分利用AnglE的功能,提升NLP任务的性能。