jinaai

/

jina-embeddings-v5-text-small-classification

like 4

Follow

Jina AI 2.02k

Hello!

Pull Request overview

• Add Sentence Transformers compatibility
• Add ST snippet to README

Details

I noticed that because this model doesn't use the adapter work like the main model, this model can actually be implemented without any trust_remote_code in Sentence Transformers. That might be pretty valuable for users. I just loaded the model with SentenceTransformer, saved it with save_pretrained("tmp"), updated the default pooling config with lasttoken pooling instead, and added a Normalize to the modules.json.

Try it out

from sentence_transformers import SentenceTransformer
import torch

model = SentenceTransformer(
    "jinaai/jina-embeddings-v5-text-small-classification",
    model_kwargs={"dtype": torch.bfloat16},  # Recommended for GPUs
    config_kwargs={"_attn_implementation": "flash_attention_2"},  # Recommended but optional
    revision="refs/pr/1",
)
# Optional: set truncate_dim in encode() to control embedding size

texts = [
    "My order hasn't arrived yet and it's been two weeks.",
    "How do I reset my password?",
    "I'd like a refund for my recent purchase.",
    "Your product exceeded my expectations. Great job!",
]

# Encode texts
embeddings = model.encode(texts)
print(embeddings.shape)
# (4, 1024)

similarity = model.similarity(embeddings, embeddings)
print(similarity)
# tensor([[1.0000, 0.6262, 0.7605, 0.6213],
#         [0.6262, 1.0000, 0.6365, 0.5854],
#         [0.7605, 0.6365, 1.0000, 0.6423],
#         [0.6213, 0.5854, 0.6423, 1.0000]])

Note the revision argument. I got this result with the script:

# tensor([[1.0000, 0.6262, 0.7605, 0.6213],
#         [0.6262, 1.0000, 0.6365, 0.5854],
#         [0.7605, 0.6365, 1.0000, 0.6423],
#         [0.6213, 0.5854, 0.6423, 1.0000]])

Compared to this output when using the original model with the classification task:

tensor([[1.0000, 0.6254, 0.7624, 0.6206],
        [0.6254, 1.0000, 0.6369, 0.5841],
        [0.7624, 0.6369, 1.0000, 0.6413],
        [0.6206, 0.5841, 0.6413, 1.0000]])

I believe the small difference is likely due to the adapter loading vs already being merged into the model.

• Tom Aarsen