Views
No views yet

jina-code-embeddings is an embedding model for code retrieval.
The model supports various types of code retrieval (text-to-code, code-to-code, code-to-text, code-to-completion) and technical question answering across 15+ programming languages.jina-code-embeddings-1.5b features:| Feature | Jina Code Embeddings 1.5B |
|---|---|
| Base Model | Qwen2.5-Coder-1.5B |
| Supported Tasks | nl2code, code2code, code2nl, code2completion, qa |
| Model DType | BFloat 16 |
| Max Sequence Length | 32768 |
| Embedding Vector Dimension | 1536 |
| Matryoshka dimensions | 128, 256, 512, 1024, 1536 |
| Pooling Strategy | Last-token pooling |
| Attention Mechanism | FlashAttention2 |
transformers>=4.53.0torch>=2.7.1sentence-transformers interface, install this package as well.1# !pip install transformers>=4.53.0 torch>=2.7.1
2
3import torch
4import torch.nn.functional as F
5
6from transformers import AutoModel, AutoTokenizer
7
8INSTRUCTION_CONFIG = {
9 "nl2code": {
10 "query": "Find the most relevant code snippet given the following query:\n",
11 "passage": "Candidate code snippet:\n"
12 },
13 "qa": {
14 "query": "Find the most relevant answer given the following question:\n",
15 "passage": "Candidate answer:\n"
16 },
17 "code2code": {
18 "query": "Find an equivalent code snippet given the following code snippet:\n",
19 "passage": "Candidate code snippet:\n"
20 },
21 "code2nl": {
22 "query": "Find the most relevant comment given the following code snippet:\n",
23 "passage": "Candidate comment:\n"
24 },
25 "code2completion": {
26 "query": "Find the most relevant completion given the following start of code snippet:\n",
27 "passage": "Candidate completion:\n"
28 }
29}
30
31MAX_LENGTH = 8192
32
33def cosine_similarity(x,y):
34 x = F.normalize(x, p=2, dim=1)
35 y = F.normalize(y, p=2, dim=1)
36 return x @ y.T
37
38def last_token_pool(last_hidden_states, attention_mask):
39 left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
40 if left_padding:
41 return last_hidden_states[:, -1]
42 else:
43 sequence_lengths = attention_mask.sum(dim=1) - 1
44 batch_size = last_hidden_states.shape[0]
45 return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
46
47def add_instruction(instruction, query):
48 return f'{instruction}{query}'
49
50# The queries and documents to embed
51queries = [
52 add_instruction(INSTRUCTION_CONFIG["nl2code"]["query"], "print hello world in python"),
53 add_instruction(INSTRUCTION_CONFIG["nl2code"]["query"], "initialize array of 5 zeros in c++")
54]
55documents = [
56 add_instruction(INSTRUCTION_CONFIG["nl2code"]["passage"], "print('Hello World!')"),
57 add_instruction(INSTRUCTION_CONFIG["nl2code"]["passage"], "int arr[5] = {0, 0, 0, 0, 0};")
58]
59all_inputs = queries + documents
60
61tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-code-embeddings-1.5b')
62model = AutoModel.from_pretrained('jinaai/jina-code-embeddings-1.5b')
63
64batch_dict = tokenizer(
65 all_inputs,
66 padding=True,
67 truncation=True,
68 max_length=MAX_LENGTH,
69 return_tensors="pt",
70)
71batch_dict.to(model.device)
72outputs = model(**batch_dict)
73embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
74query_embeddings = embeddings[:2]
75passage_embeddings = embeddings[2:]
76
77# Compute the (cosine) similarity between the query and document embeddings
78scores = cosine_similarity(query_embeddings, passage_embeddings)
79print(scores)
80# tensor([[0.7647, 0.1115],
81# [0.0930, 0.6606]], grad_fn=<MmBackward0>)1# !pip install sentence_transformers>=5.0.0 torch>=2.7.1
2
3import torch
4from sentence_transformers import SentenceTransformer
5
6# Load the model
7model = SentenceTransformer(
8 "jinaai/jina-code-embeddings-1.5b",
9 model_kwargs={
10 "torch_dtype": torch.bfloat16,
11 "attn_implementation": "flash_attention_2",
12 "device_map": "cuda"
13 },
14 tokenizer_kwargs={"padding_side": "left"},
15)
16
17# The queries and documents to embed
18queries = [
19 "print hello world in python",
20 "initialize array of 5 zeros in c++"
21]
22documents = [
23 "print('Hello World!')",
24 "int arr[5] = {0, 0, 0, 0, 0};"
25]
26
27query_embeddings = model.encode(queries, prompt_name="nl2code_query")
28document_embeddings = model.encode(documents, prompt_name="nl2code_document")
29
30# Compute the (cosine) similarity between the query and document embeddings
31similarity = model.similarity(query_embeddings, document_embeddings)
32print(similarity)
33# tensor([[0.7670, 0.1117],
34# [0.0938, 0.6607]])1
2import torch
3import torch.nn.functional as F
4from vllm import LLM
5
6INSTRUCTION_CONFIG = {
7 "nl2code": {
8 "query": "Find the most relevant code snippet given the following query:\n",
9 "passage": "Candidate code snippet:\n"
10 },
11 "qa": {
12 "query": "Find the most relevant answer given the following question:\n",
13 "passage": "Candidate answer:\n"
14 },
15 "code2code": {
16 "query": "Find an equivalent code snippet given the following code snippet:\n",
17 "passage": "Candidate code snippet:\n"
18 },
19 "code2nl": {
20 "query": "Find the most relevant comment given the following code snippet:\n",
21 "passage": "Candidate comment:\n"
22 },
23 "code2completion": {
24 "query": "Find the most relevant completion given the following start of code snippet:\n",
25 "passage": "Candidate completion:\n"
26 }
27}
28
29def add_instruction(instruction, text):
30 return f"{instruction}{text}"
31
32def cosine_similarity(x, y):
33 x = F.normalize(x, p=2, dim=1)
34 y = F.normalize(y, p=2, dim=1)
35 return x @ y.T
36
37# Build the queries and documents
38queries = [
39 add_instruction(INSTRUCTION_CONFIG["nl2code"]["query"], "print hello world in python"),
40 add_instruction(INSTRUCTION_CONFIG["nl2code"]["query"], "initialize array of 5 zeros in c++"),
41]
42documents = [
43 add_instruction(INSTRUCTION_CONFIG["nl2code"]["passage"], "print('Hello World!')"),
44 add_instruction(INSTRUCTION_CONFIG["nl2code"]["passage"], "int arr[5] = {0, 0, 0, 0, 0};"),
45]
46all_inputs = queries + documents
47
48# vLLM embedding model
49llm = LLM(
50 model="jinaai/jina-code-embeddings-1.5b",
51 task="embed"
52)
53
54# Encode with vLLM
55outputs = llm.encode(all_inputs)
56
57# Collect embeddings into a single tensor
58emb_list = []
59for out in outputs:
60 vec = out.outputs.data.detach()
61 emb_list.append(vec)
62embeddings = torch.stack(emb_list, dim=0)
63
64# Split into query and passage embeddings
65n_q = len(queries)
66query_embeddings = embeddings[:n_q]
67passage_embeddings = embeddings[n_q:]
68
69# Cosine similarity matrix (queries x documents)
70scores = cosine_similarity(query_embeddings, passage_embeddings)
71print(scores)
72# tensor([[0.7650, 0.1118],
73# [0.0937, 0.6613]])@misc{kryvosheieva2025efficientcodeembeddingscode,
title={Efficient Code Embeddings from Code Generation Models},
author={Daria Kryvosheieva and Saba Sturua and Michael Günther and Scott Martens and Han Xiao},
year={2025},
eprint={2508.21290},
archivePrefix={arXiv},
primaryClass={cs.CL},
url={https://arxiv.org/abs/2508.21290},
}