/** * Embedding functions for converting data to vectors using Transformers.js * Complete rewrite to eliminate TensorFlow.js and use ONNX-based models */ import { isBrowser } from './environment.js'; // @ts-ignore - Transformers.js is now the primary embedding library import { pipeline, env } from '@huggingface/transformers'; /** * Detect the best available GPU device for the current environment */ export async function detectBestDevice() { // Browser environment - check for WebGPU support if (isBrowser()) { if (typeof navigator !== 'undefined' && 'gpu' in navigator) { try { const adapter = await navigator.gpu?.requestAdapter(); if (adapter) { return 'webgpu'; } } catch (error) { // WebGPU not available or failed to initialize } } return 'cpu'; } // Node.js environment - check for CUDA support try { // Check if ONNX Runtime GPU packages are available // This is a simple heuristic - in production you might want more sophisticated detection const hasGpu = process.env.CUDA_VISIBLE_DEVICES !== undefined || process.env.ONNXRUNTIME_GPU_ENABLED === 'true'; return hasGpu ? 'cuda' : 'cpu'; } catch (error) { return 'cpu'; } } /** * Resolve device string to actual device configuration */ export async function resolveDevice(device = 'auto') { if (device === 'auto') { return await detectBestDevice(); } // Map 'gpu' to appropriate GPU type for current environment if (device === 'gpu') { const detected = await detectBestDevice(); return detected === 'cpu' ? 'cpu' : detected; } return device; } export class TransformerEmbedding { /** * Create a new TransformerEmbedding instance */ constructor(options = {}) { this.extractor = null; this.initialized = false; this.verbose = true; this.verbose = options.verbose !== undefined ? options.verbose : true; // PRODUCTION-READY MODEL CONFIGURATION // Priority order: explicit option > environment variable > smart default let localFilesOnly; if (options.localFilesOnly !== undefined) { // 1. Explicit option takes highest priority localFilesOnly = options.localFilesOnly; } else if (process.env.BRAINY_ALLOW_REMOTE_MODELS !== undefined) { // 2. Environment variable override localFilesOnly = process.env.BRAINY_ALLOW_REMOTE_MODELS !== 'true'; } else if (process.env.NODE_ENV === 'development') { // 3. Development mode allows remote models localFilesOnly = false; } else if (isBrowser()) { // 4. Browser defaults to allowing remote models localFilesOnly = false; } else { // 5. Node.js production: try local first, but allow remote as fallback // This is the NEW production-friendly default localFilesOnly = false; } this.options = { model: options.model || 'Xenova/all-MiniLM-L6-v2', verbose: this.verbose, cacheDir: options.cacheDir || './models', localFilesOnly: localFilesOnly, dtype: options.dtype || 'fp32', device: options.device || 'auto' }; if (this.verbose) { this.logger('log', `Embedding config: localFilesOnly=${localFilesOnly}, model=${this.options.model}, cacheDir=${this.options.cacheDir}`); } // Configure transformers.js environment if (!isBrowser()) { // Set cache directory for Node.js env.cacheDir = this.options.cacheDir; // Prioritize local models for offline operation env.allowRemoteModels = !this.options.localFilesOnly; env.allowLocalModels = true; } else { // Browser configuration // Allow both local and remote models, but prefer local if available env.allowLocalModels = true; env.allowRemoteModels = true; // Force the configuration to ensure it's applied if (this.verbose) { this.logger('log', `Browser env config - allowLocalModels: ${env.allowLocalModels}, allowRemoteModels: ${env.allowRemoteModels}, localFilesOnly: ${this.options.localFilesOnly}`); } } } /** * Get the default cache directory for models */ async getDefaultCacheDir() { if (isBrowser()) { return './models'; // Browser default } // Check for bundled models in the package const possiblePaths = [ // In the installed package './node_modules/@soulcraft/brainy/models', // In development/source './models', './dist/../models', // Alternative locations '../models', '../../models' ]; // Check if we're in Node.js and try to find the bundled models if (typeof process !== 'undefined' && process.versions?.node) { try { // Use dynamic import instead of require for ES modules compatibility const { createRequire } = await import('module'); const require = createRequire(import.meta.url); const path = require('path'); const fs = require('fs'); // Try to resolve the package location try { const brainyPackagePath = require.resolve('@soulcraft/brainy/package.json'); const brainyPackageDir = path.dirname(brainyPackagePath); const bundledModelsPath = path.join(brainyPackageDir, 'models'); if (fs.existsSync(bundledModelsPath)) { this.logger('log', `Using bundled models from package: ${bundledModelsPath}`); return bundledModelsPath; } } catch (e) { // Not installed as package, continue } // Try relative paths from current location for (const relativePath of possiblePaths) { const fullPath = path.resolve(relativePath); if (fs.existsSync(fullPath)) { this.logger('log', `Using bundled models from: ${fullPath}`); return fullPath; } } } catch (error) { // Silently fall back to default path if module detection fails } } // Fallback to default cache directory return './models'; } /** * Check if we're running in a test environment */ isTestEnvironment() { // Always use real implementation - no more mocking return false; } /** * Log message only if verbose mode is enabled */ logger(level, message, ...args) { if (level === 'error' || this.verbose) { console[level](`[TransformerEmbedding] ${message}`, ...args); } } /** * Initialize the embedding model */ async init() { if (this.initialized) { return; } // Always use real implementation - no mocking try { // Resolve device configuration and cache directory const device = await resolveDevice(this.options.device); const cacheDir = this.options.cacheDir === './models' ? await this.getDefaultCacheDir() : this.options.cacheDir; this.logger('log', `Loading Transformer model: ${this.options.model} on device: ${device}`); const startTime = Date.now(); // Load the feature extraction pipeline with GPU support const pipelineOptions = { cache_dir: cacheDir, local_files_only: isBrowser() ? false : this.options.localFilesOnly, dtype: this.options.dtype }; // Add device configuration for GPU acceleration if (device !== 'cpu') { pipelineOptions.device = device; this.logger('log', `🚀 GPU acceleration enabled: ${device}`); } if (this.verbose) { this.logger('log', `Pipeline options: ${JSON.stringify(pipelineOptions)}`); } try { this.extractor = await pipeline('feature-extraction', this.options.model, pipelineOptions); } catch (gpuError) { // Fallback to CPU if GPU initialization fails if (device !== 'cpu') { this.logger('warn', `GPU initialization failed, falling back to CPU: ${gpuError?.message || gpuError}`); const cpuOptions = { ...pipelineOptions }; delete cpuOptions.device; this.extractor = await pipeline('feature-extraction', this.options.model, cpuOptions); } else { // PRODUCTION-READY ERROR HANDLING // If local_files_only is true and models are missing, try enabling remote downloads if (pipelineOptions.local_files_only && gpuError?.message?.includes('local_files_only')) { this.logger('warn', 'Local models not found, attempting remote download as fallback...'); try { const remoteOptions = { ...pipelineOptions, local_files_only: false }; this.extractor = await pipeline('feature-extraction', this.options.model, remoteOptions); this.logger('log', '✅ Successfully downloaded and loaded model from remote'); // Update the configuration to reflect what actually worked this.options.localFilesOnly = false; } catch (remoteError) { // Both local and remote failed - throw comprehensive error const errorMsg = `Failed to load embedding model "${this.options.model}". ` + `Local models not found and remote download failed. ` + `To fix: 1) Set BRAINY_ALLOW_REMOTE_MODELS=true, ` + `2) Run "npm run download-models", or ` + `3) Use a custom embedding function.`; throw new Error(errorMsg); } } else { throw gpuError; } } } const loadTime = Date.now() - startTime; this.logger('log', `✅ Model loaded successfully in ${loadTime}ms`); this.initialized = true; } catch (error) { this.logger('error', 'Failed to initialize Transformer embedding model:', error); throw new Error(`Transformer embedding initialization failed: ${error}`); } } /** * Generate embeddings for text data */ async embed(data) { if (!this.initialized) { await this.init(); } try { // Handle different input types let textToEmbed; if (typeof data === 'string') { // Handle empty string case if (data.trim() === '') { // Return a zero vector of 384 dimensions (all-MiniLM-L6-v2 standard) return new Array(384).fill(0); } textToEmbed = [data]; } else if (Array.isArray(data) && data.every((item) => typeof item === 'string')) { // Handle empty array or array with empty strings if (data.length === 0 || data.every((item) => item.trim() === '')) { return new Array(384).fill(0); } // Filter out empty strings textToEmbed = data.filter((item) => item.trim() !== ''); if (textToEmbed.length === 0) { return new Array(384).fill(0); } } else { throw new Error('TransformerEmbedding only supports string or string[] data'); } // Ensure the extractor is available if (!this.extractor) { throw new Error('Transformer embedding model is not available'); } // Generate embeddings with mean pooling and normalization const result = await this.extractor(textToEmbed, { pooling: 'mean', normalize: true }); // Extract the embedding data let embedding; if (textToEmbed.length === 1) { // Single text input - return first embedding embedding = Array.from(result.data.slice(0, 384)); } else { // Multiple texts - return first embedding (maintain compatibility) embedding = Array.from(result.data.slice(0, 384)); } // Validate embedding dimensions if (embedding.length !== 384) { this.logger('warn', `Unexpected embedding dimension: ${embedding.length}, expected 384`); // Pad or truncate to 384 dimensions if (embedding.length < 384) { embedding = [...embedding, ...new Array(384 - embedding.length).fill(0)]; } else { embedding = embedding.slice(0, 384); } } return embedding; } catch (error) { this.logger('error', 'Error generating embeddings:', error); throw new Error(`Failed to generate embeddings: ${error}`); } } /** * Dispose of the model and free resources */ async dispose() { if (this.extractor && typeof this.extractor.dispose === 'function') { await this.extractor.dispose(); } this.extractor = null; this.initialized = false; } /** * Get the dimension of embeddings produced by this model */ getDimension() { return 384; } /** * Check if the model is initialized */ isInitialized() { return this.initialized; } } // Legacy alias for backward compatibility export const UniversalSentenceEncoder = TransformerEmbedding; /** * Create a new embedding model instance */ export function createEmbeddingModel(options) { return new TransformerEmbedding(options); } /** * Default embedding function using the lightweight transformer model */ export const defaultEmbeddingFunction = async (data) => { const embedder = new TransformerEmbedding({ verbose: false }); return await embedder.embed(data); }; /** * Create an embedding function with custom options */ export function createEmbeddingFunction(options = {}) { const embedder = new TransformerEmbedding(options); return async (data) => { return await embedder.embed(data); }; } /** * Batch embedding function for processing multiple texts efficiently */ export async function batchEmbed(texts, options = {}) { const embedder = new TransformerEmbedding(options); await embedder.init(); const embeddings = []; // Process in batches for memory efficiency const batchSize = 32; for (let i = 0; i < texts.length; i += batchSize) { const batch = texts.slice(i, i + batchSize); for (const text of batch) { const embedding = await embedder.embed(text); embeddings.push(embedding); } } await embedder.dispose(); return embeddings; } /** * Embedding functions for specific model types */ export const embeddingFunctions = { /** Default lightweight model (all-MiniLM-L6-v2, 384 dimensions) */ default: defaultEmbeddingFunction, /** Create custom embedding function */ create: createEmbeddingFunction, /** Batch processing */ batch: batchEmbed }; //# sourceMappingURL=embedding.js.map