Universal AnglE 嵌入模型

📢 WhereIsAI/UAE-Large-V1 采用 MIT 许可证。您可以在任何场景下自由使用。 如果您将其用于学术论文，可以通过 👉 引用信息 进行引用。

🤝 关注我们：

• GitHub：https://github.com/SeanLee97/AnglE
• 预印本论文：AnglE-optimized Text Embeddings
• 会议论文：AoE: Angle-optimized Embeddings for Semantic Textual Similarity (ACL24)
• 📘 文档：https://angle.readthedocs.io/en/latest/index.html

欢迎使用 AnglE 训练和推理强大的句子嵌入模型。

🏆 成就

• 📅 2024年5月16日 | AnglE 的论文被 ACL 2024 主会接收
• 📅 2023年12月4日 | 🔥 我们的通用英文句子嵌入模型 WhereIsAI/UAE-Large-V1 在 MTEB 排行榜 上以平均分 64.64 取得了 SOTA！

🧑‍🤝‍🧑 同系列模型：

• WhereIsAI/UAE-Code-Large-V1：该模型可用于代码或 GitHub Issue 相似度测量。

使用方法

1. angle_emb

python -m pip install -U angle-emb

1. 非检索任务

无需指定任何提示词。

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'
], normalize_embedding=True)

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

2. 检索任务

对于检索任务，请使用提示词 Prompts.C 处理查询（不要用于文档）。

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(Prompts.C.format(text='what is the weather?'))
doc_vecs = angle.encode([
    'The weather is great!',
    'it is rainy today.',
    'i am going to bed'
])

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

2. Sentence Transformer

from angle_emb import Prompts
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("WhereIsAI/UAE-Large-V1").cuda()

qv = model.encode(Prompts.C.format(text='what is the weather?'))
doc_vecs = model.encode([
    'The weather is great!',
    'it is rainy today.',
    'i am going to bed'
])

for dv in doc_vecs:
    print(1 - spatial.distance.cosine(qv, dv))

3. Infinity

Infinity 是一个采用 MIT 许可证的 OpenAI 兼容部署服务器。

docker run --gpus all -v $PWD/data:/app/.cache -p \"7997\":\"7997\" \
michaelf34/infinity:latest \
v2 --model-id WhereIsAI/UAE-Large-V1 --revision \"369c368f70f16a613f19f5598d4f12d9f44235d4\" --dtype float16 --batch-size 32 --device cuda --engine torch --port 7997

引用

如果您使用了我们的预训练模型，欢迎通过以下方式引用我们的工作：

@article{li2023angle,
  title={AnglE-optimized Text Embeddings},
  author={Li, Xianming and Li, Jing},
  journal={arXiv preprint arXiv:2309.12871},
  year={2023}
}