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基準的研究人員。 正是由於他們不懈的努力，定義了更好的標準，我們才能夠提高模型性能。