mixedbread-ai / mxbai-embed-xsmall-v1

pip install -U angle-emb

from angle_emb import AnglE
from angle_emb.utils import cosine_similarity

# 1. Specify preferred dimensions
dimensions = 384

# 2. Load model and set pooling strategy to avg
model = AnglE.from_pretrained(
    "mixedbread-ai/mxbai-embed-xsmall-v1",
    pooling_strategy='avg').cuda()

query = 'A man is eating a piece of bread'

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

# 3. Encode
embeddings = model.encode(docs, embedding_size=dimensions)

for doc, emb in zip(docs[1:], embeddings[1:]):
    print(f'{query} ||| {doc}', cosine_similarity(embeddings[0], emb))

python -m pip install -U sentence-transformers

from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim

# 1. Specify preferred dimensions
dimensions = 384

# 2. Load model
model = SentenceTransformer("mixedbread-ai/mxbai-embed-xsmall-v1", truncate_dim=dimensions)

query = 'A man is eating a piece of bread'

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


# 3. Encode
embeddings = model.encode(docs)

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

pip install -U transformers

from typing import Dict

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

def pooling(outputs: torch.Tensor, inputs: Dict) -> np.ndarray:
    outputs = torch.sum(
      outputs * inputs["attention_mask"][:, :, None], dim=1) / torch.sum(inputs["attention_mask"])
    return outputs.detach().cpu().numpy()

# 1. Load model
model_id = 'mixedbread-ai/mxbai-embed-xsmall-v1'
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModel.from_pretrained(model_id).cuda()

query = 'A man is eating a piece of bread'

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

# 2. Encode
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)

# 3. Compute similarity scores
similarities = cos_sim(embeddings[0], embeddings[1:])
print('similarities:', similarities)

python -m pip install batched

import uvicorn
import batched
from fastapi import FastAPI
from fastapi.responses import ORJSONResponse
from sentence_transformers import SentenceTransformer
from pydantic import BaseModel
 
app = FastAPI()
 
model = SentenceTransformer('mixedbread-ai/mxbai-embed-xsmall-v1')
model.encode = batched.aio.dynamically(model.encode)
 
class EmbeddingsRequest(BaseModel):
    input: str | list[str]
 
@app.post("/embeddings")
async def embeddings(request: EmbeddingsRequest):
    return ORJSONResponse({"embeddings": await model.encode(request.input)})
 
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)