jina-embeddings-v3-hf

import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("jinaai/jina-embeddings-v3-hf", dtype=torch.float32)

# Load and activate the retrieval_query LoRA adapter
# Available tasks: retrieval_query, retrieval_passage, separation, classification, text_matching
task = "retrieval_query"
model.load_adapter("jinaai/jina-embeddings-v3-hf", adapter_name=task, adapter_kwargs={"subfolder": task})
model.set_adapter(task)
model.eval()

tokenizer = AutoTokenizer.from_pretrained("jinaai/jina-embeddings-v3-hf")
texts = ["How is the weather today?", "What is the current weather like today?"]
encoded = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")

with torch.no_grad():
    outputs = model(**encoded)

# Mean pooling
attention_mask = encoded["attention_mask"].unsqueeze(-1)
embeddings = (outputs.last_hidden_state * attention_mask).sum(dim=1) / attention_mask.sum(dim=1)
embeddings = F.normalize(embeddings, p=2, dim=1)

print(embeddings.shape)
# Cosine similarity
print(embeddings @ embeddings.T)

@misc{sturua2024jinaembeddingsv3multilingualembeddingstask,
      title={jina-embeddings-v3: Multilingual Embeddings With Task LoRA}, 
      author={Saba Sturua and Isabelle Mohr and Mohammad Kalim Akram and Michael Günther and Bo Wang and Markus Krimmel and Feng Wang and Georgios Mastrapas and Andreas Koukounas and Andreas Koukounas and Nan Wang and Han Xiao},
      year={2024},
      eprint={2409.10173},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2409.10173}, 
}