Snowflake Arctic Embed S

🚀 Snowflake的Arctic-embed-s模型

Snowflake的Arctic-embed-s是一款专注于文本嵌入的模型，能够创建高质量的检索模型，在性能方面进行了优化，为文本检索任务提供了高效且准确的解决方案。

🚀 快速开始

使用Sentence Transformers

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

from sentence_transformers import SentenceTransformer

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

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:", query)
    for document, score in doc_score_pairs:
        print(score, document)

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

Query: what is snowflake?
0.533809 The Data Cloud!
0.49207097 Mexico City of Course!
Query: Where can I get the best tacos?
0.56592476 Mexico City of Course!
0.48255116 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-s')
model = AutoModel.from_pretrained('Snowflake/snowflake-arctic-embed-s', 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:", 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-s', {
    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.48255123876493394, 0.5659250100112143]

✨ 主要特性

• 高性能检索：snowflake-arctic-embedding模型在MTEB/BEIR排行榜的各个尺寸变体中均实现了最先进的性能。通过使用脚本进行评估，与其他顶级模型相比，每个模型尺寸类别都实现了SOTA检索准确率。
• 多阶段训练：模型通过利用现有的开源文本表示模型（如bert-base-uncased）进行训练，并在多阶段管道中进行训练，以优化其检索性能。首先，使用大量的查询 - 文档对进行训练，其中负样本是在批次内推导出来的；预训练利用了大约4亿个公共数据集和专有网络搜索数据的混合样本。预训练之后，在较小的数据集（大约100万个样本）上进行长时间训练，进一步优化模型，该数据集包含从困难有害挖掘中得出的查询、正文档和负文档的三元组。负样本的挖掘和数据整理对于检索准确率至关重要。
• 多种模型尺寸：提供了多种尺寸的模型，包括snowflake-arctic-embed-xs、snowflake-arctic-embed-s、snowflake-arctic-embed-m、snowflake-arctic-embed-m-long和snowflake-arctic-embed-l，可以根据不同的需求和预算进行选择。

📦 安装指南

使用不同的库调用模型时，安装依赖的方式如下：

• Sentence Transformers：使用pip install sentence-transformers安装。
• Huggingface transformers：使用pip install transformers安装。
• Transformers.js：使用npm i @xenova/transformers安装。

💻 使用示例

基础用法

上述使用不同库调用模型的代码示例均为基础用法，展示了如何使用不同的库来调用snowflake-arctic-embed-s模型进行文本嵌入和相似度计算。

高级用法

在实际应用中，可以根据具体需求对基础用法进行扩展，例如：

• 批量处理：可以将多个查询和文档作为一个批次进行处理，以提高处理效率。
• 多语言支持：如果需要处理多语言文本，可以选择支持多语言的模型，并进行相应的调整。

📚 详细文档

模型介绍

snowflake-arctic-embed是一套专注于创建高性能检索模型的文本嵌入模型。以下是不同尺寸模型的详细信息：

与其他闭源嵌入模型相比，最大的模型snowflake-arctic-embed-l可以作为自然替代方案：

各尺寸模型特点

• snowflake-arctic-embed-xs：基于all-MiniLM-L6-v2模型，仅有2200万个参数和384个维度，适合对延迟和总体拥有成本（TCO）要求严格的场景。尽管尺寸小，但其检索准确率接近拥有1亿参数的模型。 | 模型名称 | MTEB检索分数 (NDCG @ 10) | | ---- | ---- | | snowflake-arctic-embed-xs | 50.15 | | GIST-all-MiniLM-L6-v2 | 45.12 | | gte-tiny | 44.92 | | all-MiniLM-L6-v2 | 41.95 | | bge-micro-v2 | 42.56 |

• snowflake-arctic-embed-s：基于intfloat/e5-small-unsupervised模型，在小尺寸的情况下不牺牲检索准确率。仅有3300万个参数和384个维度，便于扩展到大型数据集。 | 模型名称 | MTEB检索分数 (NDCG @ 10) | | ---- | ---- | | snowflake-arctic-embed-s | 51.98 | | bge-small-en-v1.5 | 51.68 | | Cohere-embed-english-light-v3.0 | 51.34 | | text-embedding-3-small | 51.08 | | e5-small-v2 | 49.04 |

• snowflake-arctic-embed-m：基于intfloat/e5-base-unsupervised模型，是提供最佳检索性能且不影响推理速度的主力模型。 | 模型名称 | MTEB检索分数 (NDCG @ 10) | | ---- | ---- | | snowflake-arctic-embed-m | 54.90 | | bge-base-en-v1.5 | 53.25 | | nomic-embed-text-v1.5 | 53.25 | | GIST-Embedding-v0 | 52.31 | | gte-base | 52.31 |

• snowflake-arctic-embed-m-long：基于nomic-ai/nomic-embed-text-v1-unsupervised模型，是中型模型的长上下文变体，适合受其他模型常规512个标记上下文限制的工作负载。不使用RPE时，该模型支持多达2048个标记；使用RPE时，可扩展到8192个标记。 | 模型名称 | MTEB检索分数 (NDCG @ 10) | | ---- | ---- | | snowflake-arctic-embed-m-long | 54.83 | | nomic-embed-text-v1.5 | 53.01 | | nomic-embed-text-v1 | 52.81 |

• snowflake-arctic-embed-l：基于intfloat/e5-large-unsupervised模型，可直接替代闭源API，提供最准确的检索体验。 | 模型名称 | MTEB检索分数 (NDCG @ 10) | | ---- | ---- | | snowflake-arctic-embed-l | 55.98 | | UAE-Large-V1 | 54.66 | | bge-large-en-v1.5 | 54.29 | | mxbai-embed-large-v1 | 54.39 | | e5-Large-v2 | 50.56 |

技术报告

详细的技术报告可以在这里找到。

🔧 技术细节

训练过程

模型的训练分为两个主要阶段：

1. 预训练：使用大量的查询 - 文档对进行训练，其中负样本是在批次内推导出来的。预训练利用了大约4亿个公共数据集和专有网络搜索数据的混合样本。
2. 优化训练：在较小的数据集（大约100万个样本）上进行长时间训练，进一步优化模型，该数据集包含从困难有害挖掘中得出的查询、正文档和负文档的三元组。

负样本挖掘

负样本的挖掘和数据整理对于检索准确率至关重要。通过困难有害挖掘，可以找到与查询更相关的负样本，从而提高模型的区分能力。

📄 许可证

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