zembed-1

from sentence_transformers import SentenceTransformer

# Initialize model
model = SentenceTransformer(
    "zeroentropy/zembed-1",
    trust_remote_code=True,
    model_kwargs={"torch_dtype": "bfloat16"},
)

# Define query and documents
query = "What is backpropagation?"
documents = [
    "Backpropagation is a fundamental algorithm for training neural networks by computing gradients.",
    "Gradient descent is used to optimize model parameters during the training process.",
    "Neural network training relies on efficient computation of derivatives through backpropagation.",
]

# Encode query and documents (uses task-specific prompts automatically)
query_embeddings = model.encode_query(query)
document_embeddings = model.encode_document(documents)
# (2560,) (3, 2560)

# Compute cosine similarities
similarities = model.similarity(query_embeddings, document_embeddings)
# tensor([[0.7525, 0.5670, 0.6835]])