Jina Code Embeddings : Recherche de code SOTA à 0,5B et 1,5B

Efficient Code Embeddings from Code Generation Models

jina-code-embeddings is a novel code embedding model suite designed to retrieve code from natural language queries, perform technical question-answering, and identify semantically similar code snippets across programming languages. It makes innovative use of an autoregressive backbone pre-trained on both text and code, generating embeddings via last-token pooling. We outline the training recipe and demonstrate state-of-the-art performance despite the relatively small size of the models, validating this approach to code embedding model construction.

arXiv.orgDaria Kryvosheieva

jina-code-embeddings-1.5b - Search Foundation Models

Efficient code embeddings from code generation models

Search Foundation ModelsJina AI

jinaai/jina-code-embeddings-1.5b · Hugging Face

We’re on a journey to advance and democratize artificial intelligence through open source and open science.

Aujourd'hui, nous publions jina-code-embeddings, une nouvelle suite de modèles d'intégration de code en deux tailles (0,5 milliard et 1,5 milliard de paramètres), ainsi que des quantifications GGUF pour les deux. Construits sur des LLM de génération de code autorégressifs, ces modèles atteignent des performances de récupération de pointe malgré leur taille compacte. Ils prennent en charge plus de 15 langages de programmation, notamment Python, JavaScript, Java, C++, C#, Go, Rust, TypeScript, SQL, MATLAB, R, Swift, Kotlin, HTML/CSS, PHP, Ruby, Scala, Perl et Shell.

jina-code-embeddings atteint une performance moyenne de 78,41 % (0,5 milliard) et 79,04 % (1,5 milliard) sur 25 benchmarks de récupération de code. Le modèle 0,5B surpasse Qwen3-Embedding-0.6B de 5 points de pourcentage bien qu'il soit 20 % plus petit, tandis que la variante 1,5B équivaut à voyage-code-3 (79,23 %) et dépasse gemini-embedding-001 (77,38 %) — deux modèles propriétaires avec des architectures non divulguées.

*Closed-source models with undisclosed architecture

Les deux modèles ont été entraînés avec cinq préfixes d'instructions spécifiques à la tâche pour différents scénarios de récupération, chacun prenant en charge les rôles de requête et de document pour la récupération asymétrique. Par exemple, vous pouvez utiliser nl2code_query pour intégrer des requêtes et nl2code_document pour intégrer des documents.

tagTraining Recipe

Nous utilisons des modèles de génération de code pré-entraînés comme base d'intégration. Construits sur Qwen2.5-Coder-0.5B et 1.5B, nos modèles présentent :

Les modèles d'intégration de code traditionnels sont confrontés à un goulot d'étranglement fondamental : il n'y a tout simplement pas assez de paires commentaire-code de haute qualité pour l'apprentissage supervisé. En commençant par Qwen2.5-Coder pré-entraîné sur 5,5 milliards de tokens couvrant plus de 92 langages de programmation, nous héritons d'une compréhension sémantique approfondie des constructions de programmation, de la reconnaissance de formes inter-langues et d'une connaissance intégrée de la syntaxe et des idiomes. Le réglage fin contrastif adapte ensuite ces connaissances aux tâches de récupération avec un minimum de données alignées, contournant ainsi la pénurie de données qui limite les modèles à encodeur uniquement.

Pour les tâches sous-représentées telles que les traductions de code inter-framework, nous avons généré des données synthétiques à l'aide de LLM, chaque exemple synthétique étant validé manuellement pour la qualité. Nos données d'entraînement ont combiné les divisions d'entraînement des tâches de code MTEB existantes avec des ensembles de données publics adaptés, notamment CommitPackFT, SWE-Bench, Spider, MBPP et CodeSearchNet.

Contrairement à jina-embeddings-v3 et v4, nous n'avons pas utilisé LoRA et sommes passés directement à un post-entraînement complet. Pour les petits modèles comme les nôtres (494M et 1,54B de paramètres), l'efficacité des paramètres de LoRA devient moins intéressante : la surcharge de l'adaptateur peut en fait nuire aux performances lorsque vous avez une capacité limitée. Nous avions besoin que chaque paramètre fonctionne sur la tâche d'intégration. Même pour les scénarios multi-tâches, les préfixes d'instructions spécifiques à la tâche se sont avérés plus propres que plusieurs adaptateurs LoRA. Au lieu de modifier les configurations de pondération, nous ajoutons simplement différentes instructions, ce qui est beaucoup plus simple et plus aligné sur la façon dont les LLM traitent naturellement les informations conditionnelles.

L'entraînement a été remarquablement efficace : les deux modèles ont été entraînés à l'aide de l'apprentissage contrastif avec perte InfoNCE sur 4 GPU A100 80 Go, ce qui a pris seulement 8,3 heures pour le modèle 0,5B et 12 heures pour la variante 1,5B.

Enfin, nous avons comparé différentes stratégies de mise en commun. La mise en commun du dernier token a atteint une moyenne globale de 78,41 %, surpassant systématiquement la mise en commun moyenne (77,20 %) et la mise en commun de l'attention latente (78,27 %) dans toutes les catégories de benchmarks. Cet avantage de 1,2 point de pourcentage nous a amenés à rompre avec la tradition de mise en commun moyenne que nous avions établie dans jina-embeddings-v2, v3 et v4. À mesure que davantage de modèles de récupération sont construits sur des LLM à décodeur uniquement, la mise en commun du dernier token devient le choix naturel : la mise en commun moyenne ne s'aligne tout simplement pas bien avec les mécanismes d'attention unidirectionnelle. Bien que la mise en commun moyenne puisse fonctionner et s'entraîne souvent plus facilement dans les premières étapes (probablement en raison de son paysage d'optimisation convexe), nos expériences montrent systématiquement qu'elle se stabilise en dessous du seuil de performance que la mise en commun du dernier token atteint.

tagGetting Started

Les deux modèles fonctionnent de manière transparente via notre API Search Foundation et avec des frameworks populaires, notamment sentence-transformers, transformers et llama.cpp.

tagVia API

curl http://api.jina.ai/v1/embeddings \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $JINA_API_KEY" \
  -d @- <<EOFEOF
  {
    "model": "jina-code-embeddings-1.5b",
    "input": ["print hello world in python"],
    "task": "nl2code.passage"
  }
EOFEOF
Copier

tagVia sentence-transformers

from sentence_transformers import SentenceTransformer

# Load the model (choose 0.5b or 1.5b)
model = SentenceTransformer(
    "jinaai/jina-code-embeddings-1.5b",
    model_kwargs={"torch_dtype": "bfloat16"},
    tokenizer_kwargs={"padding_side": "left"}
)

# Natural language to code
queries = ["print hello world in python", "initialize array of 5 zeros in c++"]
documents = ["print('Hello World!')", "int arr[5] = {0, 0, 0, 0, 0};"]

# Generate embeddings with task-specific prefixes
query_embeddings = model.encode(queries, prompt_name="nl2code_query")
document_embeddings = model.encode(documents, prompt_name="nl2code_document")

# Compute similarity
similarity = model.similarity(query_embeddings, document_embeddings)
Copier

tagVia transformers

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

def last_token_pool(last_hidden_states, attention_mask):
    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
    if left_padding:
        return last_hidden_states[:, -1]
    else:
        sequence_lengths = attention_mask.sum(dim=1) - 1
        batch_size = last_hidden_states.shape[0]
        return last_hidden_states[torch.arange(batch_size), sequence_lengths]

tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-code-embeddings-1.5b')
model = AutoModel.from_pretrained('jinaai/jina-code-embeddings-1.5b')

# Apply task-specific prefix
query = "Find the most relevant code snippet given the following query:\nprint hello world"
code = "Candidate code snippet:\nprint('Hello World!')"

# Tokenize and embed
batch_dict = tokenizer([query, code], padding=True, truncation=True, return_tensors="pt")
outputs = model(**batch_dict)
embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
Copier

tagTroncature des vecteurs modèles Matriochka

Les deux modèles ont été entraînés avec l'apprentissage de la représentation Matriochka pour les dimensions [64, 128, 256, 512, 896], ce qui vous permet de tronquer les vecteurs modèles sans recalcul :

# Vecteurs modèles complets : 896d (0.5B) ou 1536d (1.5B)
full_embedding = model.encode(text)

# Tronquer vers des dimensions plus petites pour plus d'efficacité
small_embedding = full_embedding[:256]  # Fonctionne pour les deux modèles
tiny_embedding = full_embedding[:128]   # 0.5B prend en charge jusqu'à 64d
Copier

Cette flexibilité permet de faire des compromis entre performances et efficacité en fonction de vos besoins.

tagConclusion

jina-code-embeddings démontre que des vecteurs modèles de code efficaces ne nécessitent pas une échelle massive. En nous appuyant sur des modèles de génération de code et en appliquant un fine-tuning ciblé, nous obtenons des performances de pointe avec des modèles de moins de 1,5 milliard de paramètres.

Les excellents résultats obtenus avec des modèles aussi compacts (0,5B/1,5B) valident notre thèse : la bonne fondation compte plus que le nombre de paramètres. Les modèles de génération comprennent la sémantique du code - cette compréhension est directement transférée aux tâches de représentation.

Cela s'aligne sur notre vision plus large chez Jina AI : des architectures unifiées où les vecteurs modèles et la génération émergent de la même fondation, repoussant les limites de ce qui est possible avec les modèles de fondation de recherche.