Snowflake/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 All embeddings are created using Snowflake's 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… See the full description on the dataset page: https://huggingface.co/datasets/Snowflake/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 All embeddings are created using Snowflake's 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
Recall@100
Recall@1000
Loading the dataset
Loading the document embeddings
You can either load the dataset like this:
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:
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):
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")