Jina Code Embeddings: Recupero di codice SOTA a 0,5B e 1,5B

Efficient Code Embeddings from Code Generation Models

jina-code-embeddings è una nuova suite di modelli di *embeddings* di codice progettata per recuperare codice da query in linguaggio naturale, eseguire attività di question-answering tecnico e identificare snippet di codice semanticamente simili tra diversi linguaggi di programmazione. Sfrutta in modo innovativo un backbone autoregressivo pre-addestrato sia su testo che su codice, generando *embeddings* tramite il *last-token pooling*. Descriviamo la ricetta di addestramento e dimostriamo prestazioni all'avanguardia nonostante le dimensioni relativamente ridotte dei modelli, convalidando questo approccio alla costruzione di modelli di *embeddings* di codice.

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.

Oggi rilasciamo jina-code-embeddings, una nuova suite di modelli di *embeddings* di codice in due dimensioni (0.5B e 1.5B parametri), insieme a quantizzazioni GGUF per entrambi. Costruiti su *LLM* di generazione di codice autoregressivi, questi modelli raggiungono prestazioni di recupero all'avanguardia nonostante le loro dimensioni compatte. Supportano oltre 15 linguaggi di programmazione tra cui Python, JavaScript, Java, C++, C#, Go, Rust, TypeScript, SQL, MATLAB, R, Swift, Kotlin, HTML/CSS, PHP, Ruby, Scala, Perl e Shell.

jina-code-embeddings raggiunge una performance media del 78.41% (0.5B) e del 79.04% (1.5B) su 25 benchmark di recupero di codice. Il modello da 0.5B supera Qwen3-Embedding-0.6B di 5 punti percentuali pur essendo il 20% più piccolo, mentre la variante da 1.5B corrisponde a voyage-code-3 (79.23%) e supera gemini-embedding-001 (77.38%), entrambi modelli proprietari con architetture non divulgate.

*Modelli a codice chiuso con architettura non divulgata

Entrambi i modelli sono stati addestrati con cinque prefissi di istruzioni specifiche per attività per diversi scenari di recupero, ciascuno dei quali supporta sia i ruoli di query che di documento per il recupero asimmetrico. Ad esempio, è possibile utilizzare nl2code_query per incorporare le query e nl2code_document per incorporare i documenti.

tagRicetta di addestramento

Utilizziamo modelli di generazione di codice pre-addestrati come backbone di *embedding*. Costruiti su Qwen2.5-Coder-0.5B e 1.5B, i nostri modelli presentano:

I modelli di *embeddings* di codice tradizionali affrontano un collo di bottiglia fondamentale: semplicemente non ci sono abbastanza coppie commento-codice di alta qualità per l'addestramento supervisionato. Partendo da Qwen2.5-Coder pre-addestrato su 5.5 trilioni di *tokens* che coprono oltre 92 linguaggi di programmazione, ereditiamo una profonda comprensione semantica delle costruzioni di programmazione, il riconoscimento di pattern cross-linguaggio e la conoscenza integrata della sintassi e degli idiomi. La messa a punto contrastiva adatta quindi questa conoscenza per le attività di recupero con dati allineati minimi, eludendo la scarsità di dati che vincola i modelli solo encoder.

Per le attività sottorappresentate come le traduzioni di codice cross-framework, abbiamo generato dati sintetici utilizzando *LLM*, con ogni esempio sintetico convalidato manualmente per la qualità. I nostri dati di addestramento combinano le divisioni di addestramento delle attività di codice MTEB esistenti con set di dati pubblici adattati tra cui CommitPackFT, SWE-Bench, Spider, MBPP e CodeSearchNet.

A differenza di jina-embeddings-v3 e v4, non abbiamo utilizzato LoRA e siamo passati direttamente al post-training completo. Per i modelli piccoli come i nostri (494M e 1.54B parametri), l'efficienza dei parametri di LoRA diventa meno interessante: l'overhead dell'adattatore può effettivamente danneggiare le prestazioni quando si ha una capacità limitata. Avevamo bisogno che ogni parametro lavorasse sull'attività di *embedding*. Anche per gli scenari multi-task, i prefissi di istruzioni specifiche per attività si sono dimostrati più puliti rispetto a più adattatori LoRA. Invece di commutare le configurazioni dei pesi, anteponiamo semplicemente istruzioni diverse, molto più snelle e più allineate al modo in cui gli *LLM* elaborano naturalmente le informazioni condizionali.

L'addestramento è stato straordinariamente efficiente: entrambi i modelli sono stati addestrati utilizzando l'apprendimento contrastivo con perdita InfoNCE su 4 GPU A100 da 80 GB, completando in sole 8.3 ore per il modello da 0.5B e 12 ore per la variante da 1.5B.

Infine, abbiamo valutato diverse strategie di *pooling*. Il *last-token pooling* ha raggiunto una media complessiva del 78.41%, superando costantemente il *mean pooling* (77.20%) e il *latent attention pooling* (78.27%) in tutte le categorie di benchmark. Questo vantaggio di 1.2 punti percentuali ci ha portato a rompere con la tradizione del *mean pooling* che abbiamo stabilito in jina-embeddings-v2, v3 e v4. Man mano che sempre più modelli di recupero si basano su *LLM* solo decoder, il *last-token pooling* diventa la scelta naturale: il *mean pooling* semplicemente non si allinea bene con i meccanismi di attenzione unidirezionali. Sebbene il *mean pooling* possa funzionare e spesso si addestri più facilmente nelle prime fasi (probabilmente a causa del suo panorama di ottimizzazione convesso), i nostri esperimenti mostrano costantemente che si stabilizza al di sotto del limite di prestazione che il *last-token pooling* raggiunge.

tagIniziare

Entrambi i modelli funzionano perfettamente tramite la nostra Search Foundation API e con framework popolari tra cui sentence-transformers, transformers e llama.cpp

tagTramite 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
Copy

tagTramite 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)
Copy

tagTramite 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'])
Copy

tagMatryoshka Embeddings Cut-Off

Entrambi i modelli sono stati addestrati con l'apprendimento della rappresentazione di Matryoshka per le dimensioni [64, 128, 256, 512, 896], consentendoti di troncare gli embeddings senza ricalcolare:

# 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
Copy

Questa flessibilità consente di bilanciare prestazioni ed efficienza in base alle proprie esigenze.

tagConclusione

jina-code-embeddings dimostra che gli embeddings di codice efficaci non richiedono una scala massiccia. Basandoci sui modelli di generazione di codice e applicando un fine-tuning mirato, otteniamo prestazioni all'avanguardia con modelli inferiori a 1,5 miliardi di parametri.

I solidi risultati di modelli così compatti (0,5B/1,5B) convalidano la nostra tesi: la giusta base conta più del numero di parametri. I modelli di generazione comprendono la semantica del codice: questa comprensione si trasferisce direttamente alle attività di rappresentazione.

Questo è in linea con la nostra visione più ampia in Jina AI: architetture unificate in cui l'embedding e la generazione emergono dalla stessa base, spingendo i confini di ciò che è possibile con i modelli di fondazione di ricerca.