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

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