Coder Social home page Coder Social logo

arctic-embed's Introduction

Snowflake Arctic-Embed Models

Snowflake's Artic-embed

News

07/26/2024: Release preprint [2407.18887] Embedding And Clustering Your Data Can Improve Contrastive Pretraining on arXiv.

07/18/2024: Release v1.5 of our m-sized model, snowflake-arctic-embed-m-v1.5. Details about the development of this model are available in the launch post on the Snowflake engineering blog.

05/10/2024: Release the technical report on Arctic Embed

04/16/2024: Release the ** snowflake-arctic-embed ** family of text embedding models. The releases are state-of-the-art for Retrieval quality at each of their representative size profiles. Technical Report is coming shortly. For more details, please refer to our Github: Arctic-Embed.

Models

snowflake-arctic-embed is a suite of text embedding models that focuses on creating high-quality retrieval models optimized for performance.

The snowflake-arctic-embedding models achieve state-of-the-art performance on the MTEB/BEIR leaderboard for each of their size variants. Evaluation is performed using these scripts. As shown below, each class of model size achieves SOTA retrieval accuracy compared to other top models.

The models are trained by leveraging existing open-source text representation models, such as bert-base-uncased, and are trained in a multi-stage pipeline to optimize their retrieval performance. First, the models are trained with large batches of query-document pairs where negatives are derived in-batch—pretraining leverages about 400m samples of a mix of public datasets and proprietary web search data. Following pretraining models are further optimized with long training on a smaller dataset (about 1m samples) of triplets of query, positive document, and negative document derived from hard harmful mining. Mining of the negatives and data curation is crucial to retrieval accuracy. A detailed technical report will be available shortly.

Name MTEB Retrieval Score (NDCG @ 10) Parameters (Millions) Embedding Dimension
snowflake-arctic-embed-xs 50.15 22 384
snowflake-arctic-embed-s 51.98 33 384
snowflake-arctic-embed-m 54.90 110 768
snowflake-arctic-embed-m-long 54.83 137 768
snowflake-arctic-embed-l 55.98 335 1024

Aside from being great open-source models, the largest model, snowflake-arctic-embed-l, can serve as a natural replacement for closed-source embedding, as shown below.

Model Name MTEB Retrieval Score (NDCG @ 10)
snowflake-arctic-embed-l 55.98
Google-gecko-text-embedding 55.7
text-embedding-3-large 55.44
Cohere-embed-english-v3.0 55.00
bge-large-en-v1.5 54.29

This tiny model packs quite the punch. Based on the all-MiniLM-L6-v2 model with only 22m parameters and 384 dimensions, this model should meet even the strictest latency/TCO budgets. Despite its size, its retrieval accuracy is closer to that of models with 100m paramers.

Model Name MTEB Retrieval Score (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

Based on the intfloat/e5-small-unsupervised model, this small model does not trade off retrieval accuracy for its small size. With only 33m parameters and 384 dimensions, this model should easily allow scaling to large datasets.

Model Name MTEB Retrieval Score (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

Based on the intfloat/e5-base-unsupervised model, this medium model is the workhorse that provides the best retrieval performance without slowing down inference.

Model Name MTEB Retrieval Score (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

Based on the nomic-ai/nomic-embed-text-v1-unsupervised model, this long-context variant of our medium-sized model is perfect for workloads that can be constrained by the regular 512 token context of our other models. Without the use of RPE, this model supports up to 2048 tokens. With RPE, it can scale to 8192!

Model Name MTEB Retrieval Score (NDCG @ 10)
snowflake-arctic-embed-m-long 54.83
nomic-embed-text-v1.5 53.01
nomic-embed-text-v1 52.81

Based on the intfloat/e5-large-unsupervised model, this large model is a direct drop-in for closed APIs and delivers the most accurate retrieval experience.

Model Name MTEB Retrieval Score (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

Based on the google-bert/bert-base-uncased model, this variant of our medium model is trained with Matryoshka Representation Learning (MRL) loss to deliver exceptional retrieval performance even when vectors are truncated to 256 dimensions.

Model Name MTEB Retrieval Score (NDCG @ 10)
snowflake-arctic-embed-m-v1.5 55.14
snowflake-arctic-embed-m 54.91
Model Model Parameters MTEB Retrieval Score at 256 Dimensions (fraction of arctic-embed-m-v1.5)
Snowflake arctic-embed-m-v1.5 109M 54.2 (100%)
Google gecko 1200M 52.4 (97%)
OpenAI text-embedding-3-large Not Published 51.7 (96%)
Nomic nomic-embed-text-v1.5 138M 50.8 (94%)

Additionally, this model was designed to pair well with a corpus-independent scalar quantization scheme to achieve great performance even in as little as 128 bytes per vector (24x compression compared to 768 dimensional vectors stored in float32).

Model Version Dimensionality Scalar Quantization Bytes Per Vector (fraction of baseline) MTEB Retrieval Score (fraction of baseline) Vectors Per GB (improvement over baseline)
v1 768 None (float32) 3072 (100%) 54.9 (100%) 0.33M (1.0x)
v1 768 int8 768 (25%) 54.9 (100%) 1.3M (4x)
v1.5 768 int8 768 (25%) 55.1 (100%) 1.3M (4x)
v1.5 256 int8 256 (8.3%) 54.2 (99%) 3.9M (12x)
v1.5 256 int4 128 (4.2%) 53.7 (98%) 7.8M (24x)

NOTE: Good uniform scalar quantization ranges to use with this model (and which were used in the eval above), are -0.18 to +0.18 for 4bit and -0.3 to +0.3 for 8bit. For a detailed walkthrough of using integer quantization with snowflake-arctic-embed-m-v1.5, check out our example notebook.

Usage

Using Sentence Transformers

You can use the sentence-transformers package to use an snowflake-arctic-embed model. Here we show how to use our latest model, snowflake-arctic-embed-m-v1.5.

import torch
from sentence_transformers import SentenceTransformer
from torch.nn.functional import normalize

# Model constant.
MODEL_ID = "Snowflake/snowflake-arctic-embed-m-v1.5"

# Your queries and docs.
queries = ['what is snowflake?', 'Where can I get the best tacos?']
documents = ['The Data Cloud!', 'Mexico City of Course!']

# Load the model.
model = SentenceTransformer(
    MODEL_ID, model_kwargs=dict(add_pooling_layer=False),
)

# Generate text embeddings.
query_embeddings = model.encode(queries, prompt_name="query")
document_embeddings = model.encode(documents)

# Scores via dotproduct.
scores = query_embeddings @ document_embeddings.T

# Pretty-print the results.
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(f'Query: "{query}"')
    for document, score in doc_score_pairs:
        print(f'Score: {score:.4f} | Document: "{document}"')
    print()

#### OUTPUT ####
# Query: "what is snowflake?"
# Score: 0.3521 | Document: "The Data Cloud!"
# Score: 0.2358 | Document: "Mexico City of Course!"

# Query: "Where can I get the best tacos?"
# Score: 0.3884 | Document: "Mexico City of Course!"
# Score: 0.2389 | Document: "The Data Cloud!"
#

#### Variation: Truncated Embeddings ####
query_embeddings_256 = normalize(torch.from_numpy(query_embeddings)[:, :256])
doument_embeddings_256 = normalize(torch.from_numpy(document_embeddings)[:, :256])
scores_256 = query_embeddings_256 @ doument_embeddings_256.T

# Pretty-print the results.
for query, query_scores in zip(queries, scores_256):
    doc_score_pairs = sorted(zip(documents, query_scores), key=lambda x: x[1], reverse=True)
    print(f'Query: "{query}"')
    for document, score in doc_score_pairs:
        print(f'Score: {score:.4f} | Document: "{document}"')
    print()

#### OUTPUT ####
# Query: "what is snowflake?"
# Score: 0.3852 | Document: "The Data Cloud!"
# Score: 0.2721 | Document: "Mexico City of Course!"

# Query: "Where can I get the best tacos?"
# Score: 0.4337 | Document: "Mexico City of Course!"
# Score: 0.2886 | Document: "The Data Cloud!"
#

Using Huggingface transformers

You can use the transformers package to use an snowflake-arctic-embed model, too. Here we show how to use our latest model, snowflake-arctic-embed-m-v1.5. For optimal retrieval quality, use the CLS token to embed each text portion and use the query prefix below (just on the query).

import torch
from torch.nn.functional import normalize
from transformers import AutoModel, AutoTokenizer

# Model constants.
MODEL_ID = "Snowflake/snowflake-arctic-embed-m-v1.5"
QUERY_PREFIX = 'Represent this sentence for searching relevant passages: '

# Your queries and docs.
queries  = ['what is snowflake?', 'Where can I get the best tacos?']
documents = ['The Data Cloud!', 'Mexico City of Course!']

# Load the model and tokenizer.
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModel.from_pretrained(MODEL_ID, add_pooling_layer=False)
model.eval()

# Add query prefix and tokenize queries and docs.
queries_with_prefix = [f"{QUERY_PREFIX}{q}" for q in queries]
query_tokens = tokenizer(queries_with_prefix, padding=True, truncation=True, return_tensors='pt', max_length=512)
document_tokens =  tokenizer(documents, padding=True, truncation=True, return_tensors='pt', max_length=512)

# Use the model to generate text embeddings.
with torch.inference_mode():
    query_embeddings = model(**query_tokens)[0][:, 0]
    doument_embeddings = model(**document_tokens)[0][:, 0]

# Remember to normalize embeddings.
query_embeddings = normalize(query_embeddings)
doument_embeddings = normalize(doument_embeddings)

# Scores via dotproduct.
scores = query_embeddings @ document_embeddings.T

# Pretty-print the results.
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(f'Query: "{query}"')
    for document, score in doc_score_pairs:
        print(f'Score: {score:.4f} | Document: "{document}"')
    print()

#### OUTPUT ####
# Query: "what is snowflake?"
# Score: 0.3521 | Document: "The Data Cloud!"
# Score: 0.2358 | Document: "Mexico City of Course!"

# Query: "Where can I get the best tacos?"
# Score: 0.3884 | Document: "Mexico City of Course!"
# Score: 0.2389 | Document: "The Data Cloud!"
#

#### Variation: Truncated Embeddings ####
query_embeddings_256 = normalize(query_embeddings[:, :256])
doument_embeddings_256 = normalize(doument_embeddings[:, :256])
scores_256 = query_embeddings_256 @ doument_embeddings_256.T

# Pretty-print the results.
for query, query_scores in zip(queries, scores_256):
    doc_score_pairs = sorted(zip(documents, query_scores), key=lambda x: x[1], reverse=True)
    print(f'Query: "{query}"')
    for document, score in doc_score_pairs:
        print(f'Score: {score:.4f} | Document: "{document}"')
    print()

#### OUTPUT ####
# Query: "what is snowflake?"
# Score: 0.3852 | Document: "The Data Cloud!"
# Score: 0.2721 | Document: "Mexico City of Course!"

# Query: "Where can I get the best tacos?"
# Score: 0.4337 | Document: "Mexico City of Course!"
# Score: 0.2886 | Document: "The Data Cloud!"
#

Usage Note: Long Context Embedding With m-long

If you use the long context model with more than 2048 tokens, ensure that you initialize the model like below instead. This will use RoPE to allow up to 8192 tokens.

model = AutoModel.from_pretrained(
    'Snowflake/snowflake-arctic-embed-m-long',
    trust_remote_code=True,
    safe_serialization=True,
    rotary_scaling_factor=2
)

Using Transformers.js

If you haven't already, you can install the Transformers.js JavaScript library from NPM by running:

npm i @xenova/transformers

You can then compute embeddings from arctic-embed models as follows (m-long variant shown):

import { pipeline, dot } from '@xenova/transformers';

// Create feature extraction pipeline
const extractor = await pipeline('feature-extraction', 'Snowflake/snowflake-arctic-embed-m-long', {
    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.36740492125676116, 0.42407774292046635]

FAQ

TBD

Contact

Feel free to open an issue or pull request if you have any questions or suggestions about this project. You also can email Daniel Campos([email protected]).

License

Arctic is licensed under the Apache-2. The released models can be used for commercial purposes free of charge.

Acknowledgement

We want to thank the open-source community, which has provided the great building blocks upon which we could make our models. We thank our modeling engineers, Danmei Xu, Luke Merrick, Gaurav Nuti, and Daniel Campos, for making these great models possible. We thank our leadership, Himabindu Pucha, Kelvin So, Vivek Raghunathan, and Sridhar Ramaswamy, for supporting this work. We also thank the open-source community for producing the great models we could build on top of and making these releases possible. Finally, we thank the researchers who created BEIR and MTEB benchmarks. It is largely thanks to their tireless work to define what better looks like that we could improve model performance.

arctic-embed's People

Contributors

sfc-gh-lmerrick avatar sfc-gh-dcampos avatar jdanielmyers avatar

Stargazers

Andrew Kane avatar Kryspin Ziemski avatar  avatar Crystal_Alchemist avatar Eugenio Herrera-Berg avatar Andrew avatar  avatar Shahrukh Khan avatar Rohan avatar HuZuohan avatar Alexander J. Rodriguez D. avatar Akmal avatar Adoankim avatar Chih-Chiang Chang avatar Amir Soleimani avatar 刘林林 avatar Albert Yan avatar shiina avatar 爱可可-爱生活 avatar grok avatar Yotam avatar Autumn avatar  avatar Hassan ISMAIL FAWAZ avatar  avatar Ruben Martirosyan avatar Scott Soderstrom avatar Puxuan Yu avatar Vikas Chitturi - Open Source Contributor avatar Sasi Kiran Malladi avatar Chris Hart avatar Chris avatar Allen S, bookman avatar Waldek Kot avatar Shi Yu avatar Alexander Arellano avatar Joseph Y Cheng avatar Oliver Mannion avatar David Riordan avatar Brian J King avatar Simon Willison avatar Danmei Xu avatar

Watchers

Kryspin Ziemski avatar Dash Desai avatar Hassan ISMAIL FAWAZ avatar  avatar Yotam avatar Segan avatar

Forkers

ralphorimorin

arctic-embed's Issues

Can't run MTEB scripts

Looks like the import is wrong, and should just be from utils ..., importing from the local utils.py?

Traceback (most recent call last):
  File "/workspace/arctic-embed/src/mteb_retrieval.py", line 17, in <module>
    from eval.MTEB.utils import pool, logger, move_to_cuda
ModuleNotFoundError: No module named 'eval'

query_prefix in examples

Thank you for this amazing project.

I have a question about the example in the readme.

In the Using Sentence Transformers example, it does not use query_prefix but in the Huggingface example, it recommends to use query prefix. I would like to confirm if query_prefix is "required" to achieve better performance for QA retrieve tasks.

Thanks!

Recommend Projects

  • React photo React

    A declarative, efficient, and flexible JavaScript library for building user interfaces.

  • Vue.js photo Vue.js

    🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.

  • Typescript photo Typescript

    TypeScript is a superset of JavaScript that compiles to clean JavaScript output.

  • TensorFlow photo TensorFlow

    An Open Source Machine Learning Framework for Everyone

  • Django photo Django

    The Web framework for perfectionists with deadlines.

  • D3 photo D3

    Bring data to life with SVG, Canvas and HTML. 📊📈🎉

Recommend Topics

  • javascript

    JavaScript (JS) is a lightweight interpreted programming language with first-class functions.

  • web

    Some thing interesting about web. New door for the world.

  • server

    A server is a program made to process requests and deliver data to clients.

  • Machine learning

    Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.

  • Game

    Some thing interesting about game, make everyone happy.

Recommend Org

  • Facebook photo Facebook

    We are working to build community through open source technology. NB: members must have two-factor auth.

  • Microsoft photo Microsoft

    Open source projects and samples from Microsoft.

  • Google photo Google

    Google ❤️ Open Source for everyone.

  • D3 photo D3

    Data-Driven Documents codes.