Jina Code Embeddings: SOTA Code Retrieval bei 0.5B und 1.5B

Efficient Code Embeddings from Code Generation Models

jina-code-embeddings ist eine neuartige Suite von Code-Vektor-Modellen, die entwickelt wurde, um Code aus natürlichsprachlichen Abfragen abzurufen, technische Fragen zu beantworten und semantisch ähnliche Code-Snippets über Programmiersprachen hinweg zu identifizieren. Sie nutzt auf innovative Weise ein autoregressives Backbone, das sowohl auf Text als auch auf Code vortrainiert wurde, und generiert Vektor-Modelle über Last-Token-Pooling. Wir skizzieren das Trainingsrezept und demonstrieren modernste Leistung trotz der relativ geringen Größe der Modelle, was diesen Ansatz zur Konstruktion von Code-Vektor-Modellen validiert.

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.

Heute veröffentlichen wir jina-code-embeddings, eine neue Suite von Code-Vektor-Modellen in zwei Größen – 0,5B und 1,5B Parametern – zusammen mit GGUF-Quantisierungen für beide. Diese Modelle, die auf autoregressiven Code-Generierungs-LLMs basieren, erzielen trotz ihrer kompakten Größe eine hochmoderne Abrufleistung. Sie unterstützen über 15 Programmiersprachen, darunter Python, JavaScript, Java, C++, C#, Go, Rust, TypeScript, SQL, MATLAB, R, Swift, Kotlin, HTML/CSS, PHP, Ruby, Scala, Perl und Shell.

jina-code-embeddings erreicht eine durchschnittliche Leistung von 78,41 % (0,5B) und 79,04 % (1,5B) über 25 Code-Abruf-Benchmarks hinweg. Das 0,5B-Modell übertrifft Qwen3-Embedding-0.6B um 5 Prozentpunkte, obwohl es 20 % kleiner ist, während die 1,5B-Variante mit voyage-code-3 (79,23 %) übereinstimmt und gemini-embedding-001 (77,38 %) übertrifft – beides proprietäre Modelle mit nicht offengelegten Architekturen.

*Closed-Source-Modelle mit nicht offengelegter Architektur

Beide Modelle wurden mit fünf aufgabenspezifischen Instruktionspräfixen für verschiedene Abrufszenarien trainiert, die jeweils sowohl Abfrage- als auch Dokumentenrollen für den asymmetrischen Abruf unterstützen. Sie können beispielsweise nl2code_query verwenden, um Abfragen einzubetten, und nl2code_document, um Dokumente einzubetten.

tagTrainingsrezept

Wir verwenden vortrainierte Code-Generierungsmodelle als Embedding-Backbones. Unsere Modelle basieren auf Qwen2.5-Coder-0.5B und 1.5B und zeichnen sich durch Folgendes aus:

Traditionelle Code-Vektor-Modelle stehen vor einem grundlegenden Engpass: Es gibt einfach nicht genügend hochwertige Kommentar-Code-Paare für überwachtes Training. Indem wir mit Qwen2.5-Coder beginnen, das auf 5,5 Billionen Tokens vortrainiert wurde und über 92 Programmiersprachen umfasst, erben wir ein tiefes semantisches Verständnis von Programmierkonstrukten, sprachübergreifende Mustererkennung und eingebautes Wissen über Syntax und Idiome. Das kontrastive Feintuning passt dieses Wissen dann mit minimalen ausgerichteten Daten an Abrufaufgaben an und umgeht so die Datenknappheit, die Encoder-Only-Modelle einschränkt.

Für unterrepräsentierte Aufgaben wie Code-Übersetzungen über Frameworks hinweg haben wir synthetische Daten mithilfe von LLMs generiert, wobei jedes synthetische Beispiel manuell auf Qualität geprüft wurde. Unsere Trainingsdaten kombinierten bestehende MTEB-Code-Aufgaben-Trainingsaufteilungen mit angepassten öffentlichen Datensätzen, darunter CommitPackFT, SWE-Bench, Spider, MBPP und CodeSearchNet.

Im Gegensatz zu jina-embeddings-v3 und v4 haben wir LoRA nicht verwendet und sind direkt zum vollständigen Post-Training übergegangen. Für kleine Modelle wie unsere (494M und 1,54B Parameter) wird die Parametereffizienz von LoRA weniger überzeugend – der Adapter-Overhead kann die Leistung sogar beeinträchtigen, wenn man nur begrenzte Kapazität hat. Wir brauchten jeden Parameter, der an der Einbettungsaufgabe arbeitet. Selbst für Multi-Task-Szenarien erwiesen sich aufgabenspezifische Instruktionspräfixe als sauberer als mehrere LoRA-Adapter. Anstatt Gewichtskonfigurationen zu wechseln, fügen wir einfach verschiedene Anweisungen voran – viel schlanker und stärker darauf ausgerichtet, wie LLMs bedingte Informationen auf natürliche Weise verarbeiten.

Das Training war bemerkenswert effizient: Beide Modelle wurden mithilfe von kontrastivem Lernen mit InfoNCE-Verlust auf 4x A100 80GB GPUs trainiert und in nur 8,3 Stunden für das 0,5B-Modell und 12 Stunden für die 1,5B-Variante abgeschlossen.

Schließlich haben wir verschiedene Pooling-Strategien verglichen. Last-Token-Pooling erzielte einen Gesamtdurchschnitt von 78,41 % und übertraf damit konsistent Mean-Pooling (77,20 %) und Latent-Attention-Pooling (78,27 %) in allen Benchmark-Kategorien. Dieser Vorteil von 1,2 Prozentpunkten veranlasste uns, mit der Mean-Pooling-Tradition zu brechen, die wir in jina-embeddings-v2, v3 und v4 etabliert hatten. Da immer mehr Abrufmodelle auf Decoder-Only-LLMs aufbauen, wird Last-Token-Pooling zur natürlichen Wahl – Mean-Pooling passt einfach nicht gut zu unidirektionalen Aufmerksamkeitsmechanismen. Während Mean-Pooling funktionieren kann und in frühen Schritten oft leichter zu trainieren ist (wahrscheinlich aufgrund seiner konvexen Optimierungslandschaft), zeigen unsere Experimente konsistent, dass es unterhalb der Leistungsgrenze stagniert, die Last-Token-Pooling erreicht.

tagErste Schritte

Beide Modelle funktionieren nahtlos über unsere Search Foundation API und mit beliebten Frameworks wie sentence-transformers, transformers und 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
Kopieren

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)
Kopieren

tagVia transformers

```python
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'])

tagMatrjoschka-Vektormodell-Cut-Off

Beide Modelle wurden mit Matrjoschka Representation Learning für die Dimensionen [64, 128, 256, 512, 896] trainiert, was es Ihnen ermöglicht, Vektormodelle zu kürzen, ohne sie neu zu berechnen:

# Full embeddings: 896d (0.5B) or 1536d (1.5B)
full_embedding = model.encode(text)

# Truncate to smaller dimensions for efficiency
small_embedding = full_embedding[:256]  # Works for both models
tiny_embedding = full_embedding[:128]   # 0.5B supports down to 64d
Kopieren

Diese Flexibilität ermöglicht es, zwischen Leistung und Effizienz basierend auf Ihren Anforderungen abzuwägen.

tagFazit

jina-code-embeddings zeigt, dass effektive Code-Vektormodelle keine massive Skalierung erfordern. Indem wir auf Code-Generierungsmodellen aufbauen und gezieltes Fine-Tuning anwenden, erzielen wir modernste Leistung mit Modellen unter 1,5 Milliarden Parametern.

Die starken Ergebnisse von solch kompakten Modellen (0,5B/1,5B) bestätigen unsere These: Die richtige Grundlage ist wichtiger als die Anzahl der Parameter. Generierungsmodelle verstehen die Codesemantik – dieses Verständnis wird direkt auf Darstellungsaufgaben übertragen.

Dies steht im Einklang mit unserer umfassenderen Vision bei Jina AI: einheitliche Architekturen, bei denen Vektormodelle und Generierung aus derselben Grundlage hervorgehen und die Grenzen dessen verschieben, was mit Such-Basismodellen möglich ist.