@zuvia-software-solutions/code-mapper 2.3.0 → 2.3.2

package/dist/cli/analyze.js

@@ -266,17 +266,24 @@ export const analyzeCommand = async (inputPath, options) => {
                         continue;
                     try {
                         const msg = JSON.parse(line);
-                        if (msg.phase === 'loaded') {
+                        if (msg.phase === 'downloading' || msg.phase === 'converting') {
+                            updateBar(90, msg.message);
+                        }
+                        else if (msg.phase === 'loaded') {
                             updateBar(91, `Model loaded (${msg.load_ms}ms)`);
                         }
                         else if (msg.phase === 'queried') {
-                            updateBar(92, `Found ${msg.nodes} embeddable nodes`);
+                            updateBar(92, `Found ${msg.nodes} embeddable nodes${msg.skipped_tests ? ` (${msg.skipped_tests} test files skipped)` : ''}`);
                         }
                         else if (msg.phase === 'prepared') {
                             updateBar(93, `${msg.to_embed} to embed, ${msg.skipped} cached`);
                         }
+                        else if (msg.phase === 'embedding') {
+                            const scaled = 93 + Math.round((msg.progress / 100) * 4);
+                            updateBar(scaled, `Embedding... ${msg.progress}% (${msg.embedded} written)`);
+                        }
                         else if (msg.phase === 'embedded') {
-                            updateBar(96, `Embedded ${msg.count} nodes (${(msg.ms / 1000).toFixed(1)}s)`);
+                            updateBar(97, `Embedded ${msg.count} nodes (${(msg.ms / 1000).toFixed(1)}s)`);
                         }
                         else if (msg.phase === 'done') {
                             updateBar(98, `Embeddings complete (${msg.embedded} new, ${msg.skipped} cached)`);

package/models/jina-code-0.5b-mlx/config.json

@@ -0,0 +1,73 @@
+{
+  "architectures": [
+    "Qwen2ForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "eos_token_id": 151643,
+  "hidden_act": "silu",
+  "hidden_size": 896,
+  "initializer_range": 0.02,
+  "intermediate_size": 4864,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "matryoshka_dims": [
+    64,
+    128,
+    256,
+    512,
+    896
+  ],
+  "max_position_embeddings": 32768,
+  "max_window_layers": 24,
+  "model_type": "qwen2",
+  "num_attention_heads": 14,
+  "num_hidden_layers": 24,
+  "num_key_value_heads": 2,
+  "prompt_names": [
+    "query",
+    "passage"
+  ],
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000.0,
+  "sliding_window": null,
+  "task_names": [
+    "nl2code",
+    "qa",
+    "code2code",
+    "code2nl",
+    "code2completion"
+  ],
+  "tie_word_embeddings": true,
+  "tokenizer_class": "Qwen2TokenizerFast",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.53.0",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

package/models/jina-code-0.5b-mlx/model.py

@@ -0,0 +1,127 @@
+
+from dataclasses import dataclass
+from typing import Optional, List
+import mlx.core as mx
+import mlx.nn as nn
+
+@dataclass
+class ModelArgs:
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    max_position_embeddings: int
+    rope_theta: float = 1000000.0
+    tie_word_embeddings: bool = True
+
+class Attention(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        dim = args.hidden_size
+        self.n_heads = args.num_attention_heads
+        self.n_kv_heads = args.num_key_value_heads
+        self.head_dim = dim // self.n_heads
+        self.scale = self.head_dim ** -0.5
+        self.rope_theta = args.rope_theta
+        self.q_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=True)
+        self.k_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=True)
+        self.v_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=True)
+        self.o_proj = nn.Linear(self.n_heads * self.head_dim, dim, bias=False)
+
+    def __call__(self, x, mask=None):
+        B, L, D = x.shape
+        q = self.q_proj(x).reshape(B, L, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
+        k = self.k_proj(x).reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
+        v = self.v_proj(x).reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
+        q = mx.fast.rope(q, self.head_dim, traditional=False, base=self.rope_theta, scale=1.0, offset=0)
+        k = mx.fast.rope(k, self.head_dim, traditional=False, base=self.rope_theta, scale=1.0, offset=0)
+        out = mx.fast.scaled_dot_product_attention(q, k, v, mask=mask.astype(q.dtype) if mask is not None else None, scale=self.scale)
+        return self.o_proj(out.transpose(0, 2, 1, 3).reshape(B, L, -1))
+
+class MLP(nn.Module):
+    def __init__(self, dim, hidden):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden, bias=False)
+        self.down_proj = nn.Linear(hidden, dim, bias=False)
+        self.up_proj = nn.Linear(dim, hidden, bias=False)
+    def __call__(self, x):
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.self_attn = Attention(args)
+        self.mlp = MLP(args.hidden_size, args.intermediate_size)
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+    def __call__(self, x, mask=None):
+        h = x + self.self_attn(self.input_layernorm(x), mask)
+        return h + self.mlp(self.post_attention_layernorm(h))
+
+class Qwen2Model(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [TransformerBlock(args) for _ in range(args.num_hidden_layers)]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+    def __call__(self, inputs, mask=None):
+        h = self.embed_tokens(inputs)
+        for layer in self.layers:
+            h = layer(h, mask)
+        return self.norm(h)
+
+class JinaCodeEmbeddingModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        args = ModelArgs(
+            hidden_size=config["hidden_size"],
+            num_hidden_layers=config["num_hidden_layers"],
+            intermediate_size=config["intermediate_size"],
+            num_attention_heads=config["num_attention_heads"],
+            rms_norm_eps=config["rms_norm_eps"],
+            vocab_size=config["vocab_size"],
+            num_key_value_heads=config["num_key_value_heads"],
+            max_position_embeddings=config["max_position_embeddings"],
+            rope_theta=config.get("rope_theta", 1000000.0),
+        )
+        self.model = Qwen2Model(args)
+        self.config = config
+
+    def __call__(self, input_ids, attention_mask=None):
+        B, L = input_ids.shape
+        causal = mx.tril(mx.ones((L, L)))
+        causal = mx.where(causal == 0, -1e4, 0.0)[None, None, :, :]
+        if attention_mask is not None:
+            pad = mx.where(attention_mask == 0, -1e4, 0.0)[:, None, None, :]
+            mask = causal + pad
+        else:
+            mask = causal
+        h = self.model(input_ids, mask)
+        if attention_mask is not None:
+            seq_lens = mx.sum(attention_mask.astype(mx.int32), axis=1) - 1
+            embs = h[mx.arange(B), seq_lens]
+        else:
+            embs = h[:, -1, :]
+        norms = mx.linalg.norm(embs, axis=1, keepdims=True)
+        return embs / norms
+
+    def encode(self, texts, tokenizer, max_length=8192, truncate_dim=None, task="nl2code", prompt_type="query"):
+        PREFIXES = {"nl2code": {"query": "Find the most relevant code snippet given the following query:\n", "passage": "Candidate code snippet:\n"}}
+        prefix = PREFIXES.get(task, {}).get(prompt_type, "")
+        if prefix:
+            texts = [prefix + t for t in texts]
+        encodings = tokenizer.encode_batch(texts)
+        ml = min(max_length, max(len(e.ids) for e in encodings))
+        iids, amask = [], []
+        for e in encodings:
+            ids = e.ids[:ml]; m = e.attention_mask[:ml]; p = ml - len(ids)
+            if p > 0: ids = ids + [0]*p; m = m + [0]*p
+            iids.append(ids); amask.append(m)
+        embs = self(mx.array(iids), mx.array(amask))
+        if truncate_dim:
+            embs = embs[:, :truncate_dim]
+            embs = embs / mx.linalg.norm(embs, axis=1, keepdims=True)
+        return embs

package/models/mlx-embedder.py

@@ -35,27 +35,54 @@ MODEL_DIR = os.path.dirname(os.path.abspath(__file__)) + "/jina-code-0.5b-mlx"
 
 
 def ensure_model_downloaded():
-    """Download model weights from HuggingFace if not present."""
+    """Download and convert model weights from HuggingFace if not present."""
     weights_path = os.path.join(MODEL_DIR, "model.safetensors")
     if os.path.exists(weights_path):
         return
 
-    print(json.dumps({"phase": "downloading", "message": "Downloading embedding model (~1.1GB, first time only)..."}), flush=True)
+    os.makedirs(MODEL_DIR, exist_ok=True)
+    print(json.dumps({"phase": "downloading", "message": "Downloading Jina Code 0.5B embedding model (~940MB, first time only)..."}), flush=True)
+
     try:
         from huggingface_hub import hf_hub_download
         import shutil
-        repo = "jinaai/jina-embeddings-v5-text-small-retrieval-mlx"
-        for fname in ["model.safetensors", "tokenizer.json", "vocab.json", "merges.txt", "tokenizer_config.json"]:
+
+        repo = "jinaai/jina-code-embeddings-0.5b"
+
+        # Download tokenizer from the v5 MLX model (same Qwen tokenizer, has tokenizer.json)
+        tokenizer_repo = "jinaai/jina-embeddings-v5-text-small-retrieval-mlx"
+        for fname in ["tokenizer.json"]:
+            dest = os.path.join(MODEL_DIR, fname)
+            if not os.path.exists(dest):
+                path = hf_hub_download(tokenizer_repo, fname)
+                shutil.copy(path, dest)
+
+        # Download vocab/merges from the 0.5B model
+        for fname in ["vocab.json", "merges.txt"]:
             dest = os.path.join(MODEL_DIR, fname)
             if not os.path.exists(dest):
                 path = hf_hub_download(repo, fname)
                 shutil.copy(path, dest)
-        print(json.dumps({"phase": "downloaded", "message": "Model downloaded successfully"}), flush=True)
+
+        # Download and convert weights: bf16 → fp16, add 'model.' key prefix
+        print(json.dumps({"phase": "converting", "message": "Converting model weights (bf16 → fp16)..."}), flush=True)
+        raw_path = hf_hub_download(repo, "model.safetensors")
+        raw_weights = mx.load(raw_path)
+        converted = {}
+        for k, v in raw_weights.items():
+            new_key = "model." + k
+            if v.dtype == mx.bfloat16:
+                v = v.astype(mx.float16)
+            converted[new_key] = v
+        mx.save_safetensors(weights_path, converted)
+
+        print(json.dumps({"phase": "downloaded", "message": f"Model ready ({len(converted)} weights converted)"}), flush=True)
+
     except ImportError:
         raise RuntimeError(
             "Model weights not found. Install huggingface_hub to auto-download:\n"
             "  pip3 install huggingface_hub\n"
-            "Or manually download from: https://huggingface.co/jinaai/jina-embeddings-v5-text-small-retrieval-mlx"
+            "Or manually download from: https://huggingface.co/jinaai/jina-code-embeddings-0.5b"
         )
 
 
@@ -388,26 +415,81 @@ def batch_mode(db_path, dims=256, max_tokens=2048):
     unique_texts = [v["text"] for v in unique_by_hash.values()]
     deduped = len(to_embed) - len(unique_texts)
 
-    # Embed only unique texts
+    # Embed unique texts in streaming fashion — process each batch, write to DB
+    # immediately, free GPU memory. Keeps peak memory at ONE batch instead of ALL.
     t0_embed = time.time()
-    embeddings = embed_tiered(model, tokenizer, unique_texts, "retrieval.passage", dims, max_tokens)
-    embed_ms = int((time.time() - t0_embed) * 1000)
+    unique_entries = list(unique_by_hash.values())
 
-    print(json.dumps({"phase": "embedded", "count": len(unique_texts), "deduped": deduped, "ms": embed_ms}), flush=True)
+    # Tokenize + sort (same as embed_tiered but we handle the loop here)
+    is_code_model = "jina-code" in MODEL_DIR
+    if is_code_model:
+        prefix_map = {"retrieval.query": "Find the most relevant code snippet given the following query:\n", "retrieval.passage": "Candidate code snippet:\n"}
+    else:
+        prefix_map = {"retrieval.query": "Query: ", "retrieval.passage": "Document: "}
+    prefix = prefix_map.get("retrieval.passage", "")
+    prefixed = [prefix + e["text"] for e in unique_entries]
+    encodings = tokenizer.encode_batch(prefixed)
+    indexed = sorted(range(len(unique_entries)), key=lambda i: len(encodings[i].ids))
 
-    # Write to database — copy embedding to all nodes sharing the same hash
-    t0_write = time.time()
+    embedded_count = 0
     db.execute("BEGIN")
-    for i, (text_hash, entry) in enumerate(unique_by_hash.items()):
-        emb = embeddings[i]
-        if emb is None:
-            continue
-        blob = float_list_to_blob(emb)
-        for nid, th in entry["node_ids"]:
-            db.execute("INSERT OR REPLACE INTO embeddings (nodeId, embedding, textHash) VALUES (?, ?, ?)",
-                       (nid, blob, th))
+
+    i = 0
+    while i < len(indexed):
+        peek_idx = indexed[min(i + 1, len(indexed) - 1)]
+        tok_count = min(len(encodings[peek_idx].ids), max_tokens)
+        batch_size = get_batch_size_for_tokens(tok_count)
+
+        batch_indices = []
+        batch_encs = []
+        while len(batch_encs) < batch_size and i < len(indexed):
+            orig_idx = indexed[i]
+            batch_indices.append(orig_idx)
+            batch_encs.append(encodings[orig_idx])
+            i += 1
+
+        max_len = min(max_tokens, max(len(e.ids) for e in batch_encs))
+        input_ids = []
+        attention_mask = []
+        for enc in batch_encs:
+            ids = enc.ids[:max_len]
+            mask = enc.attention_mask[:max_len]
+            pad = max_len - len(ids)
+            if pad > 0:
+                ids = ids + [0] * pad
+                mask = mask + [0] * pad
+            input_ids.append(ids)
+            attention_mask.append(mask)
+
+        # Forward pass
+        embs = model(mx.array(input_ids), mx.array(attention_mask))
+        if dims and dims < embs.shape[1]:
+            embs = embs[:, :dims]
+        norms = mx.linalg.norm(embs, axis=1, keepdims=True)
+        embs = embs / norms
+        mx.eval(embs)
+
+        # Convert to Python + write to DB immediately
+        emb_list = embs.tolist()
+        del embs  # free MLX GPU memory
+
+        for j, orig_idx in enumerate(batch_indices):
+            entry = unique_entries[orig_idx]
+            blob = float_list_to_blob(emb_list[j])
+            for nid, th in entry["node_ids"]:
+                db.execute("INSERT OR REPLACE INTO embeddings (nodeId, embedding, textHash) VALUES (?, ?, ?)",
+                           (nid, blob, th))
+            embedded_count += len(entry["node_ids"])
+
+        # Progress
+        pct = i * 100 // len(indexed)
+        print(json.dumps({"phase": "embedding", "progress": pct, "embedded": embedded_count}), flush=True)
+
     db.execute("COMMIT")
-    write_ms = int((time.time() - t0_write) * 1000)
+    embed_ms = int((time.time() - t0_embed) * 1000)
+    write_ms = 0  # included in embed_ms now
+
+    print(json.dumps({"phase": "embedded", "count": len(unique_entries), "deduped": deduped, "ms": embed_ms}), flush=True)
 
     total_ms = int((time.time() - t0_total) * 1000)
     print(json.dumps({

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@zuvia-software-solutions/code-mapper",
-  "version": "2.3.0",
+  "version": "2.3.2",
   "description": "Graph-powered code intelligence for AI agents. Index any codebase, query via MCP or CLI.",
   "author": "Abhigyan Patwari",
   "license": "PolyForm-Noncommercial-1.0.0",
@@ -36,8 +36,8 @@
     "skills",
     "vendor",
     "models/mlx-embedder.py",
-    "models/jina-v5-small-mlx/model.py",
-    "models/jina-v5-small-mlx/config.json"
+    "models/jina-code-0.5b-mlx/model.py",
+    "models/jina-code-0.5b-mlx/config.json"
   ],
   "scripts": {
     "build": "tsc",

package/models/jina-v5-small-mlx/config.json

@@ -1,19 +0,0 @@
-{
-  "model_type": "qwen3",
-  "hidden_size": 1024,
-  "num_hidden_layers": 28,
-  "intermediate_size": 3072,
-  "num_attention_heads": 16,
-  "num_key_value_heads": 8,
-  "rms_norm_eps": 1e-06,
-  "vocab_size": 151936,
-  "max_position_embeddings": 32768,
-  "rope_theta": 3500000,
-  "rope_parameters": {
-    "rope_theta": 3500000,
-    "rope_type": "default"
-  },
-  "head_dim": 128,
-  "tie_word_embeddings": true,
-  "rope_scaling": null
-}

package/models/jina-v5-small-mlx/model.py

@@ -1,260 +0,0 @@
-"""
-Jina Embeddings v5 Text Small - MLX Implementation
-
-Pure MLX port of jina-embeddings-v5-text-small (Qwen3-0.6B backbone).
-Zero dependency on PyTorch or transformers.
-
-Features:
-- Causal attention (decoder architecture)
-- QKNorm (q_norm/k_norm per head)
-- Last-token pooling
-- L2 normalization
-- Matryoshka embedding dimensions: [32, 64, 128, 256, 512, 768, 1024]
-- Max sequence length: 32768 tokens
-- Embedding dimension: 1024
-
-Architecture:
-- RoPE (rope_theta from config)
-- SwiGLU MLP
-- RMSNorm
-- QKNorm (RMSNorm on Q/K per head)
-- No attention bias
-"""
-
-from dataclasses import dataclass
-from typing import Any, Dict, Optional, Union
-
-import mlx.core as mx
-import mlx.nn as nn
-
-
-@dataclass
-class ModelArgs:
-    model_type: str
-    hidden_size: int
-    num_hidden_layers: int
-    intermediate_size: int
-    num_attention_heads: int
-    rms_norm_eps: float
-    vocab_size: int
-    num_key_value_heads: int
-    max_position_embeddings: int
-    head_dim: int
-    tie_word_embeddings: bool
-    rope_parameters: Optional[Dict[str, Union[float, str]]] = None
-    rope_theta: Optional[float] = None
-    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
-
-
-class Attention(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-
-        dim = args.hidden_size
-        self.n_heads = n_heads = args.num_attention_heads
-        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
-
-        head_dim = args.head_dim
-        self.scale = head_dim**-0.5
-        self.head_dim = head_dim
-
-        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
-        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
-        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
-        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
-
-        # Qwen3 has QKNorm
-        self.q_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
-        self.k_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
-
-        # Resolve rope_theta from config
-        if args.rope_parameters and 'rope_theta' in args.rope_parameters:
-            rope_theta = float(args.rope_parameters['rope_theta'])
-        elif args.rope_theta:
-            rope_theta = float(args.rope_theta)
-        else:
-            rope_theta = 10000.0
-        self.rope_theta = rope_theta
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-    ) -> mx.array:
-        B, L, D = x.shape
-
-        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
-
-        # Reshape and apply QKNorm
-        queries = self.q_norm(queries.reshape(B, L, self.n_heads, -1)).transpose(0, 2, 1, 3)
-        keys = self.k_norm(keys.reshape(B, L, self.n_kv_heads, -1)).transpose(0, 2, 1, 3)
-        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
-
-        # RoPE via mx.fast
-        queries = mx.fast.rope(queries, self.head_dim, traditional=False, base=self.rope_theta, scale=1.0, offset=0)
-        keys = mx.fast.rope(keys, self.head_dim, traditional=False, base=self.rope_theta, scale=1.0, offset=0)
-
-        # Scaled dot-product attention (handles GQA natively)
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values,
-            mask=mask.astype(queries.dtype) if mask is not None else None,
-            scale=self.scale,
-        )
-
-        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
-        return self.o_proj(output)
-
-
-class MLP(nn.Module):
-    def __init__(self, dim, hidden_dim):
-        super().__init__()
-        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
-        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
-        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
-
-    def __call__(self, x) -> mx.array:
-        gate = nn.silu(self.gate_proj(x))
-        return self.down_proj(gate * self.up_proj(x))
-
-
-class TransformerBlock(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.self_attn = Attention(args)
-        self.mlp = MLP(args.hidden_size, args.intermediate_size)
-        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
-        self.post_attention_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-    ) -> mx.array:
-        r = self.self_attn(self.input_layernorm(x), mask)
-        h = x + r
-        r = self.mlp(self.post_attention_layernorm(h))
-        out = h + r
-        return out
-
-
-class Qwen3Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
-        self.layers = [TransformerBlock(args=args) for _ in range(args.num_hidden_layers)]
-        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
-
-    def __call__(self, inputs: mx.array, mask: Optional[mx.array] = None):
-        h = self.embed_tokens(inputs)
-        for layer in self.layers:
-            h = layer(h, mask)
-        return self.norm(h)
-
-
-class JinaEmbeddingModel(nn.Module):
-    """Jina v5-text-small embedding model with last-token pooling."""
-
-    def __init__(self, config: dict):
-        super().__init__()
-        args = ModelArgs(**config)
-        self.model = Qwen3Model(args)
-        self.config = config
-
-    def __call__(
-        self,
-        input_ids: mx.array,
-        attention_mask: Optional[mx.array] = None,
-    ):
-        batch_size, seq_len = input_ids.shape
-
-        # Causal mask (Qwen3 is a decoder model)
-        causal_mask = mx.tril(mx.ones((seq_len, seq_len)))
-        causal_mask = mx.where(causal_mask == 0, -1e4, 0.0)
-        causal_mask = causal_mask[None, None, :, :]
-
-        # Combine with padding mask
-        if attention_mask is not None:
-            padding_mask = mx.where(attention_mask == 0, -1e4, 0.0)
-            padding_mask = padding_mask[:, None, None, :]
-            mask = causal_mask + padding_mask
-        else:
-            mask = causal_mask
-
-        hidden_states = self.model(input_ids, mask)
-
-        # Last token pooling
-        if attention_mask is not None:
-            sequence_lengths = mx.sum(attention_mask, axis=1) - 1
-            batch_indices = mx.arange(hidden_states.shape[0])
-            embeddings = hidden_states[batch_indices, sequence_lengths]
-        else:
-            embeddings = hidden_states[:, -1, :]
-
-        # L2 normalization
-        norms = mx.linalg.norm(embeddings, axis=1, keepdims=True)
-        embeddings = embeddings / norms
-
-        return embeddings
-
-    def encode(
-        self,
-        texts: list[str],
-        tokenizer,
-        max_length: int = 8192,
-        truncate_dim: Optional[int] = None,
-        task_type: str = "retrieval.query",
-    ):
-        """
-        Encode texts to embeddings.
-
-        Args:
-            texts: List of input texts
-            tokenizer: Tokenizer instance (from tokenizers library)
-            max_length: Maximum sequence length
-            truncate_dim: Optional Matryoshka dimension [32, 64, 128, 256, 512, 768, 1024]
-            task_type: Task prefix ("retrieval.query", "retrieval.passage", etc.)
-
-        Returns:
-            Embeddings array [batch, dim]
-        """
-        prefix_map = {
-            "retrieval.query": "Query: ",
-            "retrieval.passage": "Document: ",
-            "classification": "Document: ",
-            "text-matching": "Document: ",
-            "clustering": "Document: ",
-        }
-        prefix = prefix_map.get(task_type, "")
-
-        if prefix:
-            texts = [prefix + text for text in texts]
-
-        encodings = tokenizer.encode_batch(texts)
-
-        max_len = min(max_length, max(len(enc.ids) for enc in encodings))
-        input_ids = []
-        attention_mask = []
-
-        for encoding in encodings:
-            ids = encoding.ids[:max_len]
-            mask = encoding.attention_mask[:max_len]
-
-            pad_len = max_len - len(ids)
-            if pad_len > 0:
-                ids = ids + [0] * pad_len
-                mask = mask + [0] * pad_len
-
-            input_ids.append(ids)
-            attention_mask.append(mask)
-
-        input_ids = mx.array(input_ids)
-        attention_mask = mx.array(attention_mask)
-
-        embeddings = self(input_ids, attention_mask)
-
-        if truncate_dim is not None:
-            embeddings = embeddings[:, :truncate_dim]
-            norms = mx.linalg.norm(embeddings, axis=1, keepdims=True)
-            embeddings = embeddings / norms
-
-        return embeddings