@zuvia-software-solutions/code-mapper 2.3.1 → 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/mlx-embedder.py

@@ -415,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.1",
+  "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",