Snowflake's Arctic-embed-m-long

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("Snowflake/snowflake-arctic-embed-m-long", trust_remote_code=True)

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.46484852 The Data Cloud!
0.3758855 Mexico City of Course!
Query: Where can I get the best tacos?
0.42407742 Mexico City of Course!
0.36740506 The Data Cloud!

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

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        help="Path to model",
        default="zhouhui/snowflake-arctic-embed-m-long",
    )
    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,trust_remote_code=True, add_pooling_layer=False, safe_serialization=True)
    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)

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

    # 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()

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

npm i @xenova/transformers

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]