Snowflake Arctic Embed M

🚀 Snowflake的Arctic-embed-m模型

Snowflake的Arctic-embed-m是一系列文本嵌入模型，专注于创建针对性能优化的高质量检索模型。这些模型在MTEB/BEIR排行榜上的各个尺寸变体中均实现了最先进的性能。

📋 目录

• 📰 新闻动态
• 🧠 模型介绍
• 💻 使用示例
• 📊 评估结果
• 📞 联系我们
• ❓ 常见问题
• 📄 许可证
• 🙏 致谢

📰 新闻动态

• 2024年12月4日：发布了snowflake-arctic-embed-l-v2.0和snowflake-arctic-embed-m-v2.0，这两款新模型专为多语言工作负载设计。它们的性能优于之前版本的Arctic Embed，建议使用这些新版本替代旧版本！
• 2024年7月26日：在arXiv上发布预印本[2407.18887] Embedding And Clustering Your Data Can Improve Contrastive Pretraining。
• 2024年7月18日：发布了snowflake-arctic-embed-m-v1.5，该模型能够生成高度可压缩的嵌入向量，即使每个向量压缩至128字节，仍能保持质量。有关该模型开发的详细信息，请参阅Snowflake工程博客上的发布文章。
• 2024年5月10日：发布了关于Arctic Embed的技术报告。
• 2024年4月16日：发布了snowflake-arctic-embed系列文本嵌入模型。这些模型在各自代表的尺寸配置文件中，检索质量达到了当前的最高水平。技术报告即将发布。更多详细信息，请参考我们的Github仓库：Arctic-Text-Embed。

🧠 模型介绍

snowflake-arctic-embed是一套文本嵌入模型，专注于创建针对性能优化的高质量检索模型。

snowflake-arctic-embedding模型在MTEB/BEIR排行榜的各个尺寸变体中均实现了最先进的性能。评估使用这些脚本进行。如下所示，每个模型尺寸类别与其他顶级模型相比，都实现了SOTA检索准确率。

这些模型通过利用现有的开源文本表示模型（如bert-base-uncased）进行训练，并在多阶段管道中进行训练，以优化其检索性能。首先，使用大量的查询 - 文档对（其中负样本来自批次内）对模型进行训练，预训练利用了约4亿个公共数据集和专有网络搜索数据的混合样本。预训练之后，模型在一个较小的数据集（约100万个样本）上进行长时间训练，该数据集包含从硬负挖掘中得出的查询、正文档和负文档三元组。负样本的挖掘和数据整理对于检索准确率至关重要。详细的技术报告可在此处找到。

除了是优秀的开源模型外，最大的模型snowflake-arctic-embed-l还可以作为闭源嵌入的自然替代品，如下所示。

snowflake-arctic-embed-xs

这个小型模型功能强大。它基于all-MiniLM-L6-v2模型，仅有2200万个参数和384个维度，该模型甚至可以满足最严格的延迟/TCO预算。尽管尺寸较小，但其检索准确率接近具有1亿个参数的模型。

snowflake-arctic-embed-s

该模型基于intfloat/e5-small-unsupervised模型，这个小型模型在不牺牲检索准确率的前提下，实现了较小的尺寸。仅有3300万个参数和384个维度，该模型可以轻松扩展到大型数据集。

snowflake-arctic-embed-m

该模型基于intfloat/e5-base-unsupervised模型，这个中型模型是主力模型，在不降低推理速度的情况下提供了最佳的检索性能。

snowflake-arctic-embed-m-long

该模型基于nomic-ai/nomic-embed-text-v1-unsupervised模型，这个中型模型的长上下文变体非常适合受其他模型常规512个标记上下文限制的工作负载。在不使用RPE的情况下，该模型支持多达2048个标记；使用RPE时，可扩展到8192个标记！

snowflake-arctic-embed-l

该模型基于intfloat/e5-large-unsupervised模型，这个大型模型可以直接替代闭源API，并提供最准确的检索体验。

💻 使用示例

使用Sentence Transformers

你可以使用sentence-transformers包来使用snowflake-arctic-embed模型，如下所示。

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("Snowflake/snowflake-arctic-embed-m")

queries = ['what is snowflake?', 'Where can I get the best tacos?']
documents = ['The Data Cloud!', 'Mexico City of Course!']

query_embeddings = model.encode(queries, prompt_name="query")
document_embeddings = model.encode(documents)

scores = query_embeddings @ document_embeddings.T
for query, query_scores in zip(queries, scores):
    doc_score_pairs = list(zip(documents, query_scores))
    doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)
    # 输出段落和得分
    print("查询:", query)
    for document, score in doc_score_pairs:
        print(score, document)

输出结果：

查询: what is snowflake?
0.2747492 The Data Cloud!
0.19998045 Mexico City of Course!
查询: Where can I get the best tacos?
0.29974818 Mexico City of Course!
0.2344071 The Data Cloud!

使用Huggingface transformers

你可以使用transformers包来使用snowflake-arctic-embed模型，如下所示。为了获得最佳的检索质量，请使用CLS标记嵌入每个文本部分，并使用以下查询前缀（仅用于查询）。

import torch
from transformers import AutoModel, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained('Snowflake/snowflake-arctic-embed-m')
model = AutoModel.from_pretrained('Snowflake/snowflake-arctic-embed-m', add_pooling_layer=False)
model.eval()

query_prefix = 'Represent this sentence for searching relevant passages: '
queries  = ['what is snowflake?', 'Where can I get the best tacos?']
queries_with_prefix = ["{}{}".format(query_prefix, i) for i in queries]
query_tokens = tokenizer(queries_with_prefix, padding=True, truncation=True, return_tensors='pt', max_length=512)

documents = ['The Data Cloud!', 'Mexico City of Course!']
document_tokens =  tokenizer(documents, padding=True, truncation=True, return_tensors='pt', max_length=512)

# 计算标记嵌入
with torch.no_grad():
    query_embeddings = model(**query_tokens)[0][:, 0]
    document_embeddings = model(**document_tokens)[0][:, 0]

# 归一化嵌入
query_embeddings = torch.nn.functional.normalize(query_embeddings, p=2, dim=1)
document_embeddings = torch.nn.functional.normalize(document_embeddings, p=2, dim=1)

scores = torch.mm(query_embeddings, document_embeddings.transpose(0, 1))
for query, query_scores in zip(queries, scores):
    doc_score_pairs = list(zip(documents, query_scores))
    doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)
    # 输出段落和得分
    print("查询:", query)
    for document, score in doc_score_pairs:
        print(score, document)

使用Transformers.js

如果你还没有安装Transformers.js JavaScript库，可以从NPM运行以下命令进行安装：

npm i @xenova/transformers

然后，你可以使用该模型计算嵌入，如下所示：

import { pipeline, dot } from '@xenova/transformers';

// 创建特征提取管道
const extractor = await pipeline('feature-extraction', 'Snowflake/snowflake-arctic-embed-m', {
    quantized: false, // 注释掉此行以使用量化版本
});

// 生成句子嵌入
const sentences = [
    'Represent this sentence for searching relevant passages: Where can I get the best tacos?',
    'The Data Cloud!',
    'Mexico City of Course!',
]
const output = await extractor(sentences, { normalize: true, pooling: 'cls' });

// 计算相似度得分
const [source_embeddings, ...document_embeddings ] = output.tolist();
const similarities = document_embeddings.map(x => dot(source_embeddings, x));
console.log(similarities); // [0.15664823859882132, 0.24481869975470627]

使用Infinity

使用Infinity和Docker进行OpenAI兼容API部署。

docker run --gpus all -v $PWD/data:/app/.cache -p "7997":"7997" \
michaelf34/infinity:0.0.70 \
v2 --model-id Snowflake/snowflake-arctic-embed-m --dtype float16 --batch-size 32 --engine torch --port 7997

❓ 常见问题

待补充。

📞 联系我们

如果您对这个项目有任何问题或建议，请随时打开一个issue或提交一个pull request。 您也可以发送电子邮件给Daniel Campos（daniel.campos@snowflake.com）。

📄 许可证

Arctic采用Apache-2许可证。发布的模型可以免费用于商业目的。

🙏 致谢

我们要感谢开源社区，他们提供了优秀的构建模块，使我们能够开发出这些模型。 我们感谢我们的建模工程师Danmei Xu、Luke Merrick、Gaurav Nuti和Daniel Campos，是他们让这些优秀的模型成为可能。 我们感谢我们的领导Himabindu Pucha、Kelvin So、Vivek Raghunathan和Sridhar Ramaswamy，感谢他们对这项工作的支持。 我们还要感谢开源社区，他们生产了优秀的模型，使我们能够在此基础上进行开发，并使这些发布成为可能。 最后，我们感谢创建BEIR和MTEB基准的研究人员。 正是由于他们不懈的努力，定义了更好的标准，我们才能够提高模型性能。