Jina Code Embeddings 1.5B - MLX

MLX port of Jina AI's code embedding model for Apple Silicon.

ArXiv | Blog

Installation

pip install mlx tokenizers huggingface_hub

Usage

import mlx.core as mx
from tokenizers import Tokenizer
from model import JinaCodeEmbeddingModel
import json

# Load config
with open("config.json") as f:
    config = json.load(f)

# Load model
model = JinaCodeEmbeddingModel(config)
weights = mx.load("model.safetensors")
model.load_weights(list(weights.items()))
mx.eval(model.parameters())

# Load tokenizer
tokenizer = Tokenizer.from_file("tokenizer.json")

# Encode a natural language query for code search
query_embeddings = model.encode(
    ["print hello world in python"],
    tokenizer,
    task="nl2code",
    prompt_type="query",
)

# Encode code passages
code_embeddings = model.encode(
    ["print('Hello World!')"],
    tokenizer,
    task="nl2code",
    prompt_type="passage",
)

mx.eval(query_embeddings, code_embeddings)

Task Types

Each task uses specific prefixes for queries and passages:

Matryoshka Dimensions

Supports Matryoshka embedding truncation to: 128, 256, 512, 1024, 1536

embeddings = model.encode(texts, tokenizer, task="nl2code", prompt_type="query", truncate_dim=256)

Model Details

• Architecture: Qwen2.5-Coder-1.5B
• Parameters: 1.54B
• Embedding dimension: 1536
• Max sequence length: 32768 tokens
• Languages: 15+ programming languages
• Optimized for: Apple Silicon (M1/M2/M3/M4) with Metal acceleration

Files

jina-code-embeddings-1.5b-mlx/
├── model.safetensors          # Model weights (float16)
├── model.py                    # Model implementation
├── config.json                 # Model configuration
├── tokenizer.json              # Tokenizer
├── tokenizer_config.json
├── vocab.json
├── merges.txt
└── README.md

Citation

@misc{kryvosheieva2025efficientcodeembeddingscode,
      title={Efficient Code Embeddings from Code Generation Models},
      author={Daria Kryvosheieva and Saba Sturua and Michael G\"unther and Scott Martens and Han Xiao},
      year={2025},
      eprint={2508.21290},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2508.21290},
}

License

CC BY-NC 4.0

Links

• Original PyTorch model: jinaai/jina-code-embeddings-1.5b
• 0.5B MLX variant: jinaai/jina-code-embeddings-0.5b-mlx
• Jina AI: https://jina.ai