jinaai

/

jina-embeddings-v5-text-small-clustering

like 6

Follow

Jina AI 1.96k

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-clustering",
    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 = [
    "We propose a novel neural network architecture for image segmentation.",
    "This paper analyzes the effects of monetary policy on inflation.",
    "Our method achieves state-of-the-art results on object detection benchmarks.",
    "We study the relationship between interest rates and housing prices.",
    "A new attention mechanism is introduced for visual recognition tasks.",
]

# Encode query and documents
embeddings = model.encode(texts)
print(embeddings.shape)
# (5, 1024)

similarity = model.similarity(embeddings, embeddings)
print(similarity)
# tensor([[1.0000, 0.2303, 0.8060, 0.2309, 0.8487],
#         [0.2303, 1.0000, 0.2552, 0.7400, 0.2543],
#         [0.8060, 0.2552, 1.0000, 0.2344, 0.8542],
#         [0.2309, 0.7400, 0.2344, 1.0000, 0.2279],
#         [0.8487, 0.2543, 0.8542, 0.2279, 1.0000]])

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

# tensor([[1.0000, 0.2303, 0.8060, 0.2309, 0.8487],
#         [0.2303, 1.0000, 0.2552, 0.7400, 0.2543],
#         [0.8060, 0.2552, 1.0000, 0.2344, 0.8542],
#         [0.2309, 0.7400, 0.2344, 1.0000, 0.2279],
#         [0.8487, 0.2543, 0.8542, 0.2279, 1.0000]])

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

tensor([[1.0000, 0.2291, 0.8063, 0.2307, 0.8483],
        [0.2291, 1.0000, 0.2551, 0.7410, 0.2542],
        [0.8063, 0.2551, 1.0000, 0.2353, 0.8544],
        [0.2307, 0.7410, 0.2353, 1.0000, 0.2286],
        [0.8483, 0.2542, 0.8544, 0.2286, 1.0000]])

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

• Tom Aarsen