@zuvia-software-solutions/code-mapper 2.2.3 → 2.3.1

package/dist/cli/analyze.js

@@ -5,7 +5,7 @@ import { execFileSync } from 'child_process';
 import v8 from 'v8';
 import cliProgress from 'cli-progress';
 import { runPipelineFromRepo } from '../core/ingestion/pipeline.js';
-import { openDb, closeDb, resetDb, getStats, insertEmbeddingsBatch, countEmbeddings } from '../core/db/adapter.js';
+import { openDb, closeDb, resetDb, getStats, insertEmbeddingsBatch, countEmbeddings, populateSearchText } from '../core/db/adapter.js';
 import { loadGraphToDb } from '../core/db/graph-loader.js';
 import { stitchRoutes } from '../core/ingestion/route-stitcher.js';
 import { toNodeId } from '../core/db/schema.js';
@@ -200,6 +200,10 @@ export const analyzeCommand = async (inputPath, options) => {
     const dbWarnings = dbResult.warnings;
     // Phase 2.5: HTTP route stitching (post-DB-load, needs content field)
     stitchRoutes(db);
+    // Phase 2.6: Populate searchText for BM25 concept matching
+    // Uses first comment + callers + module — must run after edges are loaded
+    updateBar(84, 'Building search index...');
+    populateSearchText(db);
     // Phase 3: FTS (85-90%)
     // FTS5 is auto-created by schema triggers — no manual index creation needed
     updateBar(85, 'Search indexes ready');

package/dist/core/db/adapter.d.ts

@@ -90,6 +90,12 @@ export declare function getStats(db: Database.Database): {
 export declare function insertNodesBatch(db: Database.Database, nodes: readonly NodeInsert[]): void;
 /** Batch insert edges in a single transaction. */
 export declare function insertEdgesBatch(db: Database.Database, edges: readonly EdgeInsert[]): void;
+/**
+ * Populate the searchText column for all nodes with semantic summaries.
+ * Uses first comment + callers + module to enable BM25 concept matching.
+ * Call AFTER edges are loaded (needs CALLS and MEMBER_OF edges).
+ */
+export declare function populateSearchText(db: Database.Database): void;
 /** Batch insert embeddings in a single transaction. */
 export declare function insertEmbeddingsBatch(db: Database.Database, items: readonly {
     nodeId: NodeId;

package/dist/core/db/adapter.js

@@ -145,12 +145,12 @@ const INSERT_NODE_SQL = `
     id, label, name, filePath, startLine, endLine, isExported, content, description,
     heuristicLabel, cohesion, symbolCount, keywords, enrichedBy,
     processType, stepCount, communities, entryPointId, terminalId,
-    parameterCount, returnType, nameExpanded
+    parameterCount, returnType, nameExpanded, searchText
   ) VALUES (
     @id, @label, @name, @filePath, @startLine, @endLine, @isExported, @content, @description,
     @heuristicLabel, @cohesion, @symbolCount, @keywords, @enrichedBy,
     @processType, @stepCount, @communities, @entryPointId, @terminalId,
-    @parameterCount, @returnType, @nameExpanded
+    @parameterCount, @returnType, @nameExpanded, @searchText
   )
 `;
 /** Insert or replace a node. Automatically expands name for FTS natural language matching. */
@@ -178,6 +178,7 @@ export function insertNode(db, node) {
         parameterCount: node.parameterCount ?? null,
         returnType: node.returnType ?? null,
         nameExpanded: node.nameExpanded ?? expandIdentifier(node.name ?? ''),
+        searchText: node.searchText ?? '',
     });
 }
 /** Get a node by ID. Returns undefined if not found. */
@@ -373,6 +374,7 @@ export function insertNodesBatch(db, nodes) {
                 terminalId: node.terminalId ?? null, parameterCount: node.parameterCount ?? null,
                 returnType: node.returnType ?? null,
                 nameExpanded: node.nameExpanded ?? expandIdentifier(node.name ?? ''),
+                searchText: node.searchText ?? '',
             });
         }
     });
@@ -392,6 +394,104 @@ export function insertEdgesBatch(db, edges) {
     });
     txn(edges);
 }
+/**
+ * Populate the searchText column for all nodes with semantic summaries.
+ * Uses first comment + callers + module to enable BM25 concept matching.
+ * Call AFTER edges are loaded (needs CALLS and MEMBER_OF edges).
+ */
+export function populateSearchText(db) {
+    // Extract first comment from content
+    function extractComment(content) {
+        if (!content)
+            return '';
+        const lines = content.split('\n');
+        const out = [];
+        let inBlock = false;
+        for (const l of lines) {
+            const t = l.trim();
+            if (t.startsWith('/**') || t.startsWith('/*')) {
+                inBlock = true;
+                const inner = t.replace(/^\/\*\*?\s*/, '').replace(/\*\/\s*$/, '').trim();
+                if (inner && !inner.startsWith('@'))
+                    out.push(inner);
+                if (t.includes('*/'))
+                    inBlock = false;
+                continue;
+            }
+            if (inBlock) {
+                if (t.includes('*/')) {
+                    inBlock = false;
+                    continue;
+                }
+                const inner = t.replace(/^\*\s?/, '').trim();
+                if (inner && !inner.startsWith('@'))
+                    out.push(inner);
+                if (out.length >= 3)
+                    break;
+                continue;
+            }
+            if (t.startsWith('//')) {
+                const inner = t.slice(2).trim();
+                if (inner)
+                    out.push(inner);
+                if (out.length >= 3)
+                    break;
+                continue;
+            }
+            if (out.length > 0)
+                break;
+        }
+        return out.join(' ');
+    }
+    const nodes = db.prepare("SELECT id, name, nameExpanded, content FROM nodes WHERE label IN ('Function','Class','Method','Interface','Const','TypeAlias','Enum')").all();
+    if (nodes.length === 0)
+        return;
+    // Batch fetch callers + module
+    const callerMap = new Map();
+    const moduleMap = new Map();
+    const ids = nodes.map(n => n.id);
+    for (let i = 0; i < ids.length; i += 900) {
+        const chunk = ids.slice(i, i + 900);
+        const ph = chunk.map(() => '?').join(',');
+        const callerRows = db.prepare(`SELECT e.targetId AS nid, n.name FROM edges e JOIN nodes n ON n.id = e.sourceId WHERE e.targetId IN (${ph}) AND e.type = 'CALLS' AND e.confidence >= 0.7`).all(...chunk);
+        for (const r of callerRows) {
+            if (!callerMap.has(r.nid))
+                callerMap.set(r.nid, []);
+            callerMap.get(r.nid).push(r.name);
+        }
+        const modRows = db.prepare(`SELECT e.sourceId AS nid, c.heuristicLabel AS module FROM edges e JOIN nodes c ON c.id = e.targetId WHERE e.sourceId IN (${ph}) AND e.type = 'MEMBER_OF' AND c.label = 'Community'`).all(...chunk);
+        for (const r of modRows)
+            moduleMap.set(r.nid, r.module);
+    }
+    // Build searchText and update
+    // Drop FTS triggers temporarily to avoid column-count issues during bulk update,
+    // then rebuild the FTS index in one pass (faster than per-row trigger updates)
+    db.exec("DROP TRIGGER IF EXISTS nodes_fts_au");
+    const txn = db.transaction(() => {
+        for (const node of nodes) {
+            const parts = [];
+            if (node.nameExpanded)
+                parts.push(node.nameExpanded);
+            const comment = extractComment(node.content);
+            if (comment)
+                parts.push(comment);
+            const callers = callerMap.get(node.id)?.slice(0, 5);
+            if (callers && callers.length > 0)
+                parts.push(callers.map(c => expandIdentifier(c)).join(' '));
+            const mod = moduleMap.get(node.id);
+            if (mod)
+                parts.push(mod);
+            db.prepare('UPDATE nodes SET searchText = ? WHERE id = ?').run(parts.join(' | '), node.id);
+        }
+    });
+    txn();
+    // Rebuild FTS index from scratch and recreate the trigger
+    db.exec("INSERT INTO nodes_fts(nodes_fts) VALUES('rebuild')");
+    db.exec(`CREATE TRIGGER IF NOT EXISTS nodes_fts_au AFTER UPDATE ON nodes BEGIN
+    INSERT INTO nodes_fts(nodes_fts, rowid, name, nameExpanded, searchText, filePath, content) VALUES ('delete', old.rowid, old.name, old.nameExpanded, old.searchText, old.filePath, old.content);
+    INSERT INTO nodes_fts(nodes_fts, rowid, name, nameExpanded, searchText, filePath, content) VALUES (new.rowid, new.name, new.nameExpanded, new.searchText, new.filePath, new.content);
+  END`);
+}
 /** Batch insert embeddings in a single transaction. */
 export function insertEmbeddingsBatch(db, items) {
     const stmt = db.prepare('INSERT OR REPLACE INTO embeddings (nodeId, embedding, textHash) VALUES (?, ?, ?)');

package/dist/core/db/schema.d.ts

@@ -49,6 +49,7 @@ export interface NodeRow {
     readonly parameterCount: number | null;
     readonly returnType: string | null;
     readonly nameExpanded: string;
+    readonly searchText: string;
 }
 /** An edge row as stored in the `edges` table */
 export interface EdgeRow {
@@ -91,6 +92,7 @@ export interface NodeInsert {
     readonly parameterCount?: number | null;
     readonly returnType?: string | null;
     readonly nameExpanded?: string;
+    readonly searchText?: string;
 }
 /** Fields required to insert an edge */
 export interface EdgeInsert {
@@ -105,4 +107,4 @@ export interface EdgeInsert {
 }
 /** Legacy edge table name constant (kept for compatibility) */
 export declare const REL_TABLE_NAME = "CodeRelation";
-export declare const SCHEMA_SQL = "\n-- Nodes: unified table for all code elements\nCREATE TABLE IF NOT EXISTS nodes (\n  id TEXT PRIMARY KEY,\n  label TEXT NOT NULL,\n  name TEXT NOT NULL DEFAULT '',\n  filePath TEXT NOT NULL DEFAULT '',\n  startLine INTEGER,\n  endLine INTEGER,\n  isExported INTEGER,\n  content TEXT NOT NULL DEFAULT '',\n  description TEXT NOT NULL DEFAULT '',\n  heuristicLabel TEXT,\n  cohesion REAL,\n  symbolCount INTEGER,\n  keywords TEXT,\n  enrichedBy TEXT,\n  processType TEXT,\n  stepCount INTEGER,\n  communities TEXT,\n  entryPointId TEXT,\n  terminalId TEXT,\n  parameterCount INTEGER,\n  returnType TEXT,\n  nameExpanded TEXT DEFAULT ''\n);\n\nCREATE INDEX IF NOT EXISTS idx_nodes_label ON nodes(label);\nCREATE INDEX IF NOT EXISTS idx_nodes_name ON nodes(name);\nCREATE INDEX IF NOT EXISTS idx_nodes_filePath ON nodes(filePath);\nCREATE INDEX IF NOT EXISTS idx_nodes_label_name ON nodes(label, name);\nCREATE INDEX IF NOT EXISTS idx_nodes_filePath_lines ON nodes(filePath, startLine, endLine);\n\n-- Edges: single table for all relationships\nCREATE TABLE IF NOT EXISTS edges (\n  id TEXT PRIMARY KEY,\n  sourceId TEXT NOT NULL,\n  targetId TEXT NOT NULL,\n  type TEXT NOT NULL,\n  confidence REAL NOT NULL DEFAULT 1.0,\n  reason TEXT NOT NULL DEFAULT '',\n  step INTEGER NOT NULL DEFAULT 0,\n  callLine INTEGER\n);\n\nCREATE INDEX IF NOT EXISTS idx_edges_sourceId ON edges(sourceId);\nCREATE INDEX IF NOT EXISTS idx_edges_targetId ON edges(targetId);\nCREATE INDEX IF NOT EXISTS idx_edges_type ON edges(type);\nCREATE INDEX IF NOT EXISTS idx_edges_source_type ON edges(sourceId, type);\nCREATE INDEX IF NOT EXISTS idx_edges_target_type ON edges(targetId, type);\n\n-- Embeddings: vector storage\nCREATE TABLE IF NOT EXISTS embeddings (\n  nodeId TEXT PRIMARY KEY,\n  embedding BLOB NOT NULL,\n  textHash TEXT\n);\n\n-- FTS5 virtual table (auto-updated via triggers)\nCREATE VIRTUAL TABLE IF NOT EXISTS nodes_fts USING fts5(\n  name,\n  nameExpanded,\n  filePath,\n  content,\n  content='nodes',\n  content_rowid='rowid'\n);\n\nCREATE TRIGGER IF NOT EXISTS nodes_fts_ai AFTER INSERT ON nodes BEGIN\n  INSERT INTO nodes_fts(rowid, name, nameExpanded, filePath, content) VALUES (new.rowid, new.name, new.nameExpanded, new.filePath, new.content);\nEND;\nCREATE TRIGGER IF NOT EXISTS nodes_fts_ad AFTER DELETE ON nodes BEGIN\n  INSERT INTO nodes_fts(nodes_fts, rowid, name, nameExpanded, filePath, content) VALUES ('delete', old.rowid, old.name, old.nameExpanded, old.filePath, old.content);\nEND;\nCREATE TRIGGER IF NOT EXISTS nodes_fts_au AFTER UPDATE ON nodes BEGIN\n  INSERT INTO nodes_fts(nodes_fts, rowid, name, nameExpanded, filePath, content) VALUES ('delete', old.rowid, old.name, old.nameExpanded, old.filePath, old.content);\n  INSERT INTO nodes_fts(nodes_fts, rowid, name, nameExpanded, filePath, content) VALUES (new.rowid, new.name, new.nameExpanded, new.filePath, new.content);\nEND;\n";
+export declare const SCHEMA_SQL = "\n-- Nodes: unified table for all code elements\nCREATE TABLE IF NOT EXISTS nodes (\n  id TEXT PRIMARY KEY,\n  label TEXT NOT NULL,\n  name TEXT NOT NULL DEFAULT '',\n  filePath TEXT NOT NULL DEFAULT '',\n  startLine INTEGER,\n  endLine INTEGER,\n  isExported INTEGER,\n  content TEXT NOT NULL DEFAULT '',\n  description TEXT NOT NULL DEFAULT '',\n  heuristicLabel TEXT,\n  cohesion REAL,\n  symbolCount INTEGER,\n  keywords TEXT,\n  enrichedBy TEXT,\n  processType TEXT,\n  stepCount INTEGER,\n  communities TEXT,\n  entryPointId TEXT,\n  terminalId TEXT,\n  parameterCount INTEGER,\n  returnType TEXT,\n  nameExpanded TEXT DEFAULT '',\n  searchText TEXT DEFAULT ''\n);\n\nCREATE INDEX IF NOT EXISTS idx_nodes_label ON nodes(label);\nCREATE INDEX IF NOT EXISTS idx_nodes_name ON nodes(name);\nCREATE INDEX IF NOT EXISTS idx_nodes_filePath ON nodes(filePath);\nCREATE INDEX IF NOT EXISTS idx_nodes_label_name ON nodes(label, name);\nCREATE INDEX IF NOT EXISTS idx_nodes_filePath_lines ON nodes(filePath, startLine, endLine);\n\n-- Edges: single table for all relationships\nCREATE TABLE IF NOT EXISTS edges (\n  id TEXT PRIMARY KEY,\n  sourceId TEXT NOT NULL,\n  targetId TEXT NOT NULL,\n  type TEXT NOT NULL,\n  confidence REAL NOT NULL DEFAULT 1.0,\n  reason TEXT NOT NULL DEFAULT '',\n  step INTEGER NOT NULL DEFAULT 0,\n  callLine INTEGER\n);\n\nCREATE INDEX IF NOT EXISTS idx_edges_sourceId ON edges(sourceId);\nCREATE INDEX IF NOT EXISTS idx_edges_targetId ON edges(targetId);\nCREATE INDEX IF NOT EXISTS idx_edges_type ON edges(type);\nCREATE INDEX IF NOT EXISTS idx_edges_source_type ON edges(sourceId, type);\nCREATE INDEX IF NOT EXISTS idx_edges_target_type ON edges(targetId, type);\n\n-- Embeddings: vector storage\nCREATE TABLE IF NOT EXISTS embeddings (\n  nodeId TEXT PRIMARY KEY,\n  embedding BLOB NOT NULL,\n  textHash TEXT\n);\n\n-- FTS5 virtual table (auto-updated via triggers)\nCREATE VIRTUAL TABLE IF NOT EXISTS nodes_fts USING fts5(\n  name,\n  nameExpanded,\n  searchText,\n  filePath,\n  content,\n  content='nodes',\n  content_rowid='rowid'\n);\n\nCREATE TRIGGER IF NOT EXISTS nodes_fts_ai AFTER INSERT ON nodes BEGIN\n  INSERT INTO nodes_fts(rowid, name, nameExpanded, searchText, filePath, content) VALUES (new.rowid, new.name, new.nameExpanded, new.searchText, new.filePath, new.content);\nEND;\nCREATE TRIGGER IF NOT EXISTS nodes_fts_ad AFTER DELETE ON nodes BEGIN\n  INSERT INTO nodes_fts(nodes_fts, rowid, name, nameExpanded, searchText, filePath, content) VALUES ('delete', old.rowid, old.name, old.nameExpanded, old.searchText, old.filePath, old.content);\nEND;\nCREATE TRIGGER IF NOT EXISTS nodes_fts_au AFTER UPDATE ON nodes BEGIN\n  INSERT INTO nodes_fts(nodes_fts, rowid, name, nameExpanded, searchText, filePath, content) VALUES ('delete', old.rowid, old.name, old.nameExpanded, old.searchText, old.filePath, old.content);\n  INSERT INTO nodes_fts(nodes_fts, rowid, name, nameExpanded, searchText, filePath, content) VALUES (new.rowid, new.name, new.nameExpanded, new.searchText, new.filePath, new.content);\nEND;\n";

package/dist/core/db/schema.js

@@ -79,7 +79,8 @@ CREATE TABLE IF NOT EXISTS nodes (
   terminalId TEXT,
   parameterCount INTEGER,
   returnType TEXT,
-  nameExpanded TEXT DEFAULT ''
+  nameExpanded TEXT DEFAULT '',
+  searchText TEXT DEFAULT ''
 );
 
 CREATE INDEX IF NOT EXISTS idx_nodes_label ON nodes(label);
@@ -117,6 +118,7 @@ CREATE TABLE IF NOT EXISTS embeddings (
 CREATE VIRTUAL TABLE IF NOT EXISTS nodes_fts USING fts5(
   name,
   nameExpanded,
+  searchText,
   filePath,
   content,
   content='nodes',
@@ -124,13 +126,13 @@ CREATE VIRTUAL TABLE IF NOT EXISTS nodes_fts USING fts5(
 );
 
 CREATE TRIGGER IF NOT EXISTS nodes_fts_ai AFTER INSERT ON nodes BEGIN
-  INSERT INTO nodes_fts(rowid, name, nameExpanded, filePath, content) VALUES (new.rowid, new.name, new.nameExpanded, new.filePath, new.content);
+  INSERT INTO nodes_fts(rowid, name, nameExpanded, searchText, filePath, content) VALUES (new.rowid, new.name, new.nameExpanded, new.searchText, new.filePath, new.content);
 END;
 CREATE TRIGGER IF NOT EXISTS nodes_fts_ad AFTER DELETE ON nodes BEGIN
-  INSERT INTO nodes_fts(nodes_fts, rowid, name, nameExpanded, filePath, content) VALUES ('delete', old.rowid, old.name, old.nameExpanded, old.filePath, old.content);
+  INSERT INTO nodes_fts(nodes_fts, rowid, name, nameExpanded, searchText, filePath, content) VALUES ('delete', old.rowid, old.name, old.nameExpanded, old.searchText, old.filePath, old.content);
 END;
 CREATE TRIGGER IF NOT EXISTS nodes_fts_au AFTER UPDATE ON nodes BEGIN
-  INSERT INTO nodes_fts(nodes_fts, rowid, name, nameExpanded, filePath, content) VALUES ('delete', old.rowid, old.name, old.nameExpanded, old.filePath, old.content);
-  INSERT INTO nodes_fts(nodes_fts, rowid, name, nameExpanded, filePath, content) VALUES (new.rowid, new.name, new.nameExpanded, new.filePath, new.content);
+  INSERT INTO nodes_fts(nodes_fts, rowid, name, nameExpanded, searchText, filePath, content) VALUES ('delete', old.rowid, old.name, old.nameExpanded, old.searchText, old.filePath, old.content);
+  INSERT INTO nodes_fts(nodes_fts, rowid, name, nameExpanded, searchText, filePath, content) VALUES (new.rowid, new.name, new.nameExpanded, new.searchText, new.filePath, new.content);
 END;
 `;

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

@@ -27,7 +27,7 @@ import mlx.core as mx
 import mlx.nn as nn
 from tokenizers import Tokenizer
 
-MODEL_DIR = os.path.dirname(os.path.abspath(__file__)) + "/jina-v5-small-mlx"
+MODEL_DIR = os.path.dirname(os.path.abspath(__file__)) + "/jina-code-0.5b-mlx"
 
 
 
@@ -35,27 +35,54 @@ MODEL_DIR = os.path.dirname(os.path.abspath(__file__)) + "/jina-v5-small-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"
         )
 
 
@@ -64,7 +91,10 @@ def load_model():
     ensure_model_downloaded()
 
     sys.path.insert(0, MODEL_DIR)
-    from model import JinaEmbeddingModel
+    import importlib
+    model_module = importlib.import_module("model")
+    # Support both model class names (v5 = JinaEmbeddingModel, code-0.5b = JinaCodeEmbeddingModel)
+    JinaEmbeddingModel = getattr(model_module, "JinaEmbeddingModel", None) or getattr(model_module, "JinaCodeEmbeddingModel")
 
     with open(os.path.join(MODEL_DIR, "config.json")) as f:
         config = json.load(f)
@@ -99,8 +129,17 @@ def embed_tiered(model, tokenizer, texts, task_type="retrieval.passage", truncat
     if not texts:
         return []
 
-    # Add task prefix
-    prefix_map = {"retrieval.query": "Query: ", "retrieval.passage": "Document: "}
+    # Add task prefix — auto-detect based on model type
+    # v5 (Qwen3): "Query: " / "Document: "
+    # code-0.5b (Qwen2): "Find the most relevant code snippet...\n" / "Candidate code snippet:\n"
+    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(task_type, "")
     prefixed = [prefix + t for t in texts] if prefix else texts
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@zuvia-software-solutions/code-mapper",
-  "version": "2.2.3",
+  "version": "2.3.1",
   "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