Snowflake's Arctic-embed-s

from sentence_transformers import SentenceTransformer

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

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.533809 The Data Cloud!
0.49207097 Mexico City of Course!
Query: Where can I get the best tacos?
0.56592476 Mexico City of Course!
0.48255116 The Data Cloud!

from openmind import AutoTokenizer, AutoModel, is_torch_npu_available
from openmind_hub import snapshot_download
import torch.nn.functional as F
from torch import Tensor
import openmind
import torch
import argparse

# Mean Pooling - Take attention mask into account for correct averaging
def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0]  # First element of model_output contains all token embeddings
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        help="Path to model",
        default="jeffding/snowflake-arctic-embed-s-openmind",
    )
    args = parser.parse_args()
    return args

def main():
    args = parse_args()
    model_path = args.model_name_or_path

    if is_torch_npu_available():
        device = "npu:0"
    else:
        device = "cpu"
        
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModel.from_pretrained(model_path, add_pooling_layer=False).to(device)
    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).to(device)

    documents = ['The Data Cloud!', 'Mexico City of Course!']
    document_tokens =  tokenizer(documents, padding=True, truncation=True, return_tensors='pt', max_length=512).to(device)

    # 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)
    
if __name__ == "__main__":
    main()

npm i @xenova/transformers

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

// Create feature extraction pipeline
const extractor = await pipeline('feature-extraction', 'Snowflake/snowflake-arctic-embed-s', {
    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.48255123876493394, 0.5659250100112143]