Jina Code Embeddings: Recuperación de código SOTA a 0.5B y 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.

Hoy lanzamos jina-code-embeddings, un nuevo conjunto de modelos de vectorización de código en dos tamaños (0.5B y 1.5B parámetros), junto con cuantificaciones GGUF para ambos. Construidos sobre LLM de generación de código autorregresivos, estos modelos logran un rendimiento de recuperación de última generación a pesar de su tamaño compacto. Admiten más de 15 lenguajes de programación, incluidos Python, JavaScript, Java, C++, C#, Go, Rust, TypeScript, SQL, MATLAB, R, Swift, Kotlin, HTML/CSS, PHP, Ruby, Scala, Perl y Shell.

jina-code-embeddings logra un rendimiento medio del 78.41 % (0.5B) y del 79.04 % (1.5B) en 25 puntos de referencia de recuperación de código. El modelo de 0.5B supera a Qwen3-Embedding-0.6B en 5 puntos porcentuales a pesar de ser un 20 % más pequeño, mientras que la variante de 1.5B coincide con voyage-code-3 (79.23 %) y supera a gemini-embedding-001 (77.38 %), ambos modelos patentados con arquitecturas no reveladas.

*Closed-source models with undisclosed architecture

Ambos modelos se entrenaron con cinco prefijos de instrucción específicos de la tarea para diferentes escenarios de recuperación, cada uno de los cuales admite roles de consulta y documento para la recuperación asimétrica. Por ejemplo, puede usar nl2code_query para insertar consultas y nl2code_document para insertar documentos.

tagReceta de entrenamiento

Utilizamos modelos de generación de código preentrenados como backbones de vectorización. Construidos sobre Qwen2.5-Coder-0.5B y 1.5B, nuestros modelos cuentan con:

Los modelos tradicionales de vectorización de código se enfrentan a un cuello de botella fundamental: simplemente no hay suficientes pares comentario-código de alta calidad para el entrenamiento supervisado. Al comenzar con Qwen2.5-Coder preentrenado en 5.5 billones de tokens que abarcan más de 92 lenguajes de programación, heredamos una comprensión semántica profunda de las construcciones de programación, el reconocimiento de patrones entre lenguajes y el conocimiento integrado de la sintaxis y los modismos. El ajuste fino contrastivo luego adapta este conocimiento para las tareas de recuperación con datos alineados mínimos, evitando la escasez de datos que restringe los modelos solo de codificador.

Para las tareas subrepresentadas, como las traducciones de código entre marcos, generamos datos sintéticos utilizando LLM, con cada ejemplo sintético validado manualmente en cuanto a calidad. Nuestros datos de entrenamiento combinaron divisiones de entrenamiento de tareas de código MTEB existentes con conjuntos de datos públicos adaptados, incluidos CommitPackFT, SWE-Bench, Spider, MBPP y CodeSearchNet.

A diferencia de jina-embeddings-v3 y v4, no usamos LoRA y pasamos directamente al post-entrenamiento completo. Para modelos pequeños como los nuestros (494M y 1.54B parámetros), la eficiencia de los parámetros de LoRA se vuelve menos convincente: la sobrecarga del adaptador en realidad puede afectar el rendimiento cuando tiene una capacidad limitada. Necesitábamos que cada parámetro funcionara en la tarea de vectorización. Incluso para escenarios de múltiples tareas, los prefijos de instrucción específicos de la tarea demostraron ser más limpios que múltiples adaptadores LoRA. En lugar de cambiar las configuraciones de peso, simplemente anteponemos diferentes instrucciones, mucho más ajustado y más alineado con la forma en que los LLM procesan naturalmente la información condicional.

El entrenamiento fue notablemente eficiente: ambos modelos se entrenaron utilizando el aprendizaje contrastivo con la pérdida InfoNCE en 4 GPU A100 de 80 GB, y se completó en solo 8.3 horas para el modelo de 0.5B y 12 horas para la variante de 1.5B.

Finalmente, evaluamos diferentes estrategias de agrupación. La agrupación del último token logró un promedio general del 78.41 %, superando constantemente la agrupación media (77.20 %) y la agrupación de atención latente (78.27 %) en todas las categorías de referencia. Esta ventaja de 1.2 puntos porcentuales nos llevó a romper con la tradición de agrupación media que establecimos en jina-embeddings-v2, v3 y v4. A medida que más modelos de recuperación se basan en LLM solo de decodificador, la agrupación del último token se convierte en la opción natural: la agrupación media simplemente no se alinea bien con los mecanismos de atención unidireccional. Si bien la agrupación media puede funcionar y, a menudo, se entrena más fácilmente en los primeros pasos (probablemente debido a su panorama de optimización convexo), nuestros experimentos muestran constantemente que se estanca por debajo del límite de rendimiento que logra la agrupación del último token.

tagEmpezando

Ambos modelos funcionan a la perfección a través de nuestra API de Search Foundation y con marcos populares, incluidos sentence-transformers, transformers y llama.cpp

tagA través de la 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
Copiar

tagA través de 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)
Copiar

tagA través de transformers

tagCorte de los Vectores Modelo Matryoshka

Ambos modelos fueron entrenados con el aprendizaje de representación Matryoshka para las dimensiones [64, 128, 256, 512, 896], lo que le permite truncar los Vectores Modelo sin volver a calcularlos:

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

Esta flexibilidad permite compensar el rendimiento y la eficiencia en función de sus requisitos.

tagConclusión

jina-code-embeddings demuestra que los Vectores Modelo de código eficaces *no* requieren una escala masiva. Al basarnos en los modelos de generación de código y aplicar un ajuste fino específico, logramos un rendimiento de última generación con modelos de menos de 1.500 millones de parámetros.

Los sólidos resultados de modelos tan compactos (0.5B/1.5B) validan nuestra tesis: **la base correcta importa más que el número de parámetros.** Los modelos de generación comprenden la semántica del código; esa comprensión se transfiere directamente a las tareas de representación.

Esto se alinea con nuestra visión más amplia en Jina AI: arquitecturas unificadas donde la inserción y la generación emergen de la misma base, superando los límites de lo que es posible con los modelos de base de búsqueda.