grepmax 0.1.0

package/dist/lib/workers/embeddings/colbert.js

@@ -0,0 +1,213 @@
+"use strict";
+var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    var desc = Object.getOwnPropertyDescriptor(m, k);
+    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
+      desc = { enumerable: true, get: function() { return m[k]; } };
+    }
+    Object.defineProperty(o, k2, desc);
+}) : (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    o[k2] = m[k];
+}));
+var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
+    Object.defineProperty(o, "default", { enumerable: true, value: v });
+}) : function(o, v) {
+    o["default"] = v;
+});
+var __importStar = (this && this.__importStar) || (function () {
+    var ownKeys = function(o) {
+        ownKeys = Object.getOwnPropertyNames || function (o) {
+            var ar = [];
+            for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
+            return ar;
+        };
+        return ownKeys(o);
+    };
+    return function (mod) {
+        if (mod && mod.__esModule) return mod;
+        var result = {};
+        if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
+        __setModuleDefault(result, mod);
+        return result;
+    };
+})();
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+    return new (P || (P = Promise))(function (resolve, reject) {
+        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+        step((generator = generator.apply(thisArg, _arguments || [])).next());
+    });
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.ColbertModel = void 0;
+const fs = __importStar(require("node:fs"));
+const path = __importStar(require("node:path"));
+const ort = __importStar(require("onnxruntime-node"));
+const config_1 = require("../../../config");
+const colbert_tokenizer_1 = require("../colbert-tokenizer");
+const CACHE_DIR = config_1.PATHS.models;
+const ONNX_THREADS = 1;
+const LOG_MODELS = process.env.OSGREP_DEBUG_MODELS === "1" ||
+    process.env.OSGREP_DEBUG_MODELS === "true";
+const log = (...args) => {
+    if (LOG_MODELS)
+        console.log(...args);
+};
+class ColbertModel {
+    constructor() {
+        this.session = null;
+        this.tokenizer = null;
+    }
+    load() {
+        return __awaiter(this, void 0, void 0, function* () {
+            if (this.session && this.tokenizer)
+                return;
+            this.tokenizer = new colbert_tokenizer_1.ColBERTTokenizer();
+            const basePath = path.join(CACHE_DIR, config_1.MODEL_IDS.colbert);
+            const onnxDir = path.join(basePath, "onnx");
+            const modelPath = path.join(onnxDir, "model_int8.onnx");
+            if (!fs.existsSync(modelPath)) {
+                throw new Error(`ColBERT ONNX model not found at ${modelPath}`);
+            }
+            yield this.tokenizer.init(basePath);
+            const sessionOptions = {
+                executionProviders: ["cpu"],
+                intraOpNumThreads: ONNX_THREADS,
+                interOpNumThreads: 1,
+                graphOptimizationLevel: "all",
+            };
+            log(`Worker: Loading ColBERT ONNX session from ${modelPath}`);
+            this.session = yield ort.InferenceSession.create(modelPath, sessionOptions);
+            if (!this.session) {
+                throw new Error(`ColBERT ONNX load failed; tried ${modelPath}`);
+            }
+        });
+    }
+    isReady() {
+        return !!(this.session && this.tokenizer);
+    }
+    runBatch(texts, denseVectors, vectorDimensions) {
+        return __awaiter(this, void 0, void 0, function* () {
+            var _a;
+            if (!this.session || !this.tokenizer)
+                return [];
+            const tokenizer = this.tokenizer;
+            const session = this.session;
+            const encodedBatch = yield Promise.all(texts.map((t) => tokenizer.encodeDoc(t)));
+            const maxLen = Math.max(...encodedBatch.map((e) => e.input_ids.length));
+            const batchInputIds = new BigInt64Array(texts.length * maxLen);
+            const batchAttentionMask = new BigInt64Array(texts.length * maxLen);
+            const padId = BigInt(50283);
+            for (let i = 0; i < encodedBatch.length; i++) {
+                const encoded = encodedBatch[i];
+                const offset = i * maxLen;
+                for (let j = 0; j < maxLen; j++) {
+                    if (j < encoded.input_ids.length) {
+                        batchInputIds[offset + j] = encoded.input_ids[j];
+                        batchAttentionMask[offset + j] = encoded.attention_mask[j];
+                    }
+                    else {
+                        batchInputIds[offset + j] = padId;
+                        batchAttentionMask[offset + j] = BigInt(0);
+                    }
+                }
+            }
+            const feeds = {
+                input_ids: new ort.Tensor("int64", batchInputIds, [texts.length, maxLen]),
+                attention_mask: new ort.Tensor("int64", batchAttentionMask, [
+                    texts.length,
+                    maxLen,
+                ]),
+            };
+            const sessionOut = yield session.run(feeds);
+            const outputName = session.outputNames[0];
+            const output = sessionOut[outputName];
+            if (!output) {
+                throw new Error("ColBERT session output missing embeddings tensor");
+            }
+            const data = output.data;
+            const [batch, seq, dim] = output.dims;
+            const results = [];
+            for (let b = 0; b < batch; b++) {
+                const batchOffset = b * seq * dim;
+                const originalLen = encodedBatch[b].input_ids.length;
+                const normalized = new Float32Array(originalLen * dim);
+                let maxVal = 0;
+                for (let s = 0; s < originalLen; s++) {
+                    const offset = batchOffset + s * dim;
+                    let sumSq = 0;
+                    for (let d = 0; d < dim; d++) {
+                        const val = data[offset + d];
+                        sumSq += val * val;
+                    }
+                    const norm = Math.sqrt(sumSq) || 1;
+                    for (let d = 0; d < dim; d++) {
+                        const val = data[offset + d] / norm;
+                        const idx = s * dim + d;
+                        normalized[idx] = val;
+                        if (Math.abs(val) > maxVal)
+                            maxVal = Math.abs(val);
+                    }
+                }
+                if (maxVal === 0)
+                    maxVal = 1;
+                const int8Array = new Int8Array(normalized.length);
+                for (let i = 0; i < normalized.length; i++) {
+                    int8Array[i] = Math.max(-127, Math.min(127, Math.round((normalized[i] / maxVal) * 127)));
+                }
+                const pooled = new Float32Array(dim);
+                const tokenCount = Math.max(1, originalLen);
+                for (let s = 0; s < originalLen; s++) {
+                    const tokenOffset = s * dim;
+                    for (let d = 0; d < dim; d++) {
+                        pooled[d] += normalized[tokenOffset + d];
+                    }
+                }
+                let pooledNorm = 0;
+                for (let d = 0; d < dim; d++) {
+                    pooled[d] /= tokenCount;
+                    pooledNorm += pooled[d] * pooled[d];
+                }
+                pooledNorm = Math.sqrt(pooledNorm) || 1;
+                for (let d = 0; d < dim; d++) {
+                    pooled[d] /= pooledNorm;
+                }
+                results.push({
+                    dense: (_a = denseVectors[b]) !== null && _a !== void 0 ? _a : new Float32Array(vectorDimensions).fill(0),
+                    colbert: int8Array,
+                    scale: maxVal,
+                    pooled_colbert_48d: pooled,
+                    token_ids: Array.from(encodedBatch[b].input_ids, (v) => Number(v)),
+                });
+            }
+            return results;
+        });
+    }
+    encodeQuery(text) {
+        return __awaiter(this, void 0, void 0, function* () {
+            if (!this.tokenizer)
+                throw new Error("ColBERT tokenizer not initialized");
+            const encoded = yield this.tokenizer.encodeQuery(text);
+            return {
+                input_ids: new BigInt64Array(encoded.input_ids),
+                attention_mask: new BigInt64Array(encoded.attention_mask),
+            };
+        });
+    }
+    runSession(feeds) {
+        return __awaiter(this, void 0, void 0, function* () {
+            if (!this.session)
+                throw new Error("ColBERT session not initialized");
+            return this.session.run(feeds);
+        });
+    }
+    getOutputName() {
+        if (!this.session)
+            throw new Error("ColBERT session not initialized");
+        return this.session.outputNames[0];
+    }
+}
+exports.ColbertModel = ColbertModel;

package/dist/lib/workers/embeddings/granite.js

@@ -0,0 +1,180 @@
+"use strict";
+var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    var desc = Object.getOwnPropertyDescriptor(m, k);
+    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
+      desc = { enumerable: true, get: function() { return m[k]; } };
+    }
+    Object.defineProperty(o, k2, desc);
+}) : (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    o[k2] = m[k];
+}));
+var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
+    Object.defineProperty(o, "default", { enumerable: true, value: v });
+}) : function(o, v) {
+    o["default"] = v;
+});
+var __importStar = (this && this.__importStar) || (function () {
+    var ownKeys = function(o) {
+        ownKeys = Object.getOwnPropertyNames || function (o) {
+            var ar = [];
+            for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
+            return ar;
+        };
+        return ownKeys(o);
+    };
+    return function (mod) {
+        if (mod && mod.__esModule) return mod;
+        var result = {};
+        if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
+        __setModuleDefault(result, mod);
+        return result;
+    };
+})();
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+    return new (P || (P = Promise))(function (resolve, reject) {
+        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+        step((generator = generator.apply(thisArg, _arguments || [])).next());
+    });
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.GraniteModel = void 0;
+const fs = __importStar(require("node:fs"));
+const path = __importStar(require("node:path"));
+const transformers_1 = require("@huggingface/transformers");
+const ort = __importStar(require("onnxruntime-node"));
+const config_1 = require("../../../config");
+const CACHE_DIR = config_1.PATHS.models;
+const ONNX_THREADS = 1;
+const LOG_MODELS = process.env.OSGREP_DEBUG_MODELS === "1" ||
+    process.env.OSGREP_DEBUG_MODELS === "true";
+const log = (...args) => {
+    if (LOG_MODELS)
+        console.log(...args);
+};
+class GraniteModel {
+    constructor() {
+        this.session = null;
+        this.tokenizer = null;
+        this.vectorDimensions = config_1.CONFIG.VECTOR_DIM;
+    }
+    resolvePaths() {
+        const basePath = path.join(CACHE_DIR, config_1.MODEL_IDS.embed);
+        const onnxDir = path.join(basePath, "onnx");
+        const candidates = ["model_q4.onnx", "model.onnx"];
+        for (const candidate of candidates) {
+            const candidatePath = path.join(onnxDir, candidate);
+            if (fs.existsSync(candidatePath)) {
+                return { modelPath: candidatePath, tokenizerPath: basePath };
+            }
+        }
+        throw new Error(`Granite ONNX model not found. Looked for ${candidates.join(", ")} in ${onnxDir}`);
+    }
+    load() {
+        return __awaiter(this, void 0, void 0, function* () {
+            if (this.session && this.tokenizer)
+                return;
+            const { modelPath, tokenizerPath } = this.resolvePaths();
+            log(`Worker: Loading Granite ONNX session from ${modelPath}`);
+            this.tokenizer = yield transformers_1.AutoTokenizer.from_pretrained(tokenizerPath);
+            const sessionOptions = {
+                executionProviders: ["cpu"],
+                intraOpNumThreads: ONNX_THREADS,
+                interOpNumThreads: 1,
+                graphOptimizationLevel: "all",
+            };
+            this.session = yield ort.InferenceSession.create(modelPath, sessionOptions);
+        });
+    }
+    isReady() {
+        return !!(this.session && this.tokenizer);
+    }
+    meanPool(hidden, attention, batch, seq, hiddenDim, targetDim) {
+        const vectors = [];
+        const seqFromMask = attention.length / Math.max(1, batch);
+        const usableSeq = Math.min(seq, seqFromMask);
+        const dim = Math.min(hiddenDim, targetDim);
+        for (let b = 0; b < batch; b++) {
+            const sum = new Float32Array(dim);
+            let count = 0;
+            const attOffset = b * seqFromMask;
+            const hiddenOffset = b * seq * hiddenDim;
+            for (let s = 0; s < usableSeq; s++) {
+                if (attention[attOffset + s] > 0) {
+                    count++;
+                    const tokenOffset = hiddenOffset + s * hiddenDim;
+                    for (let d = 0; d < dim; d++) {
+                        sum[d] += hidden[tokenOffset + d];
+                    }
+                }
+            }
+            if (count === 0)
+                count = 1;
+            let norm = 0;
+            for (let d = 0; d < dim; d++) {
+                sum[d] /= count;
+                norm += sum[d] * sum[d];
+            }
+            norm = Math.sqrt(norm) || 1;
+            for (let d = 0; d < dim; d++) {
+                sum[d] /= norm;
+            }
+            if (dim < targetDim) {
+                const padded = new Float32Array(targetDim);
+                padded.set(sum);
+                vectors.push(padded);
+            }
+            else {
+                vectors.push(sum);
+            }
+        }
+        return vectors;
+    }
+    runBatch(texts) {
+        return __awaiter(this, void 0, void 0, function* () {
+            var _a, _b, _c;
+            if (!this.session || !this.tokenizer)
+                return [];
+            const encoded = yield this.tokenizer(texts, {
+                padding: true,
+                truncation: true,
+                max_length: 256,
+            });
+            const inputTensor = encoded.input_ids;
+            const attentionTensor = encoded.attention_mask;
+            const inputIds = inputTensor.data;
+            const attentionMask = attentionTensor.data;
+            const seqLen = (_b = (_a = inputTensor.dims) === null || _a === void 0 ? void 0 : _a[1]) !== null && _b !== void 0 ? _b : Math.max(1, Math.floor(inputIds.length / texts.length));
+            const tokenTypeIdsRaw = encoded.token_type_ids;
+            const tokenTypeIds = tokenTypeIdsRaw &&
+                tokenTypeIdsRaw.data.length === inputIds.length &&
+                tokenTypeIdsRaw.data.length === attentionMask.length
+                ? tokenTypeIdsRaw.data
+                : new BigInt64Array(inputIds.length).fill(BigInt(0));
+            const feeds = {
+                input_ids: new ort.Tensor("int64", inputIds, [texts.length, seqLen]),
+                attention_mask: new ort.Tensor("int64", attentionMask, [
+                    texts.length,
+                    seqLen,
+                ]),
+                token_type_ids: new ort.Tensor("int64", tokenTypeIds, [
+                    texts.length,
+                    seqLen,
+                ]),
+            };
+            const sessionOut = yield this.session.run(feeds);
+            const hidden = (_c = sessionOut.last_hidden_state) !== null && _c !== void 0 ? _c : sessionOut[this.session.outputNames[0]];
+            if (!hidden) {
+                throw new Error("Granite ONNX output missing last_hidden_state");
+            }
+            const hiddenData = hidden.data;
+            const [batch, seq, dim] = hidden.dims;
+            return this.meanPool(hiddenData, attentionMask, batch, seq, dim, this.vectorDimensions);
+        });
+    }
+}
+exports.GraniteModel = GraniteModel;

package/dist/lib/workers/embeddings/mlx-client.js

@@ -0,0 +1,144 @@
+"use strict";
+/**
+ * MLX embedding server HTTP client.
+ * Tries the local MLX GPU server for dense embeddings.
+ * Returns null if the server isn't running — caller falls back to ONNX.
+ */
+var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    var desc = Object.getOwnPropertyDescriptor(m, k);
+    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
+      desc = { enumerable: true, get: function() { return m[k]; } };
+    }
+    Object.defineProperty(o, k2, desc);
+}) : (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    o[k2] = m[k];
+}));
+var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
+    Object.defineProperty(o, "default", { enumerable: true, value: v });
+}) : function(o, v) {
+    o["default"] = v;
+});
+var __importStar = (this && this.__importStar) || (function () {
+    var ownKeys = function(o) {
+        ownKeys = Object.getOwnPropertyNames || function (o) {
+            var ar = [];
+            for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
+            return ar;
+        };
+        return ownKeys(o);
+    };
+    return function (mod) {
+        if (mod && mod.__esModule) return mod;
+        var result = {};
+        if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
+        __setModuleDefault(result, mod);
+        return result;
+    };
+})();
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+    return new (P || (P = Promise))(function (resolve, reject) {
+        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+        step((generator = generator.apply(thisArg, _arguments || [])).next());
+    });
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.mlxEmbed = mlxEmbed;
+exports.resetMlxCache = resetMlxCache;
+const http = __importStar(require("node:http"));
+const MLX_PORT = parseInt(process.env.MLX_EMBED_PORT || "8100", 10);
+const MLX_HOST = "127.0.0.1";
+const MLX_TIMEOUT_MS = 10000;
+const EMBED_MODE = process.env.OSGREP_EMBED_MODE || "auto";
+let mlxAvailable = null;
+let lastCheck = 0;
+const CHECK_INTERVAL_MS = 30000;
+function postJSON(path, body) {
+    return new Promise((resolve) => {
+        const payload = JSON.stringify(body);
+        const req = http.request({
+            hostname: MLX_HOST,
+            port: MLX_PORT,
+            path,
+            method: "POST",
+            headers: {
+                "Content-Type": "application/json",
+                "Content-Length": Buffer.byteLength(payload),
+            },
+            timeout: MLX_TIMEOUT_MS,
+        }, (res) => {
+            const chunks = [];
+            res.on("data", (chunk) => chunks.push(chunk));
+            res.on("end", () => {
+                try {
+                    const data = JSON.parse(Buffer.concat(chunks).toString("utf-8"));
+                    resolve({ ok: res.statusCode === 200, data });
+                }
+                catch (_a) {
+                    resolve({ ok: false });
+                }
+            });
+        });
+        req.on("error", () => resolve({ ok: false }));
+        req.on("timeout", () => {
+            req.destroy();
+            resolve({ ok: false });
+        });
+        req.write(payload);
+        req.end();
+    });
+}
+/**
+ * Check if MLX server is reachable. Caches result for CHECK_INTERVAL_MS.
+ */
+function isMlxUp() {
+    return __awaiter(this, void 0, void 0, function* () {
+        const now = Date.now();
+        if (mlxAvailable !== null && now - lastCheck < CHECK_INTERVAL_MS) {
+            return mlxAvailable;
+        }
+        const result = yield new Promise((resolve) => {
+            const req = http.get({ hostname: MLX_HOST, port: MLX_PORT, path: "/health", timeout: 2000 }, (res) => {
+                res.resume();
+                resolve(res.statusCode === 200);
+            });
+            req.on("error", () => resolve(false));
+            req.on("timeout", () => {
+                req.destroy();
+                resolve(false);
+            });
+        });
+        mlxAvailable = result;
+        lastCheck = now;
+        return result;
+    });
+}
+/**
+ * Get dense embeddings from MLX server.
+ * Returns Float32Array[] on success, null if server unavailable.
+ */
+function mlxEmbed(texts) {
+    return __awaiter(this, void 0, void 0, function* () {
+        if (EMBED_MODE === "cpu")
+            return null;
+        if (!(yield isMlxUp()))
+            return null;
+        const { ok, data } = yield postJSON("/embed", { texts });
+        if (!ok || !(data === null || data === void 0 ? void 0 : data.vectors)) {
+            mlxAvailable = false;
+            return null;
+        }
+        return data.vectors.map((v) => new Float32Array(v));
+    });
+}
+/**
+ * Reset availability cache (e.g., after starting the server).
+ */
+function resetMlxCache() {
+    mlxAvailable = null;
+    lastCheck = 0;
+}