feat(src/utils): add GPU acceleration and improve threading for embeddings and distance calculations
- Integrated GPU acceleration via WebGL in `embedding.ts` and `distance.ts` for optimized performance. - Enhanced fallback mechanisms for CPU processing when GPU is unavailable or fails to initialize. - Added threading support for embedding and distance calculations in `embedding.ts` and `distance.ts` to improve performance and avoid blocking the main thread. - Refactored TensorFlow.js imports to dynamically load modules (`@tensorflow/tfjs-core`, `@tensorflow/tfjs-backend-webgl`, `@tensorflow/tfjs-backend-cpu`) for modular dependency usage. - Introduced GPU-accelerated batch distance calculations in `distance.ts` with appropriate error handling. - Applied consistent formatting to improve code readability and maintainability while adhering to project style guidelines.
This commit is contained in:
parent
bd123c4bb9
commit
6a8d044970
3 changed files with 348 additions and 65 deletions
|
|
@ -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<number[]> {
|
||||
// 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))
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<Vector>(workerImplementation.toString(), embedInWorker(data))
|
||||
return await executeInThread<Vector>(
|
||||
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 })
|
||||
|
|
|
|||
|
|
@ -20,6 +20,18 @@ export function executeInWebWorker<T>(fnString: string, args: any): Promise<T> {
|
|||
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 {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue