msmarco-v2.1-snowflake-arctic-embed-m-v2.0

# Snowflake Arctic Embed M V2.0 Embeddings for MSMARCO V2.1 for TREC-RAG

This dataset contains the embeddings for the MSMARCO-V2.1 dataset which is used as the corpora for [TREC RAG](https://trec-rag.github.io/)
All embeddings are created using [Snowflake's Arctic Embed M v2.0](https://huggingface.co/Snowflake/snowflake-arctic-embed-m-v2.0) and are intended to serve as a simple baseline for dense retrieval-based methods.
Note, that the embeddings are not normalized so you will need to normalize them before usage.

## Retrieval Performance
Retrieval performance for the TREC DL21-23, MSMARCOV2-Dev and Raggy Queries can be found below with BM25 as a baseline. For both systems, retrieval is at the segment level and Doc Score = Max (passage score).
Retrieval is done via a dot product and happens in BF16.

### NDCG@10
| Dataset | BM25 | Arctic-M-V2.0 (768 Dimensions) |
|---|---|---|
| Deep Learning 2021 | 0.5778 | |
| Deep Learning 2022 | 0.3576 | |
| Deep Learning 2023 | 0.3356 ||
| msmarcov2-dev | N/A | |
| msmarcov2-dev2 | N/A | |
| Raggy Queries | 0.4227 | |
| RAG 2024 | | |
### Recall@100
| Dataset | BM25 | Arctic-M-V2.0 (768 Dimensions) |
|---|---|---|
| Deep Learning 2021 | 0.3811 ||
| Deep Learning 2022 | 0.233 | |
| Deep Learning 2023 | 0.3049 | |
| msmarcov2-dev | 0.6683 | |
| msmarcov2-dev2 | 0.6771 | |
| Raggy Queries | 0.2807 | |
| RAG 2024 | | |

### Recall@1000
| Dataset | BM25 | Arctic-M-V2.0 (768 Dimensions) |
|---|---|---|
| Deep Learning 2021 | 0.7115 | |
| Deep Learning 2022 | 0.479 | |
| Deep Learning 2023 | 0.5852 | |
| msmarcov2-dev | 0.8528 | |
| msmarcov2-dev2 | 0.8577 | |
| Raggy Queries | 0.5745 | |
| RAG 2024 | | |

## Loading the dataset

### Loading the document embeddings

You can either load the dataset like this:
```python
from datasets import load_dataset
docs = load_dataset("Snowflake/msmarco-v2.1-snowflake-arctic-embed-m-v2.0", split="train")
```

Or you can also stream it without downloading it before:
```python
from datasets import load_dataset
docs = load_dataset("Snowflake/msmarco-v2.1-snowflake-arctic-embed-m-v2.0", split="train", streaming=True)
for doc in docs:
doc_id = j['docid']
url = doc['url']
text = doc['text']
emb = doc['embedding']
```

Note, The full dataset corpus is ~ 620GB so it will take a while to download and may not fit on some devices/

## Search
A full search example (on the first 1,000 paragraphs):
```python
from datasets import load_dataset
import torch
from transformers import AutoModel, AutoTokenizer
import numpy as np

top_k = 100
docs_stream = load_dataset("Snowflake/msmarco-v2.1-snowflake-arctic-embed-m-v2.0",split="train", streaming=True)

docs = []
doc_embeddings = []

for doc in docs_stream:
docs.append(doc)
doc_embeddings.append(doc['embedding'])
if len(docs) >= top_k:
break

doc_embeddings = np.asarray(doc_embeddings)

tokenizer = AutoTokenizer.from_pretrained('Snowflake/snowflake-arctic-embed-m-v2.0')
model = AutoModel.from_pretrained('Snowflake/snowflake-arctic-embed-m-v2.0', add_pooling_layer=False)
model.eval()

query_prefix = 'Represent this sentence for searching relevant passages: '
queries = ['how do you clean smoke off walls']
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)

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

# normalize embeddings
query_embeddings = torch.nn.functional.normalize(query_embeddings, p=2, dim=1)
doc_embeddings = torch.nn.functional.normalize(doc_embeddings, p=2, dim=1)

# Compute dot score between query embedding and document embeddings
dot_scores = np.matmul(query_embeddings, doc_embeddings.transpose())[0]
top_k_hits = np.argpartition(dot_scores, -top_k)[-top_k:].tolist()

# Sort top_k_hits by dot score
top_k_hits.sort(key=lambda x: dot_scores[x], reverse=True)

# Print results
print("Query:", queries[0])
for doc_id in top_k_hits:
print(docs[doc_id]['doc_id'])
print(docs[doc_id]['text'])
print(docs[doc_id]['url'], "\n")
```

# 面向TREC-RAG的MSMARCO V2.1数据集Snowflake Arctic Embed M V2.0 嵌入向量

本数据集包含MSMARCO-V2.1数据集的嵌入向量，该数据集被用作[TREC RAG](https://trec-rag.github.io/)的语料库。所有嵌入向量均采用[Snowflake Arctic Embed M V2.0](https://huggingface.co/Snowflake/snowflake-arctic-embed-m-v2.0)生成，旨在为基于稠密检索的方法提供简易基准测试基线。请注意，本数据集的嵌入向量未经过归一化处理，使用前需自行完成归一化操作。

## 检索性能

以下为以BM25作为基准基线时，TREC DL21-23、MSMARCOV2-Dev以及Raggy Queries的检索性能表现。两类系统均采用段落级检索策略，文档得分=最高段落得分。检索过程通过点积运算完成，并以BF16精度运行。

### NDCG@10
| 数据集 | BM25 | Arctic-M-V2.0（768维） |
|---|---|---|
| 深度学习2021 | 0.5778 | |
| 深度学习2022 | 0.3576 | |
| 深度学习2023 | 0.3356 | |
| msmarcov2-dev | N/A | |
| msmarcov2-dev2 | N/A | |
| Raggy Queries | 0.4227 | |
| RAG 2024 | | |
### Recall@100
| 数据集 | BM25 | Arctic-M-V2.0（768维） |
|---|---|---|
| 深度学习2021 | 0.3811 ||
| 深度学习2022 | 0.233 | |
| 深度学习2023 | 0.3049 | |
| msmarcov2-dev | 0.6683 | |
| msmarcov2-dev2 | 0.6771 | |
| Raggy Queries | 0.2807 | |
| RAG 2024 | | |

### Recall@1000
| 数据集 | BM25 | Arctic-M-V2.0（768维） |
|---|---|---|
| 深度学习2021 | 0.7115 | |
| 深度学习2022 | 0.479 | |
| 深度学习2023 | 0.5852 | |
| msmarcov2-dev | 0.8528 | |
| msmarcov2-dev2 | 0.8577 | |
| Raggy Queries | 0.5745 | |
| RAG 2024 | | |

## 数据集加载

### 文档嵌入向量加载

你可以通过以下方式加载数据集：
python
from datasets import load_dataset
docs = load_dataset("Snowflake/msmarco-v2.1-snowflake-arctic-embed-m-v2.0", split="train")

或者也可以在不提前下载的情况下流式加载数据集：
python
from datasets import load_dataset
docs = load_dataset("Snowflake/msmarco-v2.1-snowflake-arctic-embed-m-v2.0", split="train", streaming=True)
for doc in docs:
doc_id = j['docid']
url = doc['url']
text = doc['text']
emb = doc['embedding']

请注意，完整数据集语料库大小约为620GB，下载耗时较长，且部分设备可能无法容纳该数据集。

## 检索示例

以下为基于前1000个段落的完整检索示例：
python
from datasets import load_dataset
import torch
from transformers import AutoModel, AutoTokenizer
import numpy as np

top_k = 100
docs_stream = load_dataset("Snowflake/msmarco-v2.1-snowflake-arctic-embed-m-v2.0",split="train", streaming=True)

docs = []
doc_embeddings = []

for doc in docs_stream:
docs.append(doc)
doc_embeddings.append(doc['embedding'])
if len(docs) >= top_k:
break

doc_embeddings = np.asarray(doc_embeddings)

tokenizer = AutoTokenizer.from_pretrained('Snowflake/snowflake-arctic-embed-m-v2.0')
model = AutoModel.from_pretrained('Snowflake/snowflake-arctic-embed-m-v2.0', add_pooling_layer=False)
model.eval()

query_prefix = 'Represent this sentence for searching relevant passages: '
queries = ['how do you clean smoke off walls']
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)

# 计算Token嵌入向量
with torch.no_grad():
query_embeddings = model(**query_tokens)[0][:, 0]

# 对嵌入向量进行归一化
query_embeddings = torch.nn.functional.normalize(query_embeddings, p=2, dim=1)
doc_embeddings = torch.nn.functional.normalize(doc_embeddings, p=2, dim=1)

# 计算查询嵌入向量与文档嵌入向量之间的点积得分
dot_scores = np.matmul(query_embeddings, doc_embeddings.transpose())[0]
top_k_hits = np.argpartition(dot_scores, -top_k)[-top_k:].tolist()

# 按照点积得分对top_k_hits进行排序
top_k_hits.sort(key=lambda x: dot_scores[x], reverse=True)

# 打印检索结果
print("查询语句:", queries[0])
for doc_id in top_k_hits:
print(docs[doc_id]['doc_id'])
print(docs[doc_id]['text'])
print(docs[doc_id]['url'], "
")

搜集汇总

背景与挑战

背景概述

该数据集包含MSMARCO-V2.1的嵌入向量，由Snowflake Arctic Embed M v2.0模型生成，用于TREC-RAG任务作为密集检索方法的基准。嵌入维度为768且未归一化，使用前需处理，数据集总大小约90.54GB，采用Apache 2.0许可证。