soulcraft/brainy
2
0
Fork 0
Code Issues 1 Pull requests Projects Releases 335 Packages Wiki Activity Actions

feat(src/brainyData, src/utils): enhance embedding efficiency with batch processing and initialize safeguards

Browse source 
- 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:
David Snelling 2025-06-27 14:06:59 -07:00
parent e81979dc84
commit bba9a0c219
11 changed files with 729 additions and 633 deletions
Download patch file Download diff file Expand all files Collapse all files

155
src/brainyData.ts
View file

@ -24,7 +24,9 @@ import {
import {
cosineDistance,
defaultEmbeddingFunction,
euclideanDistance
defaultBatchEmbeddingFunction,
euclideanDistance,
cleanupWorkerPools
} from './utils/index.js'
import { NounType, VerbType, GraphNoun } from './types/graphTypes.js'
import {
@ -127,6 +129,7 @@ export class BrainyData<T = any> implements BrainyDataInterface<T> {
private index: HNSWIndex | HNSWIndexOptimized
private storage: StorageAdapter | null = null
private isInitialized = false
private isInitializing = false
private embeddingFunction: EmbeddingFunction
private distanceFunction: DistanceFunction
private requestPersistentStorage: boolean
@ -204,7 +207,48 @@ export class BrainyData<T = any> implements BrainyDataInterface<T> {
return
}
// Prevent recursive initialization
if (this.isInitializing) {
return
}
this.isInitializing = true
try {
// Pre-load the embedding model early to ensure it's always available
// This helps prevent issues with the Universal Sentence Encoder not being loaded
try {
console.log('Pre-loading Universal Sentence Encoder model...')
// Call embedding function directly to avoid circular dependency with embed()
await this.embeddingFunction('')
console.log('Universal Sentence Encoder model loaded successfully')
} catch (embedError) {
console.warn('Failed to pre-load Universal Sentence Encoder:', embedError)
// Try again with a retry mechanism
console.log('Retrying Universal Sentence Encoder initialization...')
try {
// Wait a moment before retrying
await new Promise(resolve => setTimeout(resolve, 1000))
// Try again with a different approach - use the non-threaded version
// This is a fallback in case the threaded version fails
const { createTensorFlowEmbeddingFunction } = await import('./utils/embedding.js')
const fallbackEmbeddingFunction = createTensorFlowEmbeddingFunction()
// Test the fallback embedding function
await fallbackEmbeddingFunction('')
// If successful, replace the embedding function
console.log('Successfully loaded Universal Sentence Encoder with fallback method')
this.embeddingFunction = fallbackEmbeddingFunction
} catch (retryError) {
console.error('All attempts to load Universal Sentence Encoder failed:', retryError)
// Continue initialization even if embedding model fails to load
// The application will need to handle missing embedding functionality
}
}
// Initialize storage if not provided in constructor
if (!this.storage) {
// Combine storage config with requestPersistentStorage for backward compatibility
@ -251,8 +295,10 @@ export class BrainyData<T = any> implements BrainyDataInterface<T> {
}
this.isInitialized = true
this.isInitializing = false
} catch (error) {
console.error('Failed to initialize BrainyData:', error)
this.isInitializing = false
throw new Error(`Failed to initialize BrainyData: ${error}`)
}
}
@ -478,6 +524,7 @@ export class BrainyData<T = any> implements BrainyDataInterface<T> {
forceEmbed?: boolean // Force using the embedding function even if input is a vector
addToRemote?: boolean // Whether to also add to the remote server if connected
concurrency?: number // Maximum number of concurrent operations (default: 4)
batchSize?: number // Maximum number of items to process in a single batch (default: 50)
} = {}
): Promise<string[]> {
await this.ensureInitialized()
@ -488,22 +535,86 @@ export class BrainyData<T = any> implements BrainyDataInterface<T> {
// Default concurrency to 4 if not specified
const concurrency = options.concurrency || 4
// Default batch size to 50 if not specified
const batchSize = options.batchSize || 50
try {
// Process items in batches to control concurrency
// Process items in batches to control concurrency and memory usage
const ids: string[] = []
const itemsToProcess = [...items] // Create a copy to avoid modifying the original array
while (itemsToProcess.length > 0) {
// Take up to 'concurrency' items to process in parallel
const batch = itemsToProcess.splice(0, concurrency)
// Take up to 'batchSize' items to process in a batch
const batch = itemsToProcess.splice(0, batchSize)
// Process this batch in parallel
const batchResults = await Promise.all(
batch.map((item) =>
this.add(item.vectorOrData, item.metadata, options)
)
// Separate items that are already vectors from those that need embedding
const vectorItems: Array<{
vectorOrData: Vector
metadata?: T
index: number
}> = []
const textItems: Array<{
text: string
metadata?: T
index: number
}> = []
// Categorize items
batch.forEach((item, index) => {
if (
Array.isArray(item.vectorOrData) &&
item.vectorOrData.every((val) => typeof val === 'number') &&
!options.forceEmbed
) {
// Item is already a vector
vectorItems.push({
vectorOrData: item.vectorOrData,
metadata: item.metadata,
index
})
} else if (typeof item.vectorOrData === 'string') {
// Item is text that needs embedding
textItems.push({
text: item.vectorOrData,
metadata: item.metadata,
index
})
} else {
// For now, treat other types as text
// In a more complete implementation, we might handle other types differently
const textRepresentation = String(item.vectorOrData)
textItems.push({
text: textRepresentation,
metadata: item.metadata,
index
})
}
})
// Process vector items (already embedded)
const vectorPromises = vectorItems.map(item =>
this.add(item.vectorOrData, item.metadata, options)
)
// Process text items in a single batch embedding operation
let textPromises: Promise<string>[] = []
if (textItems.length > 0) {
// Extract just the text for batch embedding
const texts = textItems.map(item => item.text)
// Perform batch embedding
const embeddings = await defaultBatchEmbeddingFunction(texts)
// Add each item with its embedding
textPromises = textItems.map((item, i) =>
this.add(embeddings[i], item.metadata, { ...options, forceEmbed: false })
)
}
// Combine all promises
const batchResults = await Promise.all([...vectorPromises, ...textPromises])
// Add the results to our ids array
ids.push(...batchResults)
}
@ -1729,7 +1840,28 @@ export class BrainyData<T = any> implements BrainyDataInterface<T> {
* Ensure the database is initialized
*/
private async ensureInitialized(): Promise<void> {
if (!this.isInitialized) {
if (this.isInitialized) {
return
}
if (this.isInitializing) {
// If initialization is already in progress, wait for it to complete
// by polling the isInitialized flag
let attempts = 0
const maxAttempts = 100 // Prevent infinite loop
const delay = 50 // ms
while (this.isInitializing && !this.isInitialized && attempts < maxAttempts) {
await new Promise(resolve => setTimeout(resolve, delay))
attempts++
}
if (!this.isInitialized) {
// If still not initialized after waiting, try to initialize again
await this.init()
}
} else {
// Normal case - not initialized and not initializing
await this.init()
}
}
@ -1838,6 +1970,9 @@ export class BrainyData<T = any> implements BrainyDataInterface<T> {
await this.disconnectFromRemoteServer()
}
// Clean up worker pools to release resources
cleanupWorkerPools()
// Additional cleanup could be added here in the future
this.isInitialized = false

15
src/global.d.ts vendored
View file

@ -5,11 +5,16 @@
// Extend the globalThis interface to include the __ENV__ property
declare global {
var __ENV__: {
isBrowser: boolean;
isNode: string | false;
isServerless: boolean;
};
isBrowser: boolean
isNode: string | false
isServerless: boolean
}
// Window interface extensions (worker-specific functions removed)
interface Window {
importTensorFlow?: () => Promise<any>
}
}
// This export is needed to make this file a module
export {};
export {}

52
src/hnsw/hnswIndex.ts
View file

@ -10,7 +10,7 @@ import {
Vector,
VectorDocument
} from '../coreTypes.js'
import { euclideanDistance, calculateDistancesWithGPU } from '../utils/index.js'
import { euclideanDistance, calculateDistancesBatch } from '../utils/index.js'
import { executeInThread } from '../utils/workerUtils.js'
// Default HNSW parameters
@ -83,7 +83,7 @@ export class HNSWIndex {
const vectorsOnly = vectors.map((item) => item.vector)
// Use GPU-accelerated distance calculation when possible
const distances = await calculateDistancesWithGPU(
const distances = await calculateDistancesBatch(
queryVector,
vectorsOnly,
this.distanceFunction
@ -96,51 +96,15 @@ export class HNSWIndex {
}))
} catch (error) {
console.error(
'Error in GPU-accelerated distance calculation, falling back to threaded CPU:',
'Error in GPU-accelerated distance calculation, falling back to sequential processing:',
error
)
// Fall back to threaded CPU processing if GPU acceleration fails
// Function to be executed in a worker thread
const distanceCalculator = (args: {
queryVector: Vector
vectors: Array<{ id: string; vector: Vector }>
distanceFnString: string
}) => {
const { queryVector, vectors, distanceFnString } = args
// Recreate the distance function from its string representation
const distanceFunction = new Function(
'return ' + distanceFnString
)() as DistanceFunction
// Calculate distances for all items
return vectors.map((item) => ({
id: item.id,
distance: distanceFunction(queryVector, item.vector)
}))
}
try {
// Convert the distance function to a string for serialization
const distanceFnString = this.distanceFunction.toString()
// Execute the distance calculation in a separate thread
return await executeInThread<Array<{ id: string; distance: number }>>(
distanceCalculator.toString(),
{ queryVector, vectors, distanceFnString }
)
} catch (threadError) {
console.error(
'Error in threaded distance calculation, falling back to sequential:',
threadError
)
// Fall back to sequential processing if both GPU and threaded execution fail
return vectors.map((item) => ({
id: item.id,
distance: this.distanceFunction(queryVector, item.vector)
}))
}
// Fall back to sequential processing if GPU acceleration fails
return vectors.map((item) => ({
id: item.id,
distance: this.distanceFunction(queryVector, item.vector)
}))
}
}

310
src/utils/distance.ts
View file

@ -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))
}
}

452
src/utils/embedding.ts
View file

@ -10,14 +10,15 @@ import { executeInThread } from './workerUtils.js'
* 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.
* This implementation attempts to use GPU processing when available for better performance,
* falling back to CPU processing for compatibility across all environments.
*/
export class UniversalSentenceEncoder implements EmbeddingModel {
private model: any = null
private initialized = false
private tf: any = null
private use: any = null
private backend: string = 'cpu' // Default to CPU
/**
* Initialize the embedding model
@ -47,36 +48,59 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
// Use type assertions to tell TypeScript these modules exist
this.tf = await import('@tensorflow/tfjs-core')
// Import CPU and WebGL backends
// Import CPU backend (always needed as fallback)
await import('@tensorflow/tfjs-backend-cpu')
// Try to import WebGL backend for GPU acceleration in browser environments
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')
if (typeof window !== 'undefined') {
await import('@tensorflow/tfjs-backend-webgl')
// Check if WebGL is available using setBackend instead of findBackend
try {
if (this.tf.setBackend) {
await this.tf.setBackend('webgl')
this.backend = 'webgl'
console.log('Using WebGL backend for TensorFlow.js')
} else {
console.warn(
'tf.setBackend is not available, falling back to CPU'
)
}
} catch (e) {
console.warn('WebGL backend not available, falling back to CPU:', e)
this.backend = 'cpu'
}
}
} catch (err) {
console.warn(
'WebGL backend failed to initialize, using CPU backend:',
err
)
await this.tf.setBackend('cpu')
} catch (error) {
console.warn('WebGL backend not available, falling back to CPU:', error)
this.backend = 'cpu'
}
// Set the backend
if (this.tf.setBackend) {
await this.tf.setBackend(this.backend)
}
this.use = await import('@tensorflow-models/universal-sentence-encoder')
// Log the module structure to help with debugging
console.log(
'Universal Sentence Encoder module structure in main thread:',
Object.keys(this.use),
this.use.default ? Object.keys(this.use.default) : 'No default export'
)
// Try to find the load function in different possible module structures
const loadFunction = findUSELoadFunction(this.use)
if (!loadFunction) {
throw new Error(
'Could not find Universal Sentence Encoder load function'
)
}
// Load the model
this.model = await this.use.load()
this.model = await loadFunction()
this.initialized = true
// Restore original console.warn
@ -132,6 +156,10 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
// 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(
@ -144,6 +172,70 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
}
}
/**
* Embed multiple texts into vectors using Universal Sentence Encoder
* This is more efficient than calling embed() multiple times
* @param dataArray Array of texts to embed
* @returns Array of embedding vectors
*/
public async embedBatch(dataArray: string[]): Promise<Vector[]> {
if (!this.initialized) {
await this.init()
}
try {
// Handle empty array case
if (dataArray.length === 0) {
return []
}
// Filter out empty strings and handle edge cases
const textToEmbed = dataArray.filter(
(text: string) => typeof text === 'string' && text.trim() !== ''
)
// If all strings were empty, return appropriate zero vectors
if (textToEmbed.length === 0) {
return dataArray.map(() => new Array(512).fill(0))
}
// Get embeddings for all texts in a single batch operation
const embeddings = await this.model.embed(textToEmbed)
// Convert to array
const embeddingArray = await embeddings.array()
// Dispose of the tensor to free memory
embeddings.dispose()
// Map the results back to the original array order
const results: Vector[] = []
let embeddingIndex = 0
for (let i = 0; i < dataArray.length; i++) {
const text = dataArray[i]
if (typeof text === 'string' && text.trim() !== '') {
// Use the embedding for non-empty strings
results.push(embeddingArray[embeddingIndex])
embeddingIndex++
} else {
// Use a zero vector for empty strings
results.push(new Array(512).fill(0))
}
}
return results
} catch (error) {
console.error(
'Failed to batch embed text with Universal Sentence Encoder:',
error
)
throw new Error(
`Failed to batch embed text with Universal Sentence Encoder: ${error}`
)
}
}
/**
* Dispose of the model resources
*/
@ -162,6 +254,105 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
}
}
/**
* Helper function to load the Universal Sentence Encoder model
* This tries multiple approaches to find the correct load function
* @param sentenceEncoderModule The imported module
* @returns The load function or null if not found
*/
function findUSELoadFunction(sentenceEncoderModule: any): Function | null {
// Log the module structure for debugging
console.log(
'Universal Sentence Encoder module structure:',
Object.keys(sentenceEncoderModule),
sentenceEncoderModule.default
? Object.keys(sentenceEncoderModule.default)
: 'No default export'
)
let loadFunction = null
// Try sentenceEncoderModule.load first (direct export)
if (
sentenceEncoderModule.load &&
typeof sentenceEncoderModule.load === 'function'
) {
loadFunction = sentenceEncoderModule.load
console.log('Using sentenceEncoderModule.load')
}
// Then try sentenceEncoderModule.default.load (default export)
else if (
sentenceEncoderModule.default &&
sentenceEncoderModule.default.load &&
typeof sentenceEncoderModule.default.load === 'function'
) {
loadFunction = sentenceEncoderModule.default.load
console.log('Using sentenceEncoderModule.default.load')
}
// Try sentenceEncoderModule.default directly if it's a function
else if (
sentenceEncoderModule.default &&
typeof sentenceEncoderModule.default === 'function'
) {
loadFunction = sentenceEncoderModule.default
console.log('Using sentenceEncoderModule.default as function')
}
// Try sentenceEncoderModule directly if it's a function
else if (typeof sentenceEncoderModule === 'function') {
loadFunction = sentenceEncoderModule
console.log('Using sentenceEncoderModule as function')
}
// Try additional common patterns
else if (
sentenceEncoderModule.UniversalSentenceEncoder &&
typeof sentenceEncoderModule.UniversalSentenceEncoder.load === 'function'
) {
loadFunction = sentenceEncoderModule.UniversalSentenceEncoder.load
console.log('Using sentenceEncoderModule.UniversalSentenceEncoder.load')
} else if (
sentenceEncoderModule.default &&
sentenceEncoderModule.default.UniversalSentenceEncoder &&
typeof sentenceEncoderModule.default.UniversalSentenceEncoder.load ===
'function'
) {
loadFunction = sentenceEncoderModule.default.UniversalSentenceEncoder.load
console.log(
'Using sentenceEncoderModule.default.UniversalSentenceEncoder.load'
)
}
// Try to find the load function in the module's properties
else {
// Look for any property that might be a load function
for (const key in sentenceEncoderModule) {
if (typeof sentenceEncoderModule[key] === 'function') {
// Check if the function name or key contains 'load'
const fnName = sentenceEncoderModule[key].name || key;
if (fnName.toLowerCase().includes('load')) {
loadFunction = sentenceEncoderModule[key];
console.log(`Using sentenceEncoderModule.${key} as load function`);
break;
}
}
// Also check nested objects
else if (typeof sentenceEncoderModule[key] === 'object' && sentenceEncoderModule[key] !== null) {
for (const nestedKey in sentenceEncoderModule[key]) {
if (typeof sentenceEncoderModule[key][nestedKey] === 'function') {
const fnName = sentenceEncoderModule[key][nestedKey].name || nestedKey;
if (fnName.toLowerCase().includes('load')) {
loadFunction = sentenceEncoderModule[key][nestedKey];
console.log(`Using sentenceEncoderModule.${key}.${nestedKey} as load function`);
break;
}
}
}
if (loadFunction) break;
}
}
}
return loadFunction;
}
/**
* Create an embedding function from an embedding model
* @param model Embedding model to use
@ -177,21 +368,22 @@ export function createEmbeddingFunction(
/**
* Creates a TensorFlow-based Universal Sentence Encoder embedding function
* This is the required embedding function for all text embeddings
* Uses a shared model instance for better performance across multiple calls
*/
export function createTensorFlowEmbeddingFunction(): EmbeddingFunction {
// Create a single shared instance of the model
const model = new UniversalSentenceEncoder()
let modelInitialized = false
// Create a single shared instance of the model that persists across all embedding calls
const sharedModel = new UniversalSentenceEncoder()
let sharedModelInitialized = false
export function createTensorFlowEmbeddingFunction(): EmbeddingFunction {
return async (data: any): Promise<Vector> => {
try {
// Initialize the model if it hasn't been initialized yet
if (!modelInitialized) {
await model.init()
modelInitialized = true
if (!sharedModelInitialized) {
await sharedModel.init()
sharedModelInitialized = true
}
return await model.embed(data)
return await sharedModel.embed(data)
} catch (error) {
console.error('Failed to use TensorFlow embedding:', error)
throw new Error(
@ -202,165 +394,59 @@ export function createTensorFlowEmbeddingFunction(): EmbeddingFunction {
}
/**
* Creates a TensorFlow-based Universal Sentence Encoder embedding function that runs in a separate thread
* 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 with GPU acceleration when available
* Default embedding function
* Uses UniversalSentenceEncoder for all text embeddings
* TensorFlow.js is required for this to work
* Uses CPU for compatibility
*/
export function createThreadedEmbeddingFunction(
options: { fallbackToMain?: boolean } = {}
): EmbeddingFunction {
// Create a standard embedding function to use as fallback
const standardEmbedding = createTensorFlowEmbeddingFunction()
export const defaultEmbeddingFunction: EmbeddingFunction =
createTensorFlowEmbeddingFunction()
// Flag to track if we've fallen back to main thread
let useFallback = false
/**
* Default batch embedding function
* Uses UniversalSentenceEncoder for all text embeddings
* TensorFlow.js is required for this to work
* Processes all items in a single batch operation
* Uses a shared model instance for better performance across multiple calls
*/
// Create a single shared instance of the model that persists across function calls
const sharedBatchModel = new UniversalSentenceEncoder()
let sharedBatchModelInitialized = false
return async (data: any): Promise<Vector> => {
// If we've already determined that threading doesn't work, use the fallback
if (useFallback) {
return standardEmbedding(data)
export const defaultBatchEmbeddingFunction: (
dataArray: string[]
) => Promise<Vector[]> = async (dataArray: string[]): Promise<Vector[]> => {
try {
// Initialize the model if it hasn't been initialized yet
if (!sharedBatchModelInitialized) {
await sharedBatchModel.init()
sharedBatchModelInitialized = true
}
try {
// Function to be executed in a worker thread
// This must be a regular function (not async) to avoid Promise cloning issues
const embedInWorker = (inputData: any) => {
// Return a plain object with the input data
// All async operations will be performed inside the worker
return { data: inputData }
}
// Worker implementation function that will be stringified and run in the worker
const workerImplementation = async ({ data }: { data: any }) => {
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
// TensorFlow.js will use its default EPSILON value
// 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')
}
}
// 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
}
}
// 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)
)
} 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
)
useFallback = true
return standardEmbedding(data)
}
// Otherwise, propagate the error
throw new Error(`Threaded embedding failed: ${error}`)
}
return await sharedBatchModel.embedBatch(dataArray)
} catch (error) {
console.error('Failed to use TensorFlow batch embedding:', error)
throw new Error(
`Universal Sentence Encoder batch embedding failed: ${error}`
)
}
}
/**
* Default embedding function
* Uses UniversalSentenceEncoder for all text embeddings
* TensorFlow.js is required for this to work
* 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
* Creates an embedding function that runs in a separate thread
* This is a wrapper around createEmbeddingFunction that uses executeInThread
* @param model Embedding model to use
*/
export const defaultEmbeddingFunction: EmbeddingFunction =
createThreadedEmbeddingFunction({ fallbackToMain: true })
export function createThreadedEmbeddingFunction(
model: EmbeddingModel
): EmbeddingFunction {
const embeddingFunction = createEmbeddingFunction(model)
return async (data: any): Promise<Vector> => {
// Convert the embedding function to a string
const fnString = embeddingFunction.toString()
// Execute the embedding function in a "thread" (main thread in this implementation)
return await executeInThread<Vector>(fnString, data)
}
}

5
src/utils/environment.ts
View file

@ -39,7 +39,7 @@ export function areWebWorkersAvailable(): boolean {
*/
export function areWorkerThreadsAvailable(): boolean {
if (!isNode()) return false;
try {
// Dynamic import to avoid errors in browser environments
require('worker_threads');
@ -51,7 +51,8 @@ export function areWorkerThreadsAvailable(): boolean {
/**
* Determine if threading is available in the current environment
* Always returns false since multithreading has been removed
*/
export function isThreadingAvailable(): boolean {
return areWebWorkersAvailable() || areWorkerThreadsAvailable();
return false;
}

1
src/utils/index.ts
View file

@ -1,2 +1,3 @@
export * from './distance.js'
export * from './embedding.js'
export * from './workerUtils.js'

162
src/utils/workerUtils.ts
View file

@ -1,162 +1,30 @@
/**
* Utility functions for working with Web Workers and Worker Threads
* Utility functions for executing functions (without Web Workers or Worker Threads)
* This is a replacement for the original multithreaded implementation
*/
import { isNode, isBrowser } from './environment.js'
/**
* Execute a function in a Web Worker (browser environment)
*
* @param fnString The function to execute as a string
* @param args The arguments to pass to the function
* @returns A promise that resolves with the result of the function
*/
export function executeInWebWorker<T>(fnString: string, args: any): Promise<T> {
return new Promise((resolve, reject) => {
// Create a blob URL for the worker script
const workerScript = `
// Define global variables to prevent "is not defined" errors
// This ensures that any minified variable names used by bundlers are available
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.onmessage = async function(e) {
try {
const fn = ${fnString}
const result = await fn(e.data)
self.postMessage({ success: true, data: result })
} catch (error) {
self.postMessage({
success: false,
error: error instanceof Error ? error.message : String(error)
})
}
}
`
const blob = new Blob([workerScript], { type: 'application/javascript' })
const blobURL = URL.createObjectURL(blob)
// Create a worker
const worker = new Worker(blobURL)
// Set up message handling
worker.onmessage = function (
e: MessageEvent<{ success: boolean; data: T; error?: string }>
) {
URL.revokeObjectURL(blobURL) // Clean up
worker.terminate() // Terminate the worker
if (e.data.success) {
resolve(e.data.data)
} else {
reject(new Error(e.data.error))
}
}
worker.onerror = function (error: ErrorEvent) {
URL.revokeObjectURL(blobURL) // Clean up
worker.terminate() // Terminate the worker
reject(error)
}
// Start the worker
worker.postMessage(args)
})
}
/**
* Execute a function in a Worker Thread (Node.js environment)
*
* @param fnString The function to execute as a string
* @param args The arguments to pass to the function
* @returns A promise that resolves with the result of the function
*/
export function executeInWorkerThread<T>(
fnString: string,
args: any
): Promise<T> {
return new Promise((resolve, reject) => {
try {
// Dynamic import to avoid errors in browser environments
const { Worker } = require('worker_threads')
// Create a worker script
const workerScript = `
const { parentPort } = require('worker_threads')
parentPort.once('message', async (data) => {
try {
const fn = ${fnString}
const result = await fn(data)
parentPort.postMessage({ success: true, data: result })
} catch (error) {
parentPort.postMessage({
success: false,
error: error instanceof Error ? error.message : String(error)
})
}
})
`
// Create a worker
const worker = new Worker(workerScript, { eval: true })
// Set up message handling
worker.on(
'message',
(data: { success: boolean; data: T; error?: string }) => {
worker.terminate() // Terminate the worker
if (data.success) {
resolve(data.data)
} else {
reject(new Error(data.error))
}
}
)
worker.on('error', (error: Error) => {
worker.terminate() // Terminate the worker
reject(error)
})
// Start the worker
worker.postMessage(args)
} catch (error) {
reject(error)
}
})
}
/**
* Execute a function in a separate thread based on the environment
* Execute a function directly in the main thread
*
* @param fnString The function to execute as a string
* @param args The arguments to pass to the function
* @returns A promise that resolves with the result of the function
*/
export function executeInThread<T>(fnString: string, args: any): Promise<T> {
if (isBrowser()) {
return executeInWebWorker<T>(fnString, args)
} else if (isNode()) {
return executeInWorkerThread<T>(fnString, args)
} else {
// Fall back to executing in the main thread
try {
// Parse the function from string and execute it
const fn = new Function('return ' + fnString)()
return Promise.resolve(fn(args) as T)
} catch (error) {
return Promise.reject(error)
}
}
/**
* No-op function for backward compatibility
* This function does nothing since there are no worker pools to clean up
*/
export function cleanupWorkerPools(): void {
// No-op function
console.log('Worker pools cleanup called (no-op in non-threaded version)')
}