Snowflake Arctic Embed Xs

🚀 Snowflake Arctic Embed 文本嵌入模型

Snowflake Arctic Embed 是一系列文本嵌入模型，專注於打造高性能的高質量檢索模型。這些模型在 MTEB/BEIR 排行榜上的各尺寸變體中均達到了最先進的檢索性能。

🚀 快速開始

使用 Sentence Transformers

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

from sentence_transformers import SentenceTransformer

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

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.57515126 The Data Cloud!
0.45798576 Mexico City of Course!
Query: Where can I get the best tacos?
0.5636022 Mexico City of Course!
0.5044898 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-xs')
model = AutoModel.from_pretrained('Snowflake/snowflake-arctic-embed-xs', 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)

# 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)

使用 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-xs', {
    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.5044895661144148, 0.5636021124426508]

✨ 主要特性

• 高性能檢索：snowflake-arctic-embed 模型在 MTEB/BEIR 排行榜的各尺寸變體中均達到了最先進的檢索性能。
• 多尺寸選擇：提供了不同尺寸的模型，包括 xs、s、m、m-long 和 l，可以根據不同的需求選擇合適的模型。
• 支持多語言：部分模型支持多語言工作負載。
• 可壓縮嵌入向量：部分模型能夠生成高度可壓縮的嵌入向量，即使壓縮到每個向量 128 字節，仍能保持質量。

📦 安裝指南

使用 Sentence Transformers

pip install sentence-transformers

使用 Huggingface transformers

pip install transformers

使用 Transformers.js

npm i @xenova/transformers

💻 使用示例

基礎用法

上述“快速開始”部分的代碼示例展示瞭如何使用不同的庫調用 snowflake-arctic-embed 模型進行文本嵌入和相似度計算。

📚 詳細文檔

模型介紹

snowflake-arctic-embed 是一系列文本嵌入模型，專注於創建針對性能優化的高質量檢索模型。這些模型通過利用現有的開源文本表示模型（如 bert-base-uncased）進行訓練，並在多階段管道中進行訓練，以優化其檢索性能。

模型評估

雪片 - 北極嵌入 模型在 MTEB/BEIR 排行榜的各尺寸變體中均達到了最先進的性能。評估使用了 這些腳本 進行。

模型參數

與其他模型對比

與閉源 API 模型對比

snowflake-arctic-embed-xs 與其他小模型對比

snowflake-arctic-embed-s 與其他小模型對比

snowflake-arctic-embed-m 與其他中等模型對比

snowflake-arctic-embed-m-long 與其他長上下文模型對比

snowflake-arctic-embed-l 與其他大模型對比

🔧 技術細節

模型通過利用現有的開源文本表示模型（如 bert-base-uncased）進行訓練，並在多階段管道中進行訓練，以優化其檢索性能。首先，模型使用大量的查詢 - 文檔對進行訓練，其中負樣本是在批次內推導出來的 - 預訓練利用了約 4 億個樣本，這些樣本來自公共數據集和專有網絡搜索數據的混合。預訓練之後，模型在一個較小的數據集（約 100 萬個樣本）上進行長時間訓練，該數據集包含查詢、正文檔和負文檔的三元組，負樣本通過困難有害挖掘得到。負樣本的挖掘和數據整理對於檢索準確性至關重要。詳細的技術報告可以在 這裡 找到。

📄 許可證

Arctic 採用 Apache-2 許可證。發佈的模型可以免費用於商業目的。

項目動態

• 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。

常見問題解答

待補充。

聯繫我們

如果您對這個項目有任何問題或建議，請隨時打開一個 issue 或 pull request。您也可以發送電子郵件至 Daniel Campos（daniel.campos@snowflake.com）。

致謝

我們要感謝開源社區，他們提供了優秀的基礎構建模塊，使我們能夠開發出這些模型。我們感謝我們的建模工程師 Danmei Xu、Luke Merrick、Gaurav Nuti 和 Daniel Campos，是他們讓這些優秀的模型成為可能。我們感謝我們的領導 Himabindu Pucha、Kelvin So、Vivek Raghunathan 和 Sridhar Ramaswamy，感謝他們對這項工作的支持。我們也感謝開源社區，他們開發了優秀的模型，使我們能夠在此基礎上進行構建併發布這些模型。最後，我們感謝創建 BEIR 和 MTEB 基準的研究人員。正是由於他們不懈的努力來定義更好的標準，我們才能夠提高模型性能。