msmarco-v2.1-snowflake-arctic-embed-l

# Snowflake Arctic Embed L 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 L](https://huggingface.co/Snowflake/snowflake-arctic-embed-l) and are intended to serve as a simple baseline for dense retrieval-based methods.

## 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 dot product and happens in BF16.
### NDCG@10
| Dataset | BM25 | Snowflake Arctic Embed L |
|---|---|---|
| Deep Learning 2021 | 0.5778 | 0.70682 |
| Deep Learning 2022 | 0.3576 | 0.5444 |
| Deep Learning 2023 | 0.3356 | 0.47372 |
| msmarcov2-dev | N/A | 0.35844 |
| msmarcov2-dev2 | N/A | 0.35821 |
| Raggy Queries | 0.4227 | 0.57759 |

### Recall@100
| Dataset | BM25 | Snowflake Arctic Embed L |
|---|---|---|
| Deep Learning 2021 | 0.3811 | 0.41361 |
| Deep Learning 2022 | 0.233 | 0.31351 |
| Deep Learning 2023 | 0.3049 | 0.34793 |
| msmarcov2-dev | 0.6683 | 0.85131 |
| msmarcov2-dev2 | 0.6771 | 0.84767 |
| Raggy Queries | 0.2807 | 0.36228 |

### Recall@1000
| Dataset | BM25 | Snowflake Arctic Embed L |
|---|---|---|
| Deep Learning 2021 | 0.7115 | 0.7193 |
| Deep Learning 2022 | 0.479 | 0.54566 |
| Deep Learning 2023 | 0.5852 | 0.59577 |
| msmarcov2-dev | 0.8528 | 0.93966 |
| msmarcov2-dev2 | 0.8577 | 0.93947 |
| Raggy Queries | 0.5745 | 0.63092 |

## 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-l", 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-l", 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-l",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-l')
model = AutoModel.from_pretrained('Snowflake/snowflake-arctic-embed-l', 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 L嵌入向量

本数据集包含用作[TREC RAG](https://trec-rag.github.io/)语料库的MSMARCO-V2.1数据集的嵌入向量。所有嵌入向量均由[Snowflake Arctic Embed L](https://huggingface.co/Snowflake/snowflake-arctic-embed-l)生成，旨在作为基于稠密检索方法的简易基准基线。

## 检索性能

TREC DL21-23、MSMARCOV2开发集以及Raggy Queries的检索性能可参见下文，以BM25作为基准基线。两种系统均采用分段级检索，文档得分 = 最大段落得分。检索通过点积计算，以BF16精度执行。

### NDCG@10
| 数据集 | BM25 | Snowflake Arctic Embed L |
|---|---|---|
| 深度学习2021 | 0.5778 | 0.70682 |
| 深度学习2022 | 0.3576 | 0.5444 |
| 深度学习2023 | 0.3356 | 0.47372 |
| MSMARCOV2开发集 | N/A | 0.35844 |
| MSMARCOV2开发集2 | N/A | 0.35821 |
| Raggy Queries | 0.4227 | 0.57759 |

### 召回率@100
| 数据集 | BM25 | Snowflake Arctic Embed L |
|---|---|---|
| 深度学习2021 | 0.3811 | 0.41361 |
| 深度学习2022 | 0.233 | 0.31351 |
| 深度学习2023 | 0.3049 | 0.34793 |
| MSMARCOV2开发集 | 0.6683 | 0.85131 |
| MSMARCOV2开发集2 | 0.6771 | 0.84767 |
| Raggy Queries | 0.2807 | 0.36228 |

### 召回率@1000
| 数据集 | BM25 | Snowflake Arctic Embed L |
|---|---|---|
| 深度学习2021 | 0.7115 | 0.7193 |
| 深度学习2022 | 0.479 | 0.54566 |
| 深度学习2023 | 0.5852 | 0.59577 |
| MSMARCOV2开发集 | 0.8528 | 0.93966 |
| MSMARCOV2开发集2 | 0.8577 | 0.93947 |
| Raggy Queries | 0.5745 | 0.63092 |

## 数据集加载

### 文档嵌入向量加载

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

或者也可以在不提前下载的情况下流式加载：
python
from datasets import load_dataset
docs = load_dataset("Snowflake/msmarco-v2.1-snowflake-arctic-embed-l", 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-l",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-l')
model = AutoModel.from_pretrained('Snowflake/snowflake-arctic-embed-l', 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'], "
")