@soulcraft/brainy 6.4.0 → 6.6.0

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

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

@@ -0,0 +1,82 @@
+/**
+ * WASM Embedding Engine
+ *
+ * The main embedding engine that combines all components:
+ * - WordPieceTokenizer: Text → Token IDs
+ * - ONNXInferenceEngine: Token IDs → Hidden States
+ * - EmbeddingPostProcessor: Hidden States → Normalized Embedding
+ *
+ * This replaces transformers.js with a clean, production-grade implementation.
+ *
+ * Features:
+ * - Singleton pattern (one model instance)
+ * - Lazy initialization
+ * - Batch processing support
+ * - Zero runtime dependencies
+ */
+import { EmbeddingResult, EngineStats } from './types.js';
+/**
+ * WASM-based embedding engine
+ */
+export declare class WASMEmbeddingEngine {
+    private tokenizer;
+    private inference;
+    private postProcessor;
+    private initialized;
+    private embedCount;
+    private totalProcessingTimeMs;
+    private constructor();
+    /**
+     * Get the singleton instance
+     */
+    static getInstance(): WASMEmbeddingEngine;
+    /**
+     * Initialize all components
+     */
+    initialize(): Promise<void>;
+    /**
+     * Perform actual initialization
+     */
+    private performInit;
+    /**
+     * Generate embedding for text
+     */
+    embed(text: string): Promise<number[]>;
+    /**
+     * Generate embedding with metadata
+     */
+    embedWithMetadata(text: string): Promise<EmbeddingResult>;
+    /**
+     * Batch embed multiple texts
+     */
+    embedBatch(texts: string[]): Promise<number[][]>;
+    /**
+     * Check if initialized
+     */
+    isInitialized(): boolean;
+    /**
+     * Get engine statistics
+     */
+    getStats(): EngineStats;
+    /**
+     * Dispose and free resources
+     */
+    dispose(): Promise<void>;
+    /**
+     * Reset singleton (for testing)
+     */
+    static resetInstance(): void;
+}
+export declare const wasmEmbeddingEngine: WASMEmbeddingEngine;
+/**
+ * Convenience function to get embeddings
+ */
+export declare function embed(text: string): Promise<number[]>;
+/**
+ * Convenience function for batch embeddings
+ */
+export declare function embedBatch(texts: string[]): Promise<number[][]>;
+/**
+ * Get embedding stats
+ */
+export declare function getEmbeddingStats(): EngineStats;

package/dist/embeddings/wasm/WASMEmbeddingEngine.js

@@ -0,0 +1,231 @@
+/**
+ * WASM Embedding Engine
+ *
+ * The main embedding engine that combines all components:
+ * - WordPieceTokenizer: Text → Token IDs
+ * - ONNXInferenceEngine: Token IDs → Hidden States
+ * - EmbeddingPostProcessor: Hidden States → Normalized Embedding
+ *
+ * This replaces transformers.js with a clean, production-grade implementation.
+ *
+ * Features:
+ * - Singleton pattern (one model instance)
+ * - Lazy initialization
+ * - Batch processing support
+ * - Zero runtime dependencies
+ */
+import { WordPieceTokenizer } from './WordPieceTokenizer.js';
+import { ONNXInferenceEngine } from './ONNXInferenceEngine.js';
+import { EmbeddingPostProcessor } from './EmbeddingPostProcessor.js';
+import { getAssetLoader } from './AssetLoader.js';
+import { MODEL_CONSTANTS } from './types.js';
+// Global singleton instance
+let globalInstance = null;
+let globalInitPromise = null;
+/**
+ * WASM-based embedding engine
+ */
+export class WASMEmbeddingEngine {
+    constructor() {
+        this.tokenizer = null;
+        this.inference = null;
+        this.postProcessor = null;
+        this.initialized = false;
+        this.embedCount = 0;
+        this.totalProcessingTimeMs = 0;
+        // Private constructor for singleton
+    }
+    /**
+     * Get the singleton instance
+     */
+    static getInstance() {
+        if (!globalInstance) {
+            globalInstance = new WASMEmbeddingEngine();
+        }
+        return globalInstance;
+    }
+    /**
+     * Initialize all components
+     */
+    async initialize() {
+        // Already initialized
+        if (this.initialized) {
+            return;
+        }
+        // Initialization in progress
+        if (globalInitPromise) {
+            await globalInitPromise;
+            return;
+        }
+        // Start initialization
+        globalInitPromise = this.performInit();
+        try {
+            await globalInitPromise;
+        }
+        finally {
+            globalInitPromise = null;
+        }
+    }
+    /**
+     * Perform actual initialization
+     */
+    async performInit() {
+        const startTime = Date.now();
+        console.log('🚀 Initializing WASM Embedding Engine...');
+        try {
+            const assetLoader = getAssetLoader();
+            // Verify assets exist
+            const verification = await assetLoader.verifyAssets();
+            if (!verification.valid) {
+                throw new Error(`Missing model assets:\n${verification.errors.join('\n')}\n\n` +
+                    `Expected model at: ${verification.modelPath}\n` +
+                    `Expected vocab at: ${verification.vocabPath}\n\n` +
+                    `Run 'npm run download-model' to download the model files.`);
+            }
+            // Load vocabulary and create tokenizer
+            console.log('📖 Loading vocabulary...');
+            const vocab = await assetLoader.loadVocab();
+            this.tokenizer = new WordPieceTokenizer(vocab);
+            console.log(`✅ Vocabulary loaded: ${this.tokenizer.vocabSize} tokens`);
+            // Initialize ONNX inference engine
+            console.log('🧠 Loading ONNX model...');
+            const modelPath = await assetLoader.getModelPath();
+            this.inference = new ONNXInferenceEngine({ modelPath });
+            await this.inference.initialize(modelPath);
+            console.log('✅ ONNX model loaded');
+            // Create post-processor
+            this.postProcessor = new EmbeddingPostProcessor(MODEL_CONSTANTS.HIDDEN_SIZE);
+            this.initialized = true;
+            const initTime = Date.now() - startTime;
+            console.log(`✅ WASM Embedding Engine ready in ${initTime}ms`);
+        }
+        catch (error) {
+            this.initialized = false;
+            this.tokenizer = null;
+            this.inference = null;
+            this.postProcessor = null;
+            throw new Error(`Failed to initialize WASM Embedding Engine: ${error instanceof Error ? error.message : String(error)}`);
+        }
+    }
+    /**
+     * Generate embedding for text
+     */
+    async embed(text) {
+        const result = await this.embedWithMetadata(text);
+        return result.embedding;
+    }
+    /**
+     * Generate embedding with metadata
+     */
+    async embedWithMetadata(text) {
+        // Ensure initialized
+        if (!this.initialized) {
+            await this.initialize();
+        }
+        if (!this.tokenizer || !this.inference || !this.postProcessor) {
+            throw new Error('Engine not properly initialized');
+        }
+        const startTime = Date.now();
+        // 1. Tokenize
+        const tokenized = this.tokenizer.encode(text);
+        // 2. Run inference
+        const hiddenStates = await this.inference.inferSingle(tokenized.inputIds, tokenized.attentionMask, tokenized.tokenTypeIds);
+        // 3. Post-process (mean pool + normalize)
+        const embedding = this.postProcessor.process(hiddenStates, tokenized.attentionMask, tokenized.inputIds.length);
+        const processingTimeMs = Date.now() - startTime;
+        this.embedCount++;
+        this.totalProcessingTimeMs += processingTimeMs;
+        return {
+            embedding: Array.from(embedding),
+            tokenCount: tokenized.tokenCount,
+            processingTimeMs,
+        };
+    }
+    /**
+     * Batch embed multiple texts
+     */
+    async embedBatch(texts) {
+        // Ensure initialized
+        if (!this.initialized) {
+            await this.initialize();
+        }
+        if (!this.tokenizer || !this.inference || !this.postProcessor) {
+            throw new Error('Engine not properly initialized');
+        }
+        if (texts.length === 0) {
+            return [];
+        }
+        // Tokenize all texts
+        const batch = this.tokenizer.encodeBatch(texts);
+        const seqLen = batch.inputIds[0].length;
+        // Run batch inference
+        const hiddenStates = await this.inference.infer(batch.inputIds, batch.attentionMask, batch.tokenTypeIds);
+        // Post-process each result
+        const embeddings = this.postProcessor.processBatch(hiddenStates, batch.attentionMask, texts.length, seqLen);
+        this.embedCount += texts.length;
+        return embeddings.map(e => Array.from(e));
+    }
+    /**
+     * Check if initialized
+     */
+    isInitialized() {
+        return this.initialized;
+    }
+    /**
+     * Get engine statistics
+     */
+    getStats() {
+        return {
+            initialized: this.initialized,
+            embedCount: this.embedCount,
+            totalProcessingTimeMs: this.totalProcessingTimeMs,
+            avgProcessingTimeMs: this.embedCount > 0
+                ? this.totalProcessingTimeMs / this.embedCount
+                : 0,
+            modelName: MODEL_CONSTANTS.MODEL_NAME,
+        };
+    }
+    /**
+     * Dispose and free resources
+     */
+    async dispose() {
+        if (this.inference) {
+            await this.inference.dispose();
+            this.inference = null;
+        }
+        this.tokenizer = null;
+        this.postProcessor = null;
+        this.initialized = false;
+    }
+    /**
+     * Reset singleton (for testing)
+     */
+    static resetInstance() {
+        if (globalInstance) {
+            globalInstance.dispose();
+        }
+        globalInstance = null;
+        globalInitPromise = null;
+    }
+}
+// Export singleton access
+export const wasmEmbeddingEngine = WASMEmbeddingEngine.getInstance();
+/**
+ * Convenience function to get embeddings
+ */
+export async function embed(text) {
+    return wasmEmbeddingEngine.embed(text);
+}
+/**
+ * Convenience function for batch embeddings
+ */
+export async function embedBatch(texts) {
+    return wasmEmbeddingEngine.embedBatch(texts);
+}
+/**
+ * Get embedding stats
+ */
+export function getEmbeddingStats() {
+    return wasmEmbeddingEngine.getStats();
+}
+//# sourceMappingURL=WASMEmbeddingEngine.js.map