Jina Code Embeddings: SOTA для извлечения кода при 0.5B и 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.

Сегодня мы выпускаем jina-code-embeddings, новый набор моделей векторного представления кода в двух размерах — 0.5B и 1.5B параметров — вместе с квантованиями GGUF для обеих. Созданные на основе авторегрессионных LLM для генерации кода, эти модели достигают современной производительности поиска, несмотря на свой компактный размер. Они поддерживают более 15 языков программирования, включая Python, JavaScript, Java, C++, C#, Go, Rust, TypeScript, SQL, MATLAB, R, Swift, Kotlin, HTML/CSS, PHP, Ruby, Scala, Perl и Shell.

jina-code-embeddings достигает 78.41% (0.5B) и 79.04% (1.5B) средней производительности по 25 эталонным тестам поиска кода. Модель 0.5B превосходит Qwen3-Embedding-0.6B на 5 процентных пунктов, несмотря на то, что она на 20% меньше, а вариант 1.5B соответствует voyage-code-3 (79.23%) и превосходит gemini-embedding-001 (77.38%) — обе проприетарные модели с нераскрытой архитектурой.

*Closed-source models with undisclosed architecture

Обе модели были обучены с использованием пяти специфических для каждой задачи префиксов инструкций для различных сценариев поиска, каждый из которых поддерживает роли как запроса, так и документа для асимметричного поиска. Например, вы можете использовать nl2code_query для встраивания запросов и nl2code_document для встраивания документов.

tagTraining Recipe

Мы используем предварительно обученные модели генерации кода в качестве основы для встраивания. Наши модели, построенные на основе Qwen2.5-Coder-0.5B и 1.5B, имеют следующие особенности:

Традиционные модели векторного представления кода сталкиваются с фундаментальным узким местом: просто не хватает высококачественных пар комментарий-код для контролируемого обучения. Начиная с Qwen2.5-Coder, предварительно обученной на 5.5 триллионах токенов, охватывающих более 92 языков программирования, мы наследуем глубокое семантическое понимание конструкций программирования, распознавание межъязыковых шаблонов и встроенные знания синтаксиса и идиом. Контрастная тонкая настройка затем адаптирует эти знания для задач поиска с минимальным количеством выровненных данных, обходя дефицит данных, который ограничивает модели только с кодировщиком.

Для недостаточно представленных задач, таких как межфреймворковые переводы кода, мы сгенерировали синтетические данные с использованием LLM, при этом каждый синтетический пример был вручную проверен на качество. Наши обучающие данные объединили существующие разделения обучения задач кода MTEB с адаптированными общедоступными наборами данных, включая CommitPackFT, SWE-Bench, Spider, MBPP и CodeSearchNet.

В отличие от jina-embeddings-v3 и v4, мы не использовали LoRA и сразу перешли к полной постобработке. Для небольших моделей, таких как наши (494M и 1.54B параметров), эффективность параметров LoRA становится менее убедительной — накладные расходы адаптера могут фактически ухудшить производительность, когда у вас ограничена емкость. Нам нужно было, чтобы каждый параметр работал над задачей встраивания. Даже для многозадачных сценариев префиксы инструкций, специфичные для каждой задачи, оказались более чистыми, чем несколько адаптеров LoRA. Вместо переключения конфигураций весов мы просто добавляем разные инструкции — намного более компактные и более соответствующие тому, как LLM естественно обрабатывают условную информацию.

Обучение было на удивление эффективным: обе модели были обучены с использованием контрастного обучения с потерей InfoNCE на 4x A100 80GB GPU, завершившись всего за 8.3 часа для модели 0.5B и 12 часов для варианта 1.5B.

Наконец, мы сравнили различные стратегии объединения. Объединение последнего токена достигло 78.41% общего среднего значения, стабильно превосходя среднее объединение (77.20%) и объединение скрытого внимания (78.27%) по всем категориям эталонных тестов. Это преимущество в 1.2 процентных пункта побудило нас отказаться от традиции среднего объединения, которую мы установили в jina-embeddings-v2, v3 и v4. Поскольку все больше моделей поиска строятся на основе LLM только с декодером, объединение последнего токена становится естественным выбором — среднее объединение просто плохо согласуется с однонаправленными механизмами внимания. Хотя среднее объединение может работать и часто легче обучается на ранних этапах (вероятно, из-за его выпуклого ландшафта оптимизации), наши эксперименты последовательно показывают, что оно достигает плато ниже предела производительности, которого достигает объединение последнего токена.

tagGetting Started

Обе модели беспрепятственно работают через наш Search Foundation API и с популярными фреймворками, включая sentence-transformers, transformers и 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
Копировать

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)
Копировать

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

tagОбрезание Matryoshka 向量模型

Обе модели были обучены с использованием Matryoshka representation learning для размерностей [64, 128, 256, 512, 896], что позволяет обрезать 向量模型 без пересчета:

# Полные 向量模型: 896d (0.5B) или 1536d (1.5B)
full_embedding = model.encode(text)

# Обрезайте до меньших размеров для эффективности
small_embedding = full_embedding[:256]  # Работает для обеих моделей
tiny_embedding = full_embedding[:128]   # 0.5B поддерживает до 64d
Копировать

Эта гибкость позволяет найти компромисс между производительностью и эффективностью в зависимости от ваших требований.

tagЗаключение

jina-code-embeddings демонстрирует, что для эффективных 向量模型 кода не требуется огромный масштаб. Основываясь на моделях генерации кода и применяя целевую доводку, мы достигаем современной производительности с моделями менее 1,5B параметров.

Сильные результаты таких компактных моделей (0,5B/1,5B) подтверждают наш тезис: правильный фундамент важнее количества параметров. Модели генерации понимают семантику кода — это понимание напрямую переносится на задачи представления.

Это соответствует нашему более широкому видению в Jina AI: унифицированные архитектуры, в которых 向量模型 и генерация возникают из одного и того же фундамента, расширяя границы того, что возможно с поисковыми базовыми моделями.