@soulcraft/brainy 6.5.0 → 6.6.1

package/dist/embeddings/wasm/AssetLoader.d.ts

@@ -0,0 +1,67 @@
+/**
+ * Asset Loader
+ *
+ * Resolves paths to model files (ONNX model, vocabulary) across environments.
+ * Handles Node.js, Bun, and bundled scenarios.
+ *
+ * Asset Resolution Order:
+ * 1. Environment variable: BRAINY_MODEL_PATH
+ * 2. Package-relative: node_modules/@soulcraft/brainy/assets/models/
+ * 3. Project-relative: ./assets/models/
+ */
+/**
+ * Asset loader for model files
+ */
+export declare class AssetLoader {
+    private modelDir;
+    /**
+     * Get the model directory path
+     */
+    getModelDir(): Promise<string>;
+    /**
+     * Resolve the model directory across environments
+     */
+    private resolveModelDir;
+    /**
+     * Get package root path (Node.js/Bun only)
+     */
+    private getPackageRootPath;
+    /**
+     * Check if path exists (works in Node.js/Bun)
+     */
+    private pathExists;
+    /**
+     * Get path to ONNX model file
+     */
+    getModelPath(): Promise<string>;
+    /**
+     * Get path to vocabulary file
+     */
+    getVocabPath(): Promise<string>;
+    /**
+     * Load vocabulary from JSON file
+     */
+    loadVocab(): Promise<Record<string, number>>;
+    /**
+     * Load model as ArrayBuffer (for ONNX session)
+     */
+    loadModel(): Promise<ArrayBuffer>;
+    /**
+     * Verify all required assets exist
+     */
+    verifyAssets(): Promise<{
+        valid: boolean;
+        modelPath: string;
+        vocabPath: string;
+        errors: string[];
+    }>;
+    /**
+     * Clear cached paths (for testing)
+     */
+    clearCache(): void;
+}
+/**
+ * Create asset loader instance
+ */
+export declare function createAssetLoader(): AssetLoader;
+export declare function getAssetLoader(): AssetLoader;

package/dist/embeddings/wasm/AssetLoader.js

@@ -0,0 +1,238 @@
+/**
+ * Asset Loader
+ *
+ * Resolves paths to model files (ONNX model, vocabulary) across environments.
+ * Handles Node.js, Bun, and bundled scenarios.
+ *
+ * Asset Resolution Order:
+ * 1. Environment variable: BRAINY_MODEL_PATH
+ * 2. Package-relative: node_modules/@soulcraft/brainy/assets/models/
+ * 3. Project-relative: ./assets/models/
+ */
+import { MODEL_CONSTANTS } from './types.js';
+// Cache resolved paths
+let cachedModelDir = null;
+let cachedVocab = null;
+/**
+ * Asset loader for model files
+ */
+export class AssetLoader {
+    constructor() {
+        this.modelDir = null;
+    }
+    /**
+     * Get the model directory path
+     */
+    async getModelDir() {
+        if (this.modelDir) {
+            return this.modelDir;
+        }
+        if (cachedModelDir) {
+            this.modelDir = cachedModelDir;
+            return cachedModelDir;
+        }
+        // Try to resolve model directory
+        const resolved = await this.resolveModelDir();
+        this.modelDir = resolved;
+        cachedModelDir = resolved;
+        return resolved;
+    }
+    /**
+     * Resolve the model directory across environments
+     */
+    async resolveModelDir() {
+        // 1. Check environment variable
+        if (typeof process !== 'undefined' && process.env?.BRAINY_MODEL_PATH) {
+            const envPath = process.env.BRAINY_MODEL_PATH;
+            if (await this.pathExists(envPath)) {
+                return envPath;
+            }
+        }
+        // 2. Try common locations
+        const modelName = MODEL_CONSTANTS.MODEL_NAME + '-q8';
+        const possiblePaths = [
+            // Package assets (when installed as dependency)
+            `./assets/models/${modelName}`,
+            `./node_modules/@soulcraft/brainy/assets/models/${modelName}`,
+            // Development paths
+            `../assets/models/${modelName}`,
+            // Absolute from package root
+            this.getPackageRootPath(`assets/models/${modelName}`),
+        ].filter(Boolean);
+        for (const path of possiblePaths) {
+            if (await this.pathExists(path)) {
+                return path;
+            }
+        }
+        // If no path found, return default (will error on use)
+        return `./assets/models/${modelName}`;
+    }
+    /**
+     * Get package root path (Node.js/Bun only)
+     */
+    getPackageRootPath(relativePath) {
+        if (typeof process === 'undefined') {
+            return null;
+        }
+        try {
+            // Use __dirname equivalent
+            const url = new URL(import.meta.url);
+            const currentDir = url.pathname.replace(/\/[^/]*$/, '');
+            // Go up from src/embeddings/wasm to package root
+            const packageRoot = currentDir.replace(/\/src\/embeddings\/wasm$/, '');
+            return `${packageRoot}/${relativePath}`;
+        }
+        catch {
+            return null;
+        }
+    }
+    /**
+     * Check if path exists (works in Node.js/Bun)
+     */
+    async pathExists(path) {
+        if (typeof process === 'undefined') {
+            // Browser - check via fetch
+            try {
+                const response = await fetch(path, { method: 'HEAD' });
+                return response.ok;
+            }
+            catch {
+                return false;
+            }
+        }
+        // Node.js/Bun
+        try {
+            const fs = await import('node:fs/promises');
+            await fs.access(path);
+            return true;
+        }
+        catch {
+            return false;
+        }
+    }
+    /**
+     * Get path to ONNX model file
+     */
+    async getModelPath() {
+        const dir = await this.getModelDir();
+        return `${dir}/model.onnx`;
+    }
+    /**
+     * Get path to vocabulary file
+     */
+    async getVocabPath() {
+        const dir = await this.getModelDir();
+        return `${dir}/vocab.json`;
+    }
+    /**
+     * Load vocabulary from JSON file
+     */
+    async loadVocab() {
+        if (cachedVocab) {
+            return cachedVocab;
+        }
+        const vocabPath = await this.getVocabPath();
+        if (typeof process !== 'undefined') {
+            // Node.js/Bun - read from filesystem
+            try {
+                const fs = await import('node:fs/promises');
+                const content = await fs.readFile(vocabPath, 'utf-8');
+                cachedVocab = JSON.parse(content);
+                return cachedVocab;
+            }
+            catch (error) {
+                throw new Error(`Failed to load vocabulary from ${vocabPath}: ${error instanceof Error ? error.message : String(error)}`);
+            }
+        }
+        else {
+            // Browser - fetch
+            try {
+                const response = await fetch(vocabPath);
+                if (!response.ok) {
+                    throw new Error(`HTTP ${response.status}`);
+                }
+                cachedVocab = await response.json();
+                return cachedVocab;
+            }
+            catch (error) {
+                throw new Error(`Failed to fetch vocabulary from ${vocabPath}: ${error instanceof Error ? error.message : String(error)}`);
+            }
+        }
+    }
+    /**
+     * Load model as ArrayBuffer (for ONNX session)
+     */
+    async loadModel() {
+        const modelPath = await this.getModelPath();
+        if (typeof process !== 'undefined') {
+            // Node.js/Bun - read from filesystem
+            try {
+                const fs = await import('node:fs/promises');
+                const buffer = await fs.readFile(modelPath);
+                // Convert Node.js Buffer to ArrayBuffer
+                return new Uint8Array(buffer).buffer;
+            }
+            catch (error) {
+                throw new Error(`Failed to load model from ${modelPath}: ${error instanceof Error ? error.message : String(error)}`);
+            }
+        }
+        else {
+            // Browser - fetch
+            try {
+                const response = await fetch(modelPath);
+                if (!response.ok) {
+                    throw new Error(`HTTP ${response.status}`);
+                }
+                return await response.arrayBuffer();
+            }
+            catch (error) {
+                throw new Error(`Failed to fetch model from ${modelPath}: ${error instanceof Error ? error.message : String(error)}`);
+            }
+        }
+    }
+    /**
+     * Verify all required assets exist
+     */
+    async verifyAssets() {
+        const errors = [];
+        const modelPath = await this.getModelPath();
+        const vocabPath = await this.getVocabPath();
+        if (!(await this.pathExists(modelPath))) {
+            errors.push(`Model file not found: ${modelPath}`);
+        }
+        if (!(await this.pathExists(vocabPath))) {
+            errors.push(`Vocabulary file not found: ${vocabPath}`);
+        }
+        return {
+            valid: errors.length === 0,
+            modelPath,
+            vocabPath,
+            errors,
+        };
+    }
+    /**
+     * Clear cached paths (for testing)
+     */
+    clearCache() {
+        this.modelDir = null;
+        cachedModelDir = null;
+        cachedVocab = null;
+    }
+}
+/**
+ * Create asset loader instance
+ */
+export function createAssetLoader() {
+    return new AssetLoader();
+}
+/**
+ * Singleton asset loader
+ */
+let singletonLoader = null;
+export function getAssetLoader() {
+    if (!singletonLoader) {
+        singletonLoader = new AssetLoader();
+    }
+    return singletonLoader;
+}
+//# sourceMappingURL=AssetLoader.js.map

package/dist/embeddings/wasm/EmbeddingPostProcessor.d.ts

@@ -0,0 +1,60 @@
+/**
+ * Embedding Post-Processor
+ *
+ * Converts raw ONNX model output to final embedding vectors.
+ * Implements mean pooling and L2 normalization as used by sentence-transformers.
+ *
+ * Pipeline:
+ * 1. Mean Pooling: Average token embeddings (weighted by attention mask)
+ * 2. L2 Normalization: Normalize to unit length for cosine similarity
+ */
+/**
+ * Post-processor for converting ONNX output to sentence embeddings
+ */
+export declare class EmbeddingPostProcessor {
+    private hiddenSize;
+    constructor(hiddenSize?: number);
+    /**
+     * Mean pool token embeddings weighted by attention mask
+     *
+     * @param hiddenStates - Raw model output [seqLen * hiddenSize] flattened
+     * @param attentionMask - Attention mask [seqLen] (1 for real tokens, 0 for padding)
+     * @param seqLen - Sequence length
+     * @returns Mean-pooled embedding [hiddenSize]
+     */
+    meanPool(hiddenStates: Float32Array, attentionMask: number[], seqLen: number): Float32Array;
+    /**
+     * L2 normalize embedding to unit length
+     *
+     * @param embedding - Input embedding
+     * @returns Normalized embedding with ||x|| = 1
+     */
+    normalize(embedding: Float32Array): Float32Array;
+    /**
+     * Full post-processing pipeline: mean pool then normalize
+     *
+     * @param hiddenStates - Raw model output [seqLen * hiddenSize]
+     * @param attentionMask - Attention mask [seqLen]
+     * @param seqLen - Sequence length
+     * @returns Final normalized embedding [hiddenSize]
+     */
+    process(hiddenStates: Float32Array, attentionMask: number[], seqLen: number): Float32Array;
+    /**
+     * Process batch of embeddings
+     *
+     * @param hiddenStates - Raw model output [batchSize * seqLen * hiddenSize]
+     * @param attentionMasks - Attention masks [batchSize][seqLen]
+     * @param batchSize - Number of sequences in batch
+     * @param seqLen - Sequence length (same for all in batch due to padding)
+     * @returns Array of normalized embeddings
+     */
+    processBatch(hiddenStates: Float32Array, attentionMasks: number[][], batchSize: number, seqLen: number): Float32Array[];
+    /**
+     * Convert Float32Array to number array
+     */
+    toNumberArray(embedding: Float32Array): number[];
+}
+/**
+ * Create a post-processor with default configuration
+ */
+export declare function createPostProcessor(): EmbeddingPostProcessor;

package/dist/embeddings/wasm/EmbeddingPostProcessor.js

@@ -0,0 +1,123 @@
+/**
+ * Embedding Post-Processor
+ *
+ * Converts raw ONNX model output to final embedding vectors.
+ * Implements mean pooling and L2 normalization as used by sentence-transformers.
+ *
+ * Pipeline:
+ * 1. Mean Pooling: Average token embeddings (weighted by attention mask)
+ * 2. L2 Normalization: Normalize to unit length for cosine similarity
+ */
+import { MODEL_CONSTANTS } from './types.js';
+/**
+ * Post-processor for converting ONNX output to sentence embeddings
+ */
+export class EmbeddingPostProcessor {
+    constructor(hiddenSize = MODEL_CONSTANTS.HIDDEN_SIZE) {
+        this.hiddenSize = hiddenSize;
+    }
+    /**
+     * Mean pool token embeddings weighted by attention mask
+     *
+     * @param hiddenStates - Raw model output [seqLen * hiddenSize] flattened
+     * @param attentionMask - Attention mask [seqLen] (1 for real tokens, 0 for padding)
+     * @param seqLen - Sequence length
+     * @returns Mean-pooled embedding [hiddenSize]
+     */
+    meanPool(hiddenStates, attentionMask, seqLen) {
+        const result = new Float32Array(this.hiddenSize);
+        // Sum of attention mask (number of real tokens)
+        let maskSum = 0;
+        for (let i = 0; i < seqLen; i++) {
+            maskSum += attentionMask[i];
+        }
+        // Avoid division by zero
+        if (maskSum === 0) {
+            maskSum = 1;
+        }
+        // Compute weighted sum for each dimension
+        for (let dim = 0; dim < this.hiddenSize; dim++) {
+            let sum = 0;
+            for (let pos = 0; pos < seqLen; pos++) {
+                // Get hidden state at [pos, dim]
+                const value = hiddenStates[pos * this.hiddenSize + dim];
+                // Weight by attention mask
+                sum += value * attentionMask[pos];
+            }
+            // Mean pool
+            result[dim] = sum / maskSum;
+        }
+        return result;
+    }
+    /**
+     * L2 normalize embedding to unit length
+     *
+     * @param embedding - Input embedding
+     * @returns Normalized embedding with ||x|| = 1
+     */
+    normalize(embedding) {
+        // Compute L2 norm
+        let sumSquares = 0;
+        for (let i = 0; i < embedding.length; i++) {
+            sumSquares += embedding[i] * embedding[i];
+        }
+        const norm = Math.sqrt(sumSquares);
+        // Avoid division by zero
+        if (norm === 0) {
+            return embedding;
+        }
+        // Normalize
+        const result = new Float32Array(embedding.length);
+        for (let i = 0; i < embedding.length; i++) {
+            result[i] = embedding[i] / norm;
+        }
+        return result;
+    }
+    /**
+     * Full post-processing pipeline: mean pool then normalize
+     *
+     * @param hiddenStates - Raw model output [seqLen * hiddenSize]
+     * @param attentionMask - Attention mask [seqLen]
+     * @param seqLen - Sequence length
+     * @returns Final normalized embedding [hiddenSize]
+     */
+    process(hiddenStates, attentionMask, seqLen) {
+        const pooled = this.meanPool(hiddenStates, attentionMask, seqLen);
+        return this.normalize(pooled);
+    }
+    /**
+     * Process batch of embeddings
+     *
+     * @param hiddenStates - Raw model output [batchSize * seqLen * hiddenSize]
+     * @param attentionMasks - Attention masks [batchSize][seqLen]
+     * @param batchSize - Number of sequences in batch
+     * @param seqLen - Sequence length (same for all in batch due to padding)
+     * @returns Array of normalized embeddings
+     */
+    processBatch(hiddenStates, attentionMasks, batchSize, seqLen) {
+        const results = [];
+        const sequenceSize = seqLen * this.hiddenSize;
+        for (let b = 0; b < batchSize; b++) {
+            // Extract this sequence's hidden states
+            const start = b * sequenceSize;
+            const seqHiddenStates = hiddenStates.slice(start, start + sequenceSize);
+            // Process
+            const embedding = this.process(seqHiddenStates, attentionMasks[b], seqLen);
+            results.push(embedding);
+        }
+        return results;
+    }
+    /**
+     * Convert Float32Array to number array
+     */
+    toNumberArray(embedding) {
+        return Array.from(embedding);
+    }
+}
+/**
+ * Create a post-processor with default configuration
+ */
+export function createPostProcessor() {
+    return new EmbeddingPostProcessor(MODEL_CONSTANTS.HIDDEN_SIZE);
+}
+//# sourceMappingURL=EmbeddingPostProcessor.js.map

package/dist/embeddings/wasm/ONNXInferenceEngine.d.ts

@@ -0,0 +1,55 @@
+/**
+ * ONNX Inference Engine
+ *
+ * Direct ONNX Runtime Web wrapper for running model inference.
+ * Uses WASM backend for universal compatibility (Node.js, Bun, Browser).
+ *
+ * This replaces transformers.js dependency with direct ONNX control.
+ */
+import { InferenceConfig } from './types.js';
+/**
+ * ONNX Inference Engine using onnxruntime-web
+ */
+export declare class ONNXInferenceEngine {
+    private session;
+    private initialized;
+    private modelPath;
+    private config;
+    constructor(config?: Partial<InferenceConfig>);
+    /**
+     * Initialize the ONNX session
+     */
+    initialize(modelPath?: string): Promise<void>;
+    /**
+     * Run inference on tokenized input
+     *
+     * @param inputIds - Token IDs [batchSize, seqLen]
+     * @param attentionMask - Attention mask [batchSize, seqLen]
+     * @param tokenTypeIds - Token type IDs [batchSize, seqLen] (optional, defaults to zeros)
+     * @returns Hidden states [batchSize, seqLen, hiddenSize]
+     */
+    infer(inputIds: number[][], attentionMask: number[][], tokenTypeIds?: number[][]): Promise<Float32Array>;
+    /**
+     * Infer single sequence (convenience method)
+     */
+    inferSingle(inputIds: number[], attentionMask: number[], tokenTypeIds?: number[]): Promise<Float32Array>;
+    /**
+     * Check if initialized
+     */
+    isInitialized(): boolean;
+    /**
+     * Get model input/output names (for debugging)
+     */
+    getModelInfo(): {
+        inputs: readonly string[];
+        outputs: readonly string[];
+    } | null;
+    /**
+     * Dispose of the session and free resources
+     */
+    dispose(): Promise<void>;
+}
+/**
+ * Create an inference engine with default configuration
+ */
+export declare function createInferenceEngine(modelPath: string): ONNXInferenceEngine;

package/dist/embeddings/wasm/ONNXInferenceEngine.js

@@ -0,0 +1,154 @@
+/**
+ * ONNX Inference Engine
+ *
+ * Direct ONNX Runtime Web wrapper for running model inference.
+ * Uses WASM backend for universal compatibility (Node.js, Bun, Browser).
+ *
+ * This replaces transformers.js dependency with direct ONNX control.
+ */
+import * as ort from 'onnxruntime-web';
+// Configure ONNX Runtime for WASM-only
+ort.env.wasm.numThreads = 1; // Single-threaded for stability
+ort.env.wasm.simd = true; // Enable SIMD where available
+/**
+ * ONNX Inference Engine using onnxruntime-web
+ */
+export class ONNXInferenceEngine {
+    constructor(config = {}) {
+        this.session = null;
+        this.initialized = false;
+        this.modelPath = config.modelPath ?? '';
+        this.config = {
+            modelPath: this.modelPath,
+            numThreads: config.numThreads ?? 1,
+            enableSimd: config.enableSimd ?? true,
+            enableCpuMemArena: config.enableCpuMemArena ?? false,
+        };
+    }
+    /**
+     * Initialize the ONNX session
+     */
+    async initialize(modelPath) {
+        if (this.initialized && this.session) {
+            return;
+        }
+        const path = modelPath ?? this.modelPath;
+        if (!path) {
+            throw new Error('Model path is required');
+        }
+        try {
+            // Configure session options
+            const sessionOptions = {
+                executionProviders: ['wasm'],
+                graphOptimizationLevel: 'all',
+                enableCpuMemArena: this.config.enableCpuMemArena,
+                // Additional WASM-specific options
+                executionMode: 'sequential',
+            };
+            // Load model from file path or URL
+            this.session = await ort.InferenceSession.create(path, sessionOptions);
+            this.initialized = true;
+        }
+        catch (error) {
+            this.initialized = false;
+            this.session = null;
+            throw new Error(`Failed to initialize ONNX session: ${error instanceof Error ? error.message : String(error)}`);
+        }
+    }
+    /**
+     * Run inference on tokenized input
+     *
+     * @param inputIds - Token IDs [batchSize, seqLen]
+     * @param attentionMask - Attention mask [batchSize, seqLen]
+     * @param tokenTypeIds - Token type IDs [batchSize, seqLen] (optional, defaults to zeros)
+     * @returns Hidden states [batchSize, seqLen, hiddenSize]
+     */
+    async infer(inputIds, attentionMask, tokenTypeIds) {
+        if (!this.session) {
+            throw new Error('Session not initialized. Call initialize() first.');
+        }
+        const batchSize = inputIds.length;
+        const seqLen = inputIds[0].length;
+        // Convert to BigInt64Array (ONNX int64 type)
+        const inputIdsFlat = new BigInt64Array(batchSize * seqLen);
+        const attentionMaskFlat = new BigInt64Array(batchSize * seqLen);
+        const tokenTypeIdsFlat = new BigInt64Array(batchSize * seqLen);
+        for (let b = 0; b < batchSize; b++) {
+            for (let s = 0; s < seqLen; s++) {
+                const idx = b * seqLen + s;
+                inputIdsFlat[idx] = BigInt(inputIds[b][s]);
+                attentionMaskFlat[idx] = BigInt(attentionMask[b][s]);
+                tokenTypeIdsFlat[idx] = tokenTypeIds
+                    ? BigInt(tokenTypeIds[b][s])
+                    : BigInt(0);
+            }
+        }
+        // Create ONNX tensors
+        const inputIdsTensor = new ort.Tensor('int64', inputIdsFlat, [batchSize, seqLen]);
+        const attentionMaskTensor = new ort.Tensor('int64', attentionMaskFlat, [batchSize, seqLen]);
+        const tokenTypeIdsTensor = new ort.Tensor('int64', tokenTypeIdsFlat, [batchSize, seqLen]);
+        try {
+            // Run inference
+            const feeds = {
+                input_ids: inputIdsTensor,
+                attention_mask: attentionMaskTensor,
+                token_type_ids: tokenTypeIdsTensor,
+            };
+            const results = await this.session.run(feeds);
+            // Extract last_hidden_state (the output we need for mean pooling)
+            // Model outputs: last_hidden_state [batch, seq, hidden] and pooler_output [batch, hidden]
+            const output = results.last_hidden_state ?? results.token_embeddings;
+            if (!output) {
+                throw new Error('Model did not return expected output tensor');
+            }
+            return output.data;
+        }
+        finally {
+            // Dispose tensors to free memory
+            inputIdsTensor.dispose();
+            attentionMaskTensor.dispose();
+            tokenTypeIdsTensor.dispose();
+        }
+    }
+    /**
+     * Infer single sequence (convenience method)
+     */
+    async inferSingle(inputIds, attentionMask, tokenTypeIds) {
+        return this.infer([inputIds], [attentionMask], tokenTypeIds ? [tokenTypeIds] : undefined);
+    }
+    /**
+     * Check if initialized
+     */
+    isInitialized() {
+        return this.initialized;
+    }
+    /**
+     * Get model input/output names (for debugging)
+     */
+    getModelInfo() {
+        if (!this.session) {
+            return null;
+        }
+        return {
+            inputs: this.session.inputNames,
+            outputs: this.session.outputNames,
+        };
+    }
+    /**
+     * Dispose of the session and free resources
+     */
+    async dispose() {
+        if (this.session) {
+            // Release the session
+            this.session = null;
+        }
+        this.initialized = false;
+    }
+}
+/**
+ * Create an inference engine with default configuration
+ */
+export function createInferenceEngine(modelPath) {
+    return new ONNXInferenceEngine({ modelPath });
+}
+//# sourceMappingURL=ONNXInferenceEngine.js.map