feat(src/brainyData, src/utils): enhance embedding efficiency with batch processing and initialize safeguards
- Added batch embedding support with `defaultBatchEmbeddingFunction`, leveraging shared model instances for optimized performance. - Integrated `isInitializing` flag to prevent recursive initialization and ensure smooth concurrent operation handling during `BrainyData` initialization. - Pre-loaded Universal Sentence Encoder in `BrainyData` to prevent delays during embedding. - Introduced fallback mechanisms in embedding initialization for better error resiliency and model reusability. - Updated `addBatch` with support for batchSize and refactored text/vector processing logic for clearer separation and memory management. - Improved GPU and CPU backend selection in Universal Sentence Encoder for compatibility across environments. - Enhanced memory management by cleaning tensors after embedding operations. - Updated README with instructions for batch embedding, threading updates, and GPU/CPU optimizations.
This commit is contained in:
parent
e81979dc84
commit
bba9a0c219
11 changed files with 729 additions and 633 deletions
|
|
@ -13,7 +13,10 @@ import { isThreadingAvailable } from './environment.js'
|
|||
* Lower values indicate higher similarity
|
||||
* Optimized using array methods for Node.js 23.11+
|
||||
*/
|
||||
export const euclideanDistance: DistanceFunction = (a: Vector, b: Vector): number => {
|
||||
export const euclideanDistance: DistanceFunction = (
|
||||
a: Vector,
|
||||
b: Vector
|
||||
): number => {
|
||||
if (a.length !== b.length) {
|
||||
throw new Error('Vectors must have the same dimensions')
|
||||
}
|
||||
|
|
@ -21,7 +24,7 @@ export const euclideanDistance: DistanceFunction = (a: Vector, b: Vector): numbe
|
|||
// Use array.reduce for better performance in Node.js 23.11+
|
||||
const sum = a.reduce((acc, val, i) => {
|
||||
const diff = val - b[i]
|
||||
return acc + (diff * diff)
|
||||
return acc + diff * diff
|
||||
}, 0)
|
||||
|
||||
return Math.sqrt(sum)
|
||||
|
|
@ -33,19 +36,25 @@ export const euclideanDistance: DistanceFunction = (a: Vector, b: Vector): numbe
|
|||
* Range: 0 (identical) to 2 (opposite)
|
||||
* Optimized using array methods for Node.js 23.11+
|
||||
*/
|
||||
export const cosineDistance: DistanceFunction = (a: Vector, b: Vector): number => {
|
||||
export const cosineDistance: DistanceFunction = (
|
||||
a: Vector,
|
||||
b: Vector
|
||||
): number => {
|
||||
if (a.length !== b.length) {
|
||||
throw new Error('Vectors must have the same dimensions')
|
||||
}
|
||||
|
||||
// Use array.reduce to calculate all values in a single pass
|
||||
const { dotProduct, normA, normB } = a.reduce((acc, val, i) => {
|
||||
return {
|
||||
dotProduct: acc.dotProduct + (val * b[i]),
|
||||
normA: acc.normA + (val * val),
|
||||
normB: acc.normB + (b[i] * b[i])
|
||||
}
|
||||
}, { dotProduct: 0, normA: 0, normB: 0 })
|
||||
const { dotProduct, normA, normB } = a.reduce(
|
||||
(acc, val, i) => {
|
||||
return {
|
||||
dotProduct: acc.dotProduct + val * b[i],
|
||||
normA: acc.normA + val * val,
|
||||
normB: acc.normB + b[i] * b[i]
|
||||
}
|
||||
},
|
||||
{ dotProduct: 0, normA: 0, normB: 0 }
|
||||
)
|
||||
|
||||
if (normA === 0 || normB === 0) {
|
||||
return 2 // Maximum distance for zero vectors
|
||||
|
|
@ -61,7 +70,10 @@ export const cosineDistance: DistanceFunction = (a: Vector, b: Vector): number =
|
|||
* Lower values indicate higher similarity
|
||||
* Optimized using array methods for Node.js 23.11+
|
||||
*/
|
||||
export const manhattanDistance: DistanceFunction = (a: Vector, b: Vector): number => {
|
||||
export const manhattanDistance: DistanceFunction = (
|
||||
a: Vector,
|
||||
b: Vector
|
||||
): number => {
|
||||
if (a.length !== b.length) {
|
||||
throw new Error('Vectors must have the same dimensions')
|
||||
}
|
||||
|
|
@ -76,209 +88,211 @@ export const manhattanDistance: DistanceFunction = (a: Vector, b: Vector): numbe
|
|||
* Converted to a distance metric (lower is better)
|
||||
* Optimized using array methods for Node.js 23.11+
|
||||
*/
|
||||
export const dotProductDistance: DistanceFunction = (a: Vector, b: Vector): number => {
|
||||
export const dotProductDistance: DistanceFunction = (
|
||||
a: Vector,
|
||||
b: Vector
|
||||
): number => {
|
||||
if (a.length !== b.length) {
|
||||
throw new Error('Vectors must have the same dimensions')
|
||||
}
|
||||
|
||||
// Use array.reduce for better performance in Node.js 23.11+
|
||||
const dotProduct = a.reduce((sum, val, i) => sum + (val * b[i]), 0)
|
||||
const dotProduct = a.reduce((sum, val, i) => sum + val * b[i], 0)
|
||||
|
||||
// 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
|
||||
* Batch distance calculation
|
||||
* Uses TensorFlow.js with CPU backend for optimized performance
|
||||
*
|
||||
* @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(
|
||||
export async function calculateDistancesBatch(
|
||||
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))
|
||||
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 distanceCalculator = async (args: {
|
||||
queryVector: Vector
|
||||
vectors: Vector[]
|
||||
distanceFnString: string
|
||||
}) => {
|
||||
const { queryVector, vectors, distanceFnString } = args
|
||||
|
||||
// Try to use TensorFlow.js with GPU acceleration if available
|
||||
// Use TensorFlow.js with CPU processing
|
||||
const useTensorFlow = async () => {
|
||||
try {
|
||||
// TensorFlow.js will use its default EPSILON value
|
||||
// TensorFlow.js will use its default EPSILON value
|
||||
|
||||
// Use the importTensorFlow function if available (in worker context)
|
||||
// or directly import TensorFlow.js (in main thread)
|
||||
let tf
|
||||
if (
|
||||
typeof self !== 'undefined' &&
|
||||
typeof self.importTensorFlow === 'function'
|
||||
) {
|
||||
// In worker context, use the importTensorFlow function
|
||||
tf = await self.importTensorFlow()
|
||||
} else {
|
||||
// Dynamically import TensorFlow.js core module and backends
|
||||
const tf = await import('@tensorflow/tfjs-core')
|
||||
tf = await import('@tensorflow/tfjs-core')
|
||||
|
||||
// Import CPU backend as fallback
|
||||
// Import CPU backend
|
||||
await import('@tensorflow/tfjs-backend-cpu')
|
||||
|
||||
let usingGPU = false
|
||||
// Set CPU as the backend
|
||||
await tf.setBackend('cpu')
|
||||
}
|
||||
|
||||
try {
|
||||
// Try to import and use WebGL backend (GPU)
|
||||
await import('@tensorflow/tfjs-backend-webgl')
|
||||
// Convert vectors to tensors
|
||||
const queryTensor = tf.tensor2d([queryVector])
|
||||
const vectorsTensor = tf.tensor2d(vectors)
|
||||
|
||||
// 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')
|
||||
}
|
||||
let distances: number[]
|
||||
|
||||
// Convert vectors to tensors
|
||||
const queryTensor = tf.tensor2d([queryVector])
|
||||
const vectorsTensor = tf.tensor2d(vectors)
|
||||
// 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[]
|
||||
|
||||
let distances: 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()
|
||||
)
|
||||
|
||||
// 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[]
|
||||
const queryNorm = tf.norm(queryTensor, 2, 1)
|
||||
const vectorsNorm = tf.norm(vectorsTensor, 2, 1)
|
||||
|
||||
// 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 normProduct = tf.outerProduct(
|
||||
queryNorm as any,
|
||||
vectorsNorm as any
|
||||
)
|
||||
const cosineSimilarity = tf.div(dotProduct, normProduct)
|
||||
const distancesTensor = tf.sub(tf.scalar(1), cosineSimilarity)
|
||||
|
||||
const queryNorm = tf.norm(queryTensor, 2, 1)
|
||||
const vectorsNorm = tf.norm(vectorsTensor, 2, 1)
|
||||
distances = (await (distancesTensor as any)
|
||||
.squeeze()
|
||||
.array()) as number[]
|
||||
|
||||
const normProduct = tf.outerProduct(queryNorm as any, vectorsNorm as any)
|
||||
const cosineSimilarity = tf.div(dotProduct, normProduct)
|
||||
const distancesTensor = tf.sub(tf.scalar(1), cosineSimilarity)
|
||||
// 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[]
|
||||
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)
|
||||
// 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[]
|
||||
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)
|
||||
// Clean up tensors
|
||||
queryTensor.dispose()
|
||||
vectorsTensor.dispose()
|
||||
dotProduct.dispose()
|
||||
distancesTensor.dispose()
|
||||
} else {
|
||||
// For unknown distance functions, fall back to direct CPU implementation
|
||||
throw new Error(
|
||||
'Unsupported distance function for TensorFlow optimization'
|
||||
)
|
||||
}
|
||||
|
||||
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
|
||||
return {
|
||||
distances
|
||||
}
|
||||
}
|
||||
|
||||
// Try to use TensorFlow.js with GPU acceleration
|
||||
// Try to use TensorFlow.js with CPU optimization
|
||||
try {
|
||||
return await useTensorFlow()
|
||||
} catch (error) {
|
||||
// Fall back to CPU implementation if TensorFlow.js fails
|
||||
// Fall back to direct CPU implementation if TensorFlow.js fails
|
||||
// Recreate the distance function from its string representation
|
||||
const distanceFunction = new Function('return ' + distanceFnString)() as DistanceFunction
|
||||
const distanceFunction = new Function(
|
||||
'return ' + distanceFnString
|
||||
)() as DistanceFunction
|
||||
|
||||
// Calculate distances for all vectors
|
||||
const distances = vectors.map(vector => distanceFunction(queryVector, vector))
|
||||
const distances = vectors.map((vector) =>
|
||||
distanceFunction(queryVector, vector)
|
||||
)
|
||||
|
||||
return {
|
||||
distances,
|
||||
usingGPU: false
|
||||
distances
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 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)
|
||||
}
|
||||
}
|
||||
// Threading is not available, so we'll always use the main thread implementation
|
||||
// This comment is kept for clarity about the removed code
|
||||
|
||||
// If threading is not available or failed, calculate distances in the main thread
|
||||
return vectors.map(vector => distanceFunction(queryVector, vector))
|
||||
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))
|
||||
console.error('Batch distance calculation failed:', error)
|
||||
return vectors.map((vector) => distanceFunction(queryVector, vector))
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue