@zuvia-software-solutions/code-mapper 2.4.0 → 2.5.0

package/models/mlx-embedder.py

@@ -1,604 +0,0 @@
-#!/usr/bin/env python3
-"""
-MLX-accelerated code embedder for Apple Silicon.
-
-TWO MODES:
-1. Batch mode (main use): reads nodes directly from SQLite, embeds, writes back.
-   No IPC overhead — everything happens in one process.
-   Usage: python3 mlx-embedder.py batch <db_path> [--dims 256] [--max-tokens 2048]
-
-2. Interactive mode (for MCP query embedding): reads JSON from stdin.
-   Usage: python3 mlx-embedder.py [interactive]
-
-Model: Jina Embeddings v5 Text Small Retrieval (677M params, Qwen3-0.6B backbone)
-Optimized with int4 quantization (Linear) + int6 quantization (Embedding).
-"""
-
-import sys
-import os
-import json
-import time
-import struct
-import hashlib
-
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
-import mlx.core as mx
-import mlx.nn as nn
-from tokenizers import Tokenizer
-
-MODEL_DIR = os.path.dirname(os.path.abspath(__file__)) + "/jina-code-0.5b-mlx"
-
-
-
-
-
-
-def ensure_model_downloaded():
-    """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
-
-    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-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)
-
-        # 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-code-embeddings-0.5b"
-        )
-
-
-def load_model():
-    """Load model, quantize for speed. Auto-downloads weights on first use."""
-    ensure_model_downloaded()
-
-    sys.path.insert(0, MODEL_DIR)
-    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)
-
-    model = JinaEmbeddingModel(config)
-    weights = mx.load(os.path.join(MODEL_DIR, "model.safetensors"))
-    model.load_weights(list(weights.items()))
-
-    nn.quantize(model.model, bits=4, group_size=64,
-                class_predicate=lambda _, m: isinstance(m, nn.Linear))
-    nn.quantize(model.model, bits=6, group_size=64,
-                class_predicate=lambda _, m: isinstance(m, nn.Embedding))
-    mx.eval(model.parameters())
-
-    tokenizer = Tokenizer.from_file(os.path.join(MODEL_DIR, "tokenizer.json"))
-    return model, tokenizer
-
-
-def get_batch_size_for_tokens(token_count):
-    """Optimal batch size based on actual token count."""
-    if token_count <= 64:   return 256
-    if token_count <= 128:  return 128
-    if token_count <= 256:  return 64
-    if token_count <= 512:  return 32
-    if token_count <= 1024: return 16
-    return 8
-
-
-def embed_tiered(model, tokenizer, texts, task_type="retrieval.passage", truncate_dim=256, max_tokens=2048):
-    """Embed texts with token-aware batching. Tokenizes first, batches by token count.
-    Returns embeddings in the ORIGINAL input order."""
-    if not texts:
-        return []
-
-    # 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
-
-    # Tokenize everything in one call (fast — Rust HF tokenizer)
-    encodings = tokenizer.encode_batch(prefixed)
-
-    # Sort by token length for minimal padding
-    indexed = sorted(range(len(texts)), key=lambda i: len(encodings[i].ids))
-
-    all_embeddings = [None] * len(texts)
-    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)
-
-        embs = model(mx.array(input_ids), mx.array(attention_mask))
-        if truncate_dim and truncate_dim < embs.shape[1]:
-            embs = embs[:, :truncate_dim]
-        norms = mx.linalg.norm(embs, axis=1, keepdims=True)
-        embs = embs / norms
-        mx.eval(embs)
-
-        emb_list = embs.tolist()
-        for j, orig_idx in enumerate(batch_indices):
-            all_embeddings[orig_idx] = emb_list[j]
-
-    return all_embeddings
-
-
-def float_list_to_blob(floats):
-    """Convert list of floats to a binary blob (Float32Array compatible)."""
-    return struct.pack(f'{len(floats)}f', *floats)
-
-
-def md5(text):
-    return hashlib.md5(text.encode()).hexdigest()
-
-
-# =========================================================================
-# BATCH MODE — read from SQLite, embed, write back. Zero IPC.
-# =========================================================================
-
-def batch_mode(db_path, dims=256, max_tokens=2048):
-    import sqlite3
-
-    t0_total = time.time()
-
-    # Load model
-    print(json.dumps({"phase": "loading", "message": "Loading MLX model..."}), flush=True)
-    model, tokenizer = load_model()
-    load_ms = int((time.time() - t0_total) * 1000)
-    print(json.dumps({"phase": "loaded", "load_ms": load_ms, "device": str(mx.default_device())}), flush=True)
-
-    # Open database
-    db = sqlite3.connect(db_path)
-    db.execute("PRAGMA journal_mode=WAL")
-    db.execute("PRAGMA synchronous=NORMAL")
-
-    # Ensure textHash column exists (migration)
-    try:
-        db.execute("SELECT textHash FROM embeddings LIMIT 0")
-    except sqlite3.OperationalError:
-        db.execute("ALTER TABLE embeddings ADD COLUMN textHash TEXT")
-
-    # Query embeddable nodes — skip test/fixture files (BM25 covers them)
-    labels = ('Function', 'Class', 'Method', 'Interface', 'Const', 'Enum', 'TypeAlias', 'Namespace', 'Module', 'Struct')
-    placeholders = ','.join('?' * len(labels))
-    all_rows = db.execute(
-        f"SELECT id, name, label, filePath, content, startLine, endLine, nameExpanded FROM nodes WHERE label IN ({placeholders})",
-        labels
-    ).fetchall()
-
-    # Filter out test files — they're searchable via BM25 keyword matching
-    test_patterns = ('/test/', '/tests/', '/spec/', '/fixtures/', '/__tests__/', '/__mocks__/',
-                     '.test.', '.spec.', '_test.', '_spec.')
-    rows = [r for r in all_rows if not any(p in (r[3] or '') for p in test_patterns)]
-    skipped_tests = len(all_rows) - len(rows)
-
-    print(json.dumps({"phase": "queried", "nodes": len(rows), "skipped_tests": skipped_tests}), flush=True)
-
-    if not rows:
-        print(json.dumps({"phase": "done", "embedded": 0, "skipped": 0, "ms": 0}), flush=True)
-        db.close()
-        return
-
-    # Fetch graph context (callers, callees, module) for richer embedding text
-    node_ids = [r[0] for r in rows]
-    id_set = set(node_ids)
-
-    # Batch fetch callers
-    caller_map = {}
-    callee_map = {}
-    module_map = {}
-
-    # Chunk the IN clause to avoid SQLite variable limits
-    CHUNK = 500
-    for ci in range(0, len(node_ids), CHUNK):
-        chunk_ids = node_ids[ci:ci+CHUNK]
-        ph = ','.join('?' * len(chunk_ids))
-
-        for row in db.execute(f"SELECT e.targetId, 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 LIMIT {len(chunk_ids)*3}", chunk_ids):
-            caller_map.setdefault(row[0], []).append(row[1])
-
-        for row in db.execute(f"SELECT e.sourceId, n.name FROM edges e JOIN nodes n ON n.id = e.targetId WHERE e.sourceId IN ({ph}) AND e.type = 'CALLS' AND e.confidence >= 0.7 LIMIT {len(chunk_ids)*3}", chunk_ids):
-            callee_map.setdefault(row[0], []).append(row[1])
-
-        for row in db.execute(f"SELECT e.sourceId, c.heuristicLabel 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' LIMIT {len(chunk_ids)}", chunk_ids):
-            module_map[row[0]] = row[1]
-
-    # Batch fetch import names per file
-    import_map = {}
-    for ci in range(0, len(node_ids), CHUNK):
-        chunk_ids = node_ids[ci:ci+CHUNK]
-        ph = ','.join('?' * len(chunk_ids))
-        # Get unique file paths for these nodes
-        file_paths = [r[3] for r in rows if r[0] in set(chunk_ids)]
-        unique_files = list(set(file_paths))
-        if unique_files:
-            fph = ','.join('?' * len(unique_files))
-            for row in db.execute(
-                f"SELECT DISTINCT n.filePath, tn.name FROM nodes n JOIN edges e ON e.sourceId = n.id AND e.type = 'IMPORTS' JOIN nodes tn ON tn.id = e.targetId WHERE n.filePath IN ({fph}) LIMIT {len(unique_files)*10}",
-                unique_files
-            ):
-                import_map.setdefault(row[0], []).append(row[1])
-
-    print(json.dumps({"phase": "context", "with_callers": len(caller_map), "with_module": len(module_map), "with_imports": len(import_map)}), flush=True)
-
-    # Get existing text hashes for skip detection
-    existing_hashes = {}
-    for row in db.execute("SELECT nodeId, textHash FROM embeddings WHERE textHash IS NOT NULL"):
-        existing_hashes[row[0]] = row[1]
-
-    # Generate embedding texts + hashes
-    # Optimized: semantic summary (name + comment + signature + context)
-    # instead of raw code dump. 55% fewer tokens, equal search quality.
-    to_embed = []  # (node_id, text, hash)
-    skipped = 0
-
-    def extract_first_comment(content):
-        """Extract JSDoc/comment as natural language description (max 3 lines)."""
-        if not content:
-            return ""
-        lines = content.split("\n")
-        comment_lines = []
-        in_block = False
-        for l in lines:
-            t = l.strip()
-            if t.startswith("/**") or t.startswith("/*"):
-                in_block = True
-                inner = t.lstrip("/").lstrip("*").strip().rstrip("*/").strip()
-                if inner and not inner.startswith("@"):
-                    comment_lines.append(inner)
-                if "*/" in t:
-                    in_block = False
-                continue
-            if in_block:
-                if "*/" in t:
-                    in_block = False
-                    continue
-                inner = t.lstrip("*").strip()
-                if inner and not inner.startswith("@"):
-                    comment_lines.append(inner)
-                if len(comment_lines) >= 3:
-                    break
-                continue
-            if t.startswith("//"):
-                inner = t[2:].strip()
-                if inner:
-                    comment_lines.append(inner)
-                if len(comment_lines) >= 3:
-                    break
-                continue
-            if t.startswith("#") and not t.startswith("#!"):
-                inner = t[1:].strip()
-                if inner:
-                    comment_lines.append(inner)
-                if len(comment_lines) >= 3:
-                    break
-                continue
-            if comment_lines:
-                break
-        return " ".join(comment_lines)
-
-    def extract_signature(content, label):
-        """Extract code signature without full body."""
-        if not content:
-            return ""
-        lines = content.split("\n")
-        if label == "Interface":
-            return "\n".join(lines[:30]).strip() if len(lines) <= 30 else "\n".join(lines[:30]) + "\n  // ..."
-        if label == "Class":
-            sigs = []
-            for l in lines[:60]:
-                t = l.strip()
-                if not t or t.startswith("//") or t.startswith("*") or t.startswith("/*"):
-                    continue
-                if any(kw in t for kw in ("class ", "private ", "public ", "protected ", "readonly ", "static ", "abstract ")):
-                    sigs.append(t)
-                if len(sigs) >= 20:
-                    break
-            return "\n".join(sigs)
-        return "\n".join(lines[:min(8, len(lines))]).strip()
-
-    for row in rows:
-        nid, name, label, filePath, content, startLine, endLine, nameExpanded = row
-        content = content or ""
-        file_name = filePath.rsplit('/', 1)[-1] if filePath else ""
-
-        # Build semantic embedding text
-        parts = [f"{label}: {name}"]
-
-        # nameExpanded: natural language bridge (e.g. "checkStaleness" → "check staleness")
-        if nameExpanded and nameExpanded != name.lower():
-            parts.append(nameExpanded)
-
-        # First comment as natural language description
-        comment = extract_first_comment(content)
-        if comment:
-            parts.append(comment)
-
-        # Directory context
-        dir_parts = filePath.rsplit('/', 2)
-        dir_context = '/'.join(dir_parts[:-1])[-40:] if '/' in filePath else ''
-
-        # File + module location
-        loc = f"File: {file_name}"
-        if dir_context:
-            loc += f" in {dir_context}"
-        module = module_map.get(nid, "")
-        if module:
-            loc += f" | Module: {module}"
-        parts.append(loc)
-
-        # Import context
-        file_imports = import_map.get(filePath, [])[:5]
-        if file_imports:
-            parts.append(f"Imports: {', '.join(file_imports)}")
-
-        # Graph context
-        callers = caller_map.get(nid, [])[:5]
-        callees = callee_map.get(nid, [])[:5]
-        if callers:
-            parts.append(f"Called by: {', '.join(callers)}")
-        if callees:
-            parts.append(f"Calls: {', '.join(callees)}")
-
-        # Code signature (not full body)
-        sig = extract_signature(content, label)
-        if sig:
-            parts.extend(["", sig])
-
-        text = '\n'.join(parts)
-        text_hash = md5(text)
-
-        # Skip if hash unchanged
-        if existing_hashes.get(nid) == text_hash:
-            skipped += 1
-            continue
-
-        to_embed.append((nid, text, text_hash))
-
-    print(json.dumps({"phase": "prepared", "to_embed": len(to_embed), "skipped": skipped}), flush=True)
-
-    if not to_embed:
-        print(json.dumps({"phase": "done", "embedded": 0, "skipped": skipped, "ms": int((time.time() - t0_total) * 1000)}), flush=True)
-        db.close()
-        return
-
-    # Deduplicate — embed unique texts only, copy vectors to duplicates.
-    # Identical embedding texts produce identical vectors; no quality loss.
-    unique_by_hash = {}  # text_hash -> { text, node_ids: [(nid, text_hash)] }
-    for nid, text, text_hash in to_embed:
-        if text_hash in unique_by_hash:
-            unique_by_hash[text_hash]["node_ids"].append((nid, text_hash))
-        else:
-            unique_by_hash[text_hash] = {"text": text, "node_ids": [(nid, text_hash)]}
-    unique_texts = [v["text"] for v in unique_by_hash.values()]
-    deduped = len(to_embed) - len(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()
-    unique_entries = list(unique_by_hash.values())
-
-    # 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))
-
-    embedded_count = 0
-    db.execute("BEGIN")
-
-    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")
-    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({
-        "phase": "done",
-        "embedded": len(to_embed),
-        "skipped": skipped,
-        "embed_ms": embed_ms,
-        "write_ms": write_ms,
-        "total_ms": total_ms,
-    }), flush=True)
-
-    db.close()
-
-
-# =========================================================================
-# INTERACTIVE MODE — stdin/stdout JSON for MCP query embedding
-# =========================================================================
-
-def interactive_mode():
-    t0 = time.time()
-    model, tokenizer = load_model()
-    load_ms = int((time.time() - t0) * 1000)
-
-    print(json.dumps({
-        "status": "ready",
-        "model": "jina-v5-text-small-retrieval",
-        "device": str(mx.default_device()),
-        "load_ms": load_ms,
-        "precision": "int4-g64",
-    }), flush=True)
-
-    for line in sys.stdin:
-        line = line.strip()
-        if not line:
-            continue
-
-        try:
-            req = json.loads(line)
-        except json.JSONDecodeError:
-            print(json.dumps({"error": "Invalid JSON"}), flush=True)
-            continue
-
-        if "cmd" in req:
-            if req["cmd"] == "ping":
-                print(json.dumps({"status": "ready"}), flush=True)
-            elif req["cmd"] == "quit":
-                break
-            continue
-
-        texts = req.get("texts", [])
-        prompt_type = req.get("type", "passage")
-        dims = req.get("dims", 256)
-        task_type = "retrieval.query" if prompt_type == "query" else "retrieval.passage"
-
-        t0 = time.time()
-        try:
-            embeddings = embed_tiered(model, tokenizer, texts, task_type, dims)
-            elapsed_ms = int((time.time() - t0) * 1000)
-            print(json.dumps({
-                "embeddings": embeddings,
-                "count": len(embeddings),
-                "dims": dims,
-                "ms": elapsed_ms,
-            }), flush=True)
-        except Exception as e:
-            print(json.dumps({"error": str(e)}), flush=True)
-
-
-# =========================================================================
-# MAIN
-# =========================================================================
-
-if __name__ == "__main__":
-    if len(sys.argv) >= 3 and sys.argv[1] == "batch":
-        db_path = sys.argv[2]
-        dims = 256
-        max_tokens = 2048
-        for i, arg in enumerate(sys.argv[3:], 3):
-            if arg == "--dims" and i + 1 < len(sys.argv):
-                dims = int(sys.argv[i + 1])
-            if arg == "--max-tokens" and i + 1 < len(sys.argv):
-                max_tokens = int(sys.argv[i + 1])
-        batch_mode(db_path, dims, max_tokens)
-    else:
-        interactive_mode()