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許可證。發佈的模型可以免費用於商業目的。