mixedbread-ai/mxbai-embed-large-v1

这里我们提供了几种生成句子嵌入的方法。请注意，如果您想将其用于检索任务，需要为查询提供提示词 Represent this sentence for searching relevant passages:。除此之外不需要任何提示词。我们的模型还支持 Matryoshka 表示学习和二进制量化。

快速开始

这里我们提供了几种生成句子嵌入的方法。请注意，如果您想将其用于检索任务，需要为查询提供提示词 Represent this sentence for searching relevant passages: 。除此之外不需要任何提示词。

sentence-transformers

python -m pip install -U sentence-transformers

from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim
from sentence_transformers.quantization import quantize_embeddings

# 1. 指定首选维度
dimensions = 512

# 2. 加载模型
model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1", truncate_dim=dimensions)

# 用于查询检索任务的提示词：
# query_prompt = 'Represent this sentence for searching relevant passages: '

query = "A man is eating a piece of bread"
docs = [
    "A man is eating food.",
    "A man is eating pasta.",
    "The girl is carrying a baby.",
    "A man is riding a horse.",
]

# 2. 编码
query_embedding = model.encode(query, prompt_name="query")
# 等效替代方案：
# query_embedding = model.encode(query_prompt + query)
# query_embedding = model.encode(query, prompt=query_prompt)

docs_embeddings = model.encode(docs)

# 可选：对嵌入进行量化
binary_query_embedding = quantize_embeddings(query_embedding, precision="ubinary")
binary_docs_embeddings = quantize_embeddings(docs_embeddings, precision="ubinary")

similarities = cos_sim(query_embedding, docs_embeddings)
print('similarities:', similarities)

Transformers

from typing import Dict

import torch
import numpy as np
from transformers import AutoModel, AutoTokenizer
from sentence_transformers.util import cos_sim

# 对于检索，您需要传递此提示词。请在我们的博客文章中了解更多。
def transform_query(query: str) -> str:
    """ 对于检索，为查询添加提示词（文档不需要）。
    """
    return f'Represent this sentence for searching relevant passages: {query}'

# 该模型与 cls 池化（默认）配合使用效果很好，但也支持 mean 池化。
def pooling(outputs: torch.Tensor, inputs: Dict,  strategy: str = 'cls') -> np.ndarray:
    if strategy == 'cls':
        outputs = outputs[:, 0]
    elif strategy == 'mean':
        outputs = torch.sum(
            outputs * inputs["attention_mask"][:, :, None], dim=1) / torch.sum(inputs["attention_mask"], dim=1, keepdim=True)
    else:
        raise NotImplementedError
    return outputs.detach().cpu().numpy()

# 1. 加载模型
model_id = 'mixedbread-ai/mxbai-embed-large-v1'
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModel.from_pretrained(model_id).cuda()


docs = [
    transform_query('A man is eating a piece of bread'),
    "A man is eating food.",
    "A man is eating pasta.",
    "The girl is carrying a baby.",
    "A man is riding a horse.",
]

# 2. 编码
inputs = tokenizer(docs, padding=True, return_tensors='pt')
for k, v in inputs.items():
    inputs[k] = v.cuda()
outputs = model(**inputs).last_hidden_state
embeddings = pooling(outputs, inputs, 'cls')

similarities = cos_sim(embeddings[0], embeddings[1:])
print('similarities:', similarities)

Transformers.js

如果您还没有安装，可以从 NPM 安装 Transformers.js JavaScript 库：

npm i @huggingface/transformers

然后您可以像这样使用模型计算嵌入：

import { pipeline, cos_sim } from "@huggingface/transformers";

// 创建特征提取管道
const extractor = await pipeline("feature-extraction", "mixedbread-ai/mxbai-embed-large-v1", {
    dtype: "fp32", // 选项: "fp32", "fp16", "q8"
});

// 生成句子嵌入
const docs = [
    "Represent this sentence for searching relevant passages: A man is eating a piece of bread",
    "A man is eating food.",
    "A man is eating pasta.",
    "The girl is carrying a baby.",
    "A man is riding a horse.",
]
const output = await extractor(docs, { pooling: "cls" });

// 计算相似度分数
const [source_embeddings, ...document_embeddings ] = output.tolist();
const similarities = document_embeddings.map(x => cos_sim(source_embeddings, x));
console.log(similarities); // [0.7919578577247139, 0.6369278664248345, 0.16512018371357193, 0.3620778366720027]

使用 API

您可以通过我们的 API 使用该模型，如下所示：

from mixedbread_ai.client import MixedbreadAI, EncodingFormat
from sklearn.metrics.pairwise import cosine_similarity
import os

mxbai = MixedbreadAI(api_key="{MIXEDBREAD_API_KEY}")

english_sentences = [
    'What is the capital of Australia?',
    'Canberra is the capital of Australia.'
] 

res = mxbai.embeddings(
     input=english_sentences,
     model="mixedbread-ai/mxbai-embed-large-v1",
     normalized=True,
     encoding_format=[EncodingFormat.FLOAT, EncodingFormat.UBINARY, EncodingFormat.INT_8],
     dimensions=512
)

encoded_embeddings = res.data[0].embedding
print(res.dimensions, encoded_embeddings.ubinary, encoded_embeddings.float_, encoded_embeddings.int_8)

API 原生支持 int8 和二进制量化！查看文档了解更多详细信息。

Infinity

docker run --gpus all -v $PWD/data:/app/.cache -p "7997":"7997" \
michaelf34/infinity:0.0.68 \
v2 --model-id mixedbread-ai/mxbai-embed-large-v1 --revision "main" --dtype float16 --engine torch --port 7997

评估

截至 2024 年 3 月，我们的模型在 MTEB 上实现了 Bert-large 尺寸模型的 SOTA 性能。它优于 OpenAI 的 text-embedding-3-large 等商业模型，并与比其大 20 倍的模型（如 echo-mistral-7b）的性能相当。我们的模型在训练时没有使用 MTEB 数据，这表明我们的模型在多个领域、任务和文本长度上都能很好地泛化。我们知道该模型存在一些局限性，这些将在 v2 版本中修复。

更多详细信息请参阅我们的博客文章。

Matryoshka 和二进制量化

以其常用形式（浮点数组）使用的嵌入在大规模使用时占用大量内存。解决这个问题的两种方法是 Matryoshka 表示学习（MRL）和（二进制）量化。MRL 减少嵌入的维度数量，而二进制量化将每个维度的值从 float32 转换为较低精度（int8 甚至二进制）。该模型同时支持这两种方法！

您还可以更进一步，将 MRL 和量化结合使用。二进制量化和 MRL 的结合可以显著减少嵌入的内存使用。这在使用向量数据库时可以大幅降低成本。您可以在我们的博客文章中阅读更多关于该技术及其优势的信息。

社区

请加入我们的 Discord 社区，分享您的反馈和想法！我们在这里提供帮助，也很高兴与您交流。

许可证

Apache 2.0

引用

@online{emb2024mxbai,
  title={Open Source Strikes Bread - New Fluffy Embeddings Model},
  author={Sean Lee and Aamir Shakir and Darius Koenig and Julius Lipp},
  year={2024},
  url={https://www.mixedbread.ai/blog/mxbai-embed-large-v1},
}

@article{li2023angle,
  title={AnglE-optimized Text Embeddings},
  author={Li, Xianming and Li, Jing},
  journal={arXiv preprint arXiv:2309.12871},
  year={2023}
}