diff --git a/src/utils/distance.ts b/src/utils/distance.ts index a7c98560..8931d7fd 100644 --- a/src/utils/distance.ts +++ b/src/utils/distance.ts @@ -1,9 +1,12 @@ /** * Distance functions for vector similarity calculations * Optimized for Node.js 23.11+ using enhanced array methods + * GPU-accelerated versions available for high-performance computing */ import { DistanceFunction, Vector } from '../coreTypes.js' +import { executeInThread } from './workerUtils.js' +import { isThreadingAvailable } from './environment.js' /** * Calculates the Euclidean distance between two vectors @@ -41,8 +44,8 @@ export const cosineDistance: DistanceFunction = (a: Vector, b: Vector): number = dotProduct: acc.dotProduct + (val * b[i]), normA: acc.normA + (val * val), normB: acc.normB + (b[i] * b[i]) - }; - }, { dotProduct: 0, normA: 0, normB: 0 }); + } + }, { dotProduct: 0, normA: 0, normB: 0 }) if (normA === 0 || normB === 0) { return 2 // Maximum distance for zero vectors @@ -84,3 +87,198 @@ export const dotProductDistance: DistanceFunction = (a: Vector, b: Vector): numb // Convert to a distance metric (lower is better) return -dotProduct } + +/** + * GPU-accelerated batch distance calculation + * Uses TensorFlow.js with WebGL backend when available for optimal performance + * Falls back to CPU processing when GPU is not available + * + * @param queryVector The query vector to compare against all vectors + * @param vectors Array of vectors to compare against + * @param distanceFunction The distance function to use + * @returns Promise resolving to array of distances + */ +export async function calculateDistancesWithGPU( + queryVector: Vector, + vectors: Vector[], + distanceFunction: DistanceFunction = euclideanDistance +): Promise { + // For small batches, use the standard distance function + if (vectors.length < 10) { + return vectors.map(vector => distanceFunction(queryVector, vector)) + } + + try { + // Function to be executed in a worker thread + const distanceCalculator = async ( + args: { + queryVector: Vector, + vectors: Vector[], + distanceFnString: string + } + ) => { + const { queryVector, vectors, distanceFnString } = args + + // Try to use TensorFlow.js with GPU acceleration if available + const useTensorFlow = async () => { + try { + // TensorFlow.js will use its default EPSILON value + + // Dynamically import TensorFlow.js core module and backends + const tf = await import('@tensorflow/tfjs-core') + + // Import CPU backend as fallback + await import('@tensorflow/tfjs-backend-cpu') + + let usingGPU = false + + try { + // Try to import and use WebGL backend (GPU) + await import('@tensorflow/tfjs-backend-webgl') + + // Check if WebGL is available and set it as the backend + if (await tf.findBackend('webgl') || await tf.ready().then(() => tf.findBackend('webgl'))) { + await tf.setBackend('webgl') + usingGPU = true + } else { + await tf.setBackend('cpu') + } + } catch (err) { + // If WebGL fails, use CPU + await tf.setBackend('cpu') + } + + // Convert vectors to tensors + const queryTensor = tf.tensor2d([queryVector]) + const vectorsTensor = tf.tensor2d(vectors) + + let distances: number[] + + // Calculate distances based on the distance function type + if (distanceFnString.includes('euclideanDistance')) { + // Euclidean distance using GPU-optimized operations + // Formula: sqrt(sum((a - b)^2)) + const expanded = tf.sub((queryTensor as any).expandDims(1), (vectorsTensor as any).expandDims(0)) + const squaredDiff = tf.square(expanded) + const sumSquaredDiff = tf.sum(squaredDiff, -1) + const distancesTensor = tf.sqrt(sumSquaredDiff) + distances = await (distancesTensor as any).squeeze().array() as number[] + + // Clean up tensors + queryTensor.dispose() + vectorsTensor.dispose() + expanded.dispose() + squaredDiff.dispose() + sumSquaredDiff.dispose() + distancesTensor.dispose() + } else if (distanceFnString.includes('cosineDistance')) { + // Cosine distance using GPU-optimized operations + // Formula: 1 - (a·b / (||a|| * ||b||)) + const dotProduct = tf.matMul(queryTensor, (vectorsTensor as any).transpose()) + + const queryNorm = tf.norm(queryTensor, 2, 1) + const vectorsNorm = tf.norm(vectorsTensor, 2, 1) + + const normProduct = tf.outerProduct(queryNorm as any, vectorsNorm as any) + const cosineSimilarity = tf.div(dotProduct, normProduct) + const distancesTensor = tf.sub(tf.scalar(1), cosineSimilarity) + + distances = await (distancesTensor as any).squeeze().array() as number[] + + // Clean up tensors + queryTensor.dispose() + vectorsTensor.dispose() + dotProduct.dispose() + queryNorm.dispose() + vectorsNorm.dispose() + normProduct.dispose() + cosineSimilarity.dispose() + distancesTensor.dispose() + } else if (distanceFnString.includes('manhattanDistance')) { + // Manhattan distance using GPU-optimized operations + // Formula: sum(|a - b|) + const diff = tf.sub((queryTensor as any).expandDims(1), (vectorsTensor as any).expandDims(0)) + const absDiff = tf.abs(diff) + const distancesTensor = tf.sum(absDiff, -1) + + distances = await (distancesTensor as any).squeeze().array() as number[] + + // Clean up tensors + queryTensor.dispose() + vectorsTensor.dispose() + diff.dispose() + absDiff.dispose() + distancesTensor.dispose() + } else if (distanceFnString.includes('dotProductDistance')) { + // Dot product distance using GPU-optimized operations + // Formula: -sum(a * b) + const dotProduct = tf.matMul(queryTensor, (vectorsTensor as any).transpose()) + const distancesTensor = tf.neg(dotProduct) + + distances = await (distancesTensor as any).squeeze().array() as number[] + + // Clean up tensors + queryTensor.dispose() + vectorsTensor.dispose() + dotProduct.dispose() + distancesTensor.dispose() + } else { + // For unknown distance functions, fall back to CPU implementation + throw new Error('Unsupported distance function for GPU acceleration') + } + + return { + distances, + usingGPU + } + } catch (error) { + // If TensorFlow.js fails, fall back to CPU implementation + throw error + } + } + + // Try to use TensorFlow.js with GPU acceleration + try { + return await useTensorFlow() + } catch (error) { + // Fall back to CPU implementation if TensorFlow.js fails + // Recreate the distance function from its string representation + const distanceFunction = new Function('return ' + distanceFnString)() as DistanceFunction + + // Calculate distances for all vectors + const distances = vectors.map(vector => distanceFunction(queryVector, vector)) + + return { + distances, + usingGPU: false + } + } + } + + // Execute the distance calculation in a separate thread if threading is available + if (isThreadingAvailable()) { + try { + // Convert the distance function to a string for serialization + const distanceFnString = distanceFunction.toString() + + // Execute in a separate thread + const result = await executeInThread<{ distances: number[], usingGPU: boolean }>( + distanceCalculator.toString(), + { queryVector, vectors, distanceFnString } + ) + + return result.distances + } catch (error) { + // Fall back to main thread if threading fails + console.warn('Threaded distance calculation failed, falling back to main thread:', error) + } + } + + // If threading is not available or failed, calculate distances in the main thread + return vectors.map(vector => distanceFunction(queryVector, vector)) + } catch (error) { + // If anything fails, fall back to the standard distance function + console.error('GPU-accelerated distance calculation failed:', error) + return vectors.map(vector => distanceFunction(queryVector, vector)) + } +} diff --git a/src/utils/embedding.ts b/src/utils/embedding.ts index 495a49fb..4c8c8c46 100644 --- a/src/utils/embedding.ts +++ b/src/utils/embedding.ts @@ -9,6 +9,9 @@ import { executeInThread } from './workerUtils.js' * TensorFlow Universal Sentence Encoder embedding model * This model provides high-quality text embeddings using TensorFlow.js * The required TensorFlow.js dependencies are automatically installed with this package + * + * This implementation will use GPU acceleration via WebGL when available, + * falling back to CPU processing when GPU is not available or fails to initialize. */ export class UniversalSentenceEncoder implements EmbeddingModel { private model: any = null @@ -38,31 +41,38 @@ export class UniversalSentenceEncoder implements EmbeddingModel { originalWarn(message, ...optionalParams) } - // Define EPSILON flag before TensorFlow.js is loaded - // This prevents the "Cannot evaluate flag 'EPSILON': no evaluation function found" error - if (typeof window !== 'undefined') { - ;(window as any).EPSILON = 1e-7 - // Define the flag with an evaluation function for TensorFlow.js - ;(window as any).ENV = (window as any).ENV || {} - ;(window as any).ENV.flagRegistry = - (window as any).ENV.flagRegistry || {} - ;(window as any).ENV.flagRegistry.EPSILON = { - evaluationFn: () => 1e-7 - } - } else if (typeof global !== 'undefined') { - ;(global as any).EPSILON = 1e-7 - // Define the flag with an evaluation function for TensorFlow.js - ;(global as any).ENV = (global as any).ENV || {} - ;(global as any).ENV.flagRegistry = - (global as any).ENV.flagRegistry || {} - ;(global as any).ENV.flagRegistry.EPSILON = { - evaluationFn: () => 1e-7 + // TensorFlow.js will use its default EPSILON value + + // Dynamically import TensorFlow.js core module and backends + // Use type assertions to tell TypeScript these modules exist + this.tf = await import('@tensorflow/tfjs-core') + + // Import CPU and WebGL backends + await import('@tensorflow/tfjs-backend-cpu') + + try { + // Try to import and use WebGL backend first (GPU) + await import('@tensorflow/tfjs-backend-webgl') + + // Check if WebGL is available and set it as the backend + if ( + (await this.tf.findBackend('webgl')) || + (await this.tf.ready().then(() => this.tf.findBackend('webgl'))) + ) { + console.log('Using WebGL backend (GPU acceleration)') + await this.tf.setBackend('webgl') + } else { + console.log('WebGL backend not available, falling back to CPU') + await this.tf.setBackend('cpu') } + } catch (err) { + console.warn( + 'WebGL backend failed to initialize, using CPU backend:', + err + ) + await this.tf.setBackend('cpu') } - // Dynamically import TensorFlow.js and Universal Sentence Encoder - // Use type assertions to tell TypeScript these modules exist - this.tf = await import('@tensorflow/tfjs') this.use = await import('@tensorflow-models/universal-sentence-encoder') // Load the model @@ -193,9 +203,13 @@ export function createTensorFlowEmbeddingFunction(): EmbeddingFunction { /** * Creates a TensorFlow-based Universal Sentence Encoder embedding function that runs in a separate thread - * This provides better performance for CPU-intensive embedding operations + * This provides better performance for embedding operations by: + * 1. Using GPU acceleration via WebGL when available + * 2. Running in a separate thread to avoid blocking the main thread + * 3. Falling back to CPU processing when GPU is not available + * * @param options Configuration options - * @returns An embedding function that runs in a separate thread + * @returns An embedding function that runs in a separate thread with GPU acceleration when available */ export function createThreadedEmbeddingFunction( options: { fallbackToMain?: boolean } = {} @@ -223,56 +237,113 @@ export function createThreadedEmbeddingFunction( // Worker implementation function that will be stringified and run in the worker const workerImplementation = async ({ data }: { data: any }) => { - // We need to dynamically import TensorFlow.js and USE in the worker - const tf = await import('@tensorflow/tfjs') - const use = await import('@tensorflow-models/universal-sentence-encoder') + try { + // We need to dynamically import TensorFlow.js core module and USE in the worker + // Use a variable name that won't conflict with any minified variables - // Load the model - const model = await use.load() + // TensorFlow.js will use its default EPSILON value - // Handle different input types - let textToEmbed: string[] - if (typeof data === 'string') { - if (data.trim() === '') { - return new Array(512).fill(0) + // Import TensorFlow.js modules + const tf = await import('@tensorflow/tfjs-core') + + // Import CPU backend in the worker + await import('@tensorflow/tfjs-backend-cpu') + + try { + // Try to import and use WebGL backend first (GPU) + await import('@tensorflow/tfjs-backend-webgl') + + // Check if WebGL is available and set it as the backend + if ( + (tf.findBackend && (await tf.findBackend('webgl'))) || + (tf.ready && + (await tf + .ready() + .then(() => tf.findBackend && tf.findBackend('webgl')))) + ) { + console.log('Worker: Using WebGL backend (GPU acceleration)') + if (tf.setBackend) { + await tf.setBackend('webgl') + } + } else { + console.log( + 'Worker: WebGL backend not available, falling back to CPU' + ) + if (tf.setBackend) { + await tf.setBackend('cpu') + } + } + } catch (err) { + console.warn( + 'Worker: WebGL backend failed to initialize, using CPU backend:', + err + ) + if (tf.setBackend) { + await tf.setBackend('cpu') + } } - textToEmbed = [data] - } else if ( - Array.isArray(data) && - data.every((item) => typeof item === 'string') - ) { - if (data.length === 0 || data.every((item) => item.trim() === '')) { - return new Array(512).fill(0) - } - textToEmbed = data.filter((item) => item.trim() !== '') - if (textToEmbed.length === 0) { - return new Array(512).fill(0) - } - } else { - throw new Error( - 'UniversalSentenceEncoder only supports string or string[] data' + + // Import the Universal Sentence Encoder + const sentenceEncoderModule = await import( + '@tensorflow-models/universal-sentence-encoder' ) + + // Load the model directly from the module + const model = await sentenceEncoderModule.load() + + // Handle different input types + let textToEmbed: string[] + if (typeof data === 'string') { + if (data.trim() === '') { + return new Array(512).fill(0) + } + textToEmbed = [data] + } else if ( + Array.isArray(data) && + data.every((item) => typeof item === 'string') + ) { + if (data.length === 0 || data.every((item) => item.trim() === '')) { + return new Array(512).fill(0) + } + textToEmbed = data.filter((item) => item.trim() !== '') + if (textToEmbed.length === 0) { + return new Array(512).fill(0) + } + } else { + throw new Error( + 'UniversalSentenceEncoder only supports string or string[] data' + ) + } + + // Get embeddings + const embeddings = await model.embed(textToEmbed) + + // Convert to array and return the first embedding + const embeddingArray = await embeddings.array() + + // Dispose of the tensor to free memory + embeddings.dispose() + + return embeddingArray[0] + } catch (error) { + console.error('Worker error:', error) + throw error } - - // Get embeddings - const embeddings = await model.embed(textToEmbed) - - // Convert to array and return the first embedding - const embeddingArray = await embeddings.array() - - // Dispose of the tensor to free memory - embeddings.dispose() - - return embeddingArray[0] } // Execute the embedding function in a separate thread // Pass the worker implementation as a string to avoid Promise cloning issues - return await executeInThread(workerImplementation.toString(), embedInWorker(data)) + return await executeInThread( + workerImplementation.toString(), + embedInWorker(data) + ) } catch (error) { // If threading fails and fallback is enabled, use the standard embedding function if (options.fallbackToMain) { - console.warn('Threaded embedding failed, falling back to main thread:', error) + console.warn( + 'Threaded embedding failed, falling back to main thread:', + error + ) useFallback = true return standardEmbedding(data) } @@ -287,7 +358,9 @@ export function createThreadedEmbeddingFunction( * Default embedding function * Uses UniversalSentenceEncoder for all text embeddings * TensorFlow.js is required for this to work - * Uses threading when available for better performance + * Uses GPU acceleration via WebGL when available for optimal performance + * Uses threading when available to avoid blocking the main thread + * Falls back to CPU processing when GPU is not available */ export const defaultEmbeddingFunction: EmbeddingFunction = createThreadedEmbeddingFunction({ fallbackToMain: true }) diff --git a/src/utils/workerUtils.ts b/src/utils/workerUtils.ts index 46d80f51..b19ae1d2 100644 --- a/src/utils/workerUtils.ts +++ b/src/utils/workerUtils.ts @@ -20,6 +20,18 @@ export function executeInWebWorker(fnString: string, args: any): Promise { self.index = {} self.index$1 = {} self.index$2 = {} + self.universalSentenceEncoder_esm = {} + self.universalSentenceEncoder = {} + self.tfjs = {} + self.tfjs_core = {} + self.tfjs_backend_cpu = {} + self.tfjs_backend_webgl = {} + + // Additional variables that might be generated by bundlers + self.use = {} + self.tf = {} + self.sentenceEncoder = {} + self.sentenceEncoderModule = {} self.onmessage = async function(e) { try {