Views
No views yet
Represent this sentence for searching relevant passages: for query if you want to use it for retrieval. Besides that you don't need any prompt. Our model also supports Matryoshka Representation Learning and binary quantization.Represent this sentence for searching relevant passages: for query if you want to use it for retrieval. Besides that you don't need any prompt.python -m pip install -U sentence-transformers1from sentence_transformers import SentenceTransformer
2from sentence_transformers.util import cos_sim
3from sentence_transformers.quantization import quantize_embeddings
4
5# 1. Specify preffered dimensions
6dimensions = 512
7
8# 2. load model
9model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1", truncate_dim=dimensions)
10
11# The prompt used for query retrieval tasks:
12# query_prompt = 'Represent this sentence for searching relevant passages: '
13
14query = "A man is eating a piece of bread"
15docs = [
16 "A man is eating food.",
17 "A man is eating pasta.",
18 "The girl is carrying a baby.",
19 "A man is riding a horse.",
20]
21
22# 2. Encode
23query_embedding = model.encode(query, prompt_name="query")
24# Equivalent Alternatives:
25# query_embedding = model.encode(query_prompt + query)
26# query_embedding = model.encode(query, prompt=query_prompt)
27
28docs_embeddings = model.encode(docs)
29
30# Optional: Quantize the embeddings
31binary_query_embedding = quantize_embeddings(query_embedding, precision="ubinary")
32binary_docs_embeddings = quantize_embeddings(docs_embeddings, precision="ubinary")
33
34similarities = cos_sim(query_embedding, docs_embeddings)
35print('similarities:', similarities)1from typing import Dict
2
3import torch
4import numpy as np
5from transformers import AutoModel, AutoTokenizer
6from sentence_transformers.util import cos_sim
7
8# For retrieval you need to pass this prompt. Please find our more in our blog post.
9def transform_query(query: str) -> str:
10 """ For retrieval, add the prompt for query (not for documents).
11 """
12 return f'Represent this sentence for searching relevant passages: {query}'
13
14# The model works really well with cls pooling (default) but also with mean pooling.
15def pooling(outputs: torch.Tensor, inputs: Dict, strategy: str = 'cls') -> np.ndarray:
16 if strategy == 'cls':
17 outputs = outputs[:, 0]
18 elif strategy == 'mean':
19 outputs = torch.sum(
20 outputs * inputs["attention_mask"][:, :, None], dim=1) / torch.sum(inputs["attention_mask"], dim=1, keepdim=True)
21 else:
22 raise NotImplementedError
23 return outputs.detach().cpu().numpy()
24
25# 1. load model
26model_id = 'mixedbread-ai/mxbai-embed-large-v1'
27tokenizer = AutoTokenizer.from_pretrained(model_id)
28model = AutoModel.from_pretrained(model_id).cuda()
29
30
31docs = [
32 transform_query('A man is eating a piece of bread'),
33 "A man is eating food.",
34 "A man is eating pasta.",
35 "The girl is carrying a baby.",
36 "A man is riding a horse.",
37]
38
39# 2. encode
40inputs = tokenizer(docs, padding=True, return_tensors='pt')
41for k, v in inputs.items():
42 inputs[k] = v.cuda()
43outputs = model(**inputs).last_hidden_state
44embeddings = pooling(outputs, inputs, 'cls')
45
46similarities = cos_sim(embeddings[0], embeddings[1:])
47print('similarities:', similarities)npm i @huggingface/transformers1import { pipeline, cos_sim } from "@huggingface/transformers";
2
3// Create a feature extraction pipeline
4const extractor = await pipeline("feature-extraction", "mixedbread-ai/mxbai-embed-large-v1", {
5 dtype: "fp32", // Options: "fp32", "fp16", "q8"
6});
7
8// Generate sentence embeddings
9const docs = [
10 "Represent this sentence for searching relevant passages: A man is eating a piece of bread",
11 "A man is eating food.",
12 "A man is eating pasta.",
13 "The girl is carrying a baby.",
14 "A man is riding a horse.",
15]
16const output = await extractor(docs, { pooling: "cls" });
17
18// Compute similarity scores
19const [source_embeddings, ...document_embeddings ] = output.tolist();
20const similarities = document_embeddings.map(x => cos_sim(source_embeddings, x));
21console.log(similarities); // [0.7919578577247139, 0.6369278664248345, 0.16512018371357193, 0.3620778366720027]1from mixedbread_ai.client import MixedbreadAI, EncodingFormat
2from sklearn.metrics.pairwise import cosine_similarity
3import os
4
5mxbai = MixedbreadAI(api_key="{MIXEDBREAD_API_KEY}")
6
7english_sentences = [
8 'What is the capital of Australia?',
9 'Canberra is the capital of Australia.'
10]
11
12res = mxbai.embeddings(
13 input=english_sentences,
14 model="mixedbread-ai/mxbai-embed-large-v1",
15 normalized=True,
16 encoding_format=[EncodingFormat.FLOAT, EncodingFormat.UBINARY, EncodingFormat.INT_8],
17 dimensions=512
18)
19
20encoded_embeddings = res.data[0].embedding
21print(res.dimensions, encoded_embeddings.ubinary, encoded_embeddings.float_, encoded_embeddings.int_8)1docker run --gpus all -v $PWD/data:/app/.cache -p "7997":"7997" \
2michaelf34/infinity:0.0.68 \
3v2 --model-id mixedbread-ai/mxbai-embed-large-v1 --revision "main" --dtype float16 --engine torch --port 7997| Model | Avg (56 datasets) | Classification (12 datasets) | Clustering (11 datasets) | PairClassification (3 datasets) | Reranking (4 datasets) | Retrieval (15 datasets) | STS (10 datasets) | Summarization (1 dataset) |
|---|---|---|---|---|---|---|---|---|
| mxbai-embed-large-v1 | 64.68 | 75.64 | 46.71 | 87.2 | 60.11 | 54.39 | 85.00 | 32.71 |
| bge-large-en-v1.5 | 64.23 | 75.97 | 46.08 | 87.12 | 60.03 | 54.29 | 83.11 | 31.61 |
| mxbai-embed-2d-large-v1 | 63.25 | 74.14 | 46.07 | 85.89 | 58.94 | 51.42 | 84.9 | 31.55 |
| nomic-embed-text-v1 | 62.39 | 74.12 | 43.91 | 85.15 | 55.69 | 52.81 | 82.06 | 30.08 |
| jina-embeddings-v2-base-en | 60.38 | 73.45 | 41.73 | 85.38 | 56.98 | 47.87 | 80.7 | 31.6 |
| Proprietary Models | ||||||||
| OpenAI text-embedding-3-large | 64.58 | 75.45 | 49.01 | 85.72 | 59.16 | 55.44 | 81.73 | 29.92 |
| Cohere embed-english-v3.0 | 64.47 | 76.49 | 47.43 | 85.84 | 58.01 | 55.00 | 82.62 | 30.18 |
| OpenAI text-embedding-ada-002 | 60.99 | 70.93 | 45.90 | 84.89 | 56.32 | 49.25 | 80.97 | 30.80 |
1@online{emb2024mxbai,
2 title={Open Source Strikes Bread - New Fluffy Embeddings Model},
3 author={Sean Lee and Aamir Shakir and Darius Koenig and Julius Lipp},
4 year={2024},
5 url={https://www.mixedbread.ai/blog/mxbai-embed-large-v1},
6}
7
8@article{li2023angle,
9 title={AnglE-optimized Text Embeddings},
10 author={Li, Xianming and Li, Jing},
11 journal={arXiv preprint arXiv:2309.12871},
12 year={2023}
13}