🚀 Snowflake Arctic-Embed-M-Long

Snowflake Arctic-Embed-M-Long 是一套文本嵌入模型，专注于创建针对性能优化的高质量检索模型。该模型在 MTEB/BEIR 排行榜上各尺寸变体均达到了最先进的性能，可用于商业用途。

🚀 快速开始

使用 Sentence Transformers

你可以使用 sentence-transformers 包来使用 snowflake-arctic-embed 模型，示例代码如下：

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("Snowflake/snowflake-arctic-embed-m-long", trust_remote_code=True)

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)
    # Output passages & scores
    print("Query:", query)
    for document, score in doc_score_pairs:
        print(score, document)

运行上述代码，输出结果如下：

Query: what is snowflake?
0.46484852 The Data Cloud!
0.3758855 Mexico City of Course!
Query: Where can I get the best tacos?
0.42407742 Mexico City of Course!
0.36740506 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-long')
model = AutoModel.from_pretrained('Snowflake/snowflake-arctic-embed-m-long', trust_remote_code=True, add_pooling_layer=False, safe_serialization=True)
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)

# Compute token embeddings
with torch.no_grad():
    query_embeddings = model(**query_tokens)[0][:, 0]
    document_embeddings = model(**document_tokens)[0][:, 0]

# normalize embeddings
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)
    #Output passages & scores
    print("Query:", query)
    for document, score in doc_score_pairs:
        print(score, document)

如果你使用的是支持超过 2048 个标记的长上下文模型，请按以下方式初始化模型，这将使用 RPE 支持最多 8192 个标记：

model = AutoModel.from_pretrained('Snowflake/snowflake-arctic-embed-m-long', trust_remote_code=True, safe_serialization=True, rotary_scaling_factor=2)

使用 Transformers.js

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

npm i @xenova/transformers

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

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

// Create feature extraction pipeline
const extractor = await pipeline('feature-extraction', 'Snowflake/snowflake-arctic-embed-m-long', {
    quantized: false, // Comment out this line to use the quantized version
});

// Generate sentence embeddings
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' });

// Compute similarity scores
const [source_embeddings, ...document_embeddings ] = output.tolist();
const similarities = document_embeddings.map(x => dot(source_embeddings, x));
console.log(similarities); // [0.36740492125676116, 0.42407774292046635]

✨ 主要特性

• 高性能检索：snowflake-arctic-embedding 模型在 MTEB/BEIR 排行榜上各尺寸变体均达到了最先进的检索性能。
• 多模型选择：提供了不同尺寸的模型，包括 snowflake-arctic-embed-xs、snowflake-arctic-embed-s、snowflake-arctic-embed-m、snowflake-arctic-embed-m-long 和 snowflake-arctic-embed-l，可以满足不同的性能和精度需求。
• 长上下文支持：snowflake-arctic-embed-m-long 模型基于 nomic-ai/nomic-embed-text-v1-unsupervised 模型，无需使用 RPE 即可支持最多 2048 个标记，使用 RPE 时可扩展到 8192 个标记。

📦 模型介绍

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

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

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

模型性能对比

各尺寸模型在 MTEB 检索得分（NDCG @ 10）对比

与其他闭源嵌入模型对比

各模型详细介绍

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，并提供最准确的检索体验。

📚 常见问题

待补充

💻 使用示例

基础用法

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("Snowflake/snowflake-arctic-embed-m-long", trust_remote_code=True)

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)
    # Output passages & scores
    print("Query:", query)
    for document, score in doc_score_pairs:
        print(score, document)

高级用法

import torch
from transformers import AutoModel, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained('Snowflake/snowflake-arctic-embed-m-long')
model = AutoModel.from_pretrained('Snowflake/snowflake-arctic-embed-m-long', trust_remote_code=True, add_pooling_layer=False, safe_serialization=True)
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)

# Compute token embeddings
with torch.no_grad():
    query_embeddings = model(**query_tokens)[0][:, 0]
    document_embeddings = model(**document_tokens)[0][:, 0]

# normalize embeddings
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)
    #Output passages & scores
    print("Query:", query)
    for document, score in doc_score_pairs:
        print(score, document)

📄 许可证

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

🔧 致谢

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