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

**feat(docs): add comprehensive documentation for model bundling and robust loading**

Browse source 
- Introduced new documentation files under `docs/`:
  - `model-bundling-analysis.md`: Provides detailed analysis of current, bundled, hybrid, and dynamic model loading approaches, including pros, cons, and recommendations.
  - `model-management.md`: Explains how Brainy manages Universal Sentence Encoder models, including setup, usage, and troubleshooting.
  - `optional-model-bundling.md`: Details the `@soulcraft/brainy-models` package for offline reliability with pre-bundled models.

- Added `src/utils/robustModelLoader.ts`:
  - Implements enhanced model loading with retry mechanisms, timeout handling, fallback URLs, and optional local model bundling.
  - Supports Node.js and browser environments with exponential backoff logic.

- Key Updates:
  - **Hybrid Loading Strategy**: Recommended for balancing reliability and flexibility via hybrid online/offline mechanisms.
  - **Enhanced Fallback Scenarios**: Robust loader improves network-dependent reliability for embedding workflows.
  - **Offline Reliability Support**: Optional model bundling eliminates dependency on external services, supporting air-gapped and edge environments.

**Purpose**: Introduce a hybrid model loading approach with robust options for
This commit is contained in:
David Snelling 2025-08-01 15:35:29 -07:00
parent 563b983fcc
commit 42571c5883
6 changed files with 1098 additions and 214 deletions
Download patch file Download diff file Expand all files Collapse all files

341
src/utils/embedding.ts
View file

@ -5,6 +5,12 @@
import { EmbeddingFunction, EmbeddingModel, Vector } from '../coreTypes.js'
import { executeInThread } from './workerUtils.js'
import { isBrowser } from './environment.js'
import {
RobustModelLoader,
ModelLoadOptions,
createRobustModelLoader,
getUniversalSentenceEncoderFallbacks
} from './robustModelLoader.js'
/**
* TensorFlow Universal Sentence Encoder embedding model
@ -14,6 +20,11 @@ import { isBrowser } from './environment.js'
* This implementation attempts to use GPU processing when available for better performance,
* falling back to CPU processing for compatibility across all environments.
*/
export interface UniversalSentenceEncoderOptions extends ModelLoadOptions {
/** Whether to enable verbose logging */
verbose?: boolean
}
export class UniversalSentenceEncoder implements EmbeddingModel {
private model: any = null
private initialized = false
@ -21,13 +32,26 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
private use: any = null
private backend: string = 'cpu' // Default to CPU
private verbose: boolean = true // Whether to log non-essential messages
private robustLoader: RobustModelLoader
/**
* Create a new UniversalSentenceEncoder instance
* @param options Configuration options
* @param options Configuration options including reliability settings
*/
constructor(options: { verbose?: boolean } = {}) {
constructor(options: UniversalSentenceEncoderOptions = {}) {
this.verbose = options.verbose !== undefined ? options.verbose : true
// Create robust model loader with enhanced reliability features
this.robustLoader = createRobustModelLoader({
maxRetries: options.maxRetries ?? 3,
initialRetryDelay: options.initialRetryDelay ?? 1000,
maxRetryDelay: options.maxRetryDelay ?? 30000,
timeout: options.timeout ?? 60000,
useExponentialBackoff: options.useExponentialBackoff ?? true,
fallbackUrls: options.fallbackUrls ?? getUniversalSentenceEncoderFallbacks(),
verbose: this.verbose,
preferLocalModel: options.preferLocalModel ?? true
})
}
/**
@ -116,140 +140,36 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
}
/**
* Load the Universal Sentence Encoder model with retry logic
* This helps handle network failures and JSON parsing errors from TensorFlow Hub
* @param loadFunction The function to load the model
* @param maxRetries Maximum number of retry attempts
* @param baseDelay Base delay in milliseconds for exponential backoff
* Load the Universal Sentence Encoder model with robust retry and fallback mechanisms
* @param loadFunction The function to load the model from TensorFlow Hub
*/
private async loadModelWithRetry(
loadFunction: () => Promise<EmbeddingModel>,
maxRetries: number = 3,
baseDelay: number = 1000
private async loadModelFromLocal(
loadFunction: () => Promise<EmbeddingModel>
): Promise<EmbeddingModel> {
let lastError: Error | null = null
// Define alternative model URLs to try if the default one fails
const alternativeLoadFunctions: Array<() => Promise<EmbeddingModel>> = []
// Try to create alternative load functions using different model URLs
if (this.use) {
// Add alternative model URLs to try
const alternativeUrls = [
'https://storage.googleapis.com/tfjs-models/savedmodel/universal_sentence_encoder/model.json',
'https://tfhub.dev/tensorflow/tfjs-model/universal-sentence-encoder-lite/1/default/1/model.json',
'https://tfhub.dev/tensorflow/tfjs-model/universal-sentence-encoder/1/default/1/model.json',
'https://tfhub.dev/tensorflow/tfjs-model/universal-sentence-encoder/1/default/1',
'https://tfhub.dev/tensorflow/universal-sentence-encoder/4',
'https://tfhub.dev/tensorflow/universal-sentence-encoder/4/default/1/model.json'
]
// Create load functions for each alternative URL
for (const url of alternativeUrls) {
if (this.use.load) {
alternativeLoadFunctions.push(() => this.use!.load(url))
} else if (this.use.default && this.use.default.load) {
alternativeLoadFunctions.push(() => this.use!.default.load(url))
}
this.logger('log', 'Loading Universal Sentence Encoder model with robust loader...')
try {
// Use the robust model loader to handle all retry logic, timeouts, and fallbacks
const model = await this.robustLoader.loadModel(
loadFunction,
'universal-sentence-encoder'
)
this.logger('log', 'Successfully loaded Universal Sentence Encoder model')
return model
} catch (error) {
const errorMessage = error instanceof Error ? error.message : String(error)
this.logger('error', `Failed to load Universal Sentence Encoder model: ${errorMessage}`)
// Log loading statistics for debugging
const stats = this.robustLoader.getLoadingStats()
if (Object.keys(stats).length > 0) {
this.logger('log', 'Loading attempt statistics:', stats)
}
throw error
}
// First try with the original load function
for (let attempt = 0; attempt <= maxRetries; attempt++) {
try {
this.logger(
'log',
attempt === 0
? 'Loading Universal Sentence Encoder model...'
: `Retrying Universal Sentence Encoder model loading (attempt ${attempt + 1}/${maxRetries + 1})...`
)
const model = await loadFunction()
if (attempt > 0) {
this.logger(
'log',
'Universal Sentence Encoder model loaded successfully after retry'
)
}
return model
} catch (error) {
lastError = error as Error
const errorMessage = lastError.message || String(lastError)
// Check if this is a network-related error that might benefit from retry
const isRetryableError =
errorMessage.includes('Failed to parse model JSON') ||
errorMessage.includes('Failed to fetch') ||
errorMessage.includes('Network error') ||
errorMessage.includes('ENOTFOUND') ||
errorMessage.includes('ECONNRESET') ||
errorMessage.includes('ETIMEDOUT') ||
errorMessage.includes('JSON') ||
errorMessage.includes('model.json') ||
errorMessage.includes('byte length') ||
errorMessage.includes('tensor should have') ||
errorMessage.includes('shape') ||
errorMessage.includes('dimensions')
if (attempt < maxRetries && isRetryableError) {
const delay = baseDelay * Math.pow(2, attempt) // Exponential backoff
this.logger(
'warn',
`Universal Sentence Encoder model loading failed (attempt ${attempt + 1}): ${errorMessage}. Retrying in ${delay}ms...`
)
await new Promise((resolve) => setTimeout(resolve, delay))
} else {
// Either we've exhausted retries or this is not a retryable error
if (attempt >= maxRetries) {
this.logger(
'warn',
`Universal Sentence Encoder model loading failed after ${maxRetries + 1} attempts. Last error: ${errorMessage}. Trying alternative URLs...`
)
// Try alternative URLs if available
if (alternativeLoadFunctions.length > 0) {
for (let i = 0; i < alternativeLoadFunctions.length; i++) {
try {
this.logger(
'log',
`Trying alternative model URL ${i + 1}/${alternativeLoadFunctions.length}...`
)
const model = await alternativeLoadFunctions[i]()
this.logger(
'log',
`Successfully loaded Universal Sentence Encoder from alternative URL ${i + 1}`
)
return model
} catch (altError) {
this.logger(
'warn',
`Failed to load from alternative URL ${i + 1}: ${altError}`
)
// Continue to the next alternative
}
}
}
// If we get here, all alternatives failed
this.logger(
'error',
`Universal Sentence Encoder model loading failed after trying all alternatives. Last error: ${errorMessage}`
)
} else {
this.logger(
'error',
`Universal Sentence Encoder model loading failed with non-retryable error: ${errorMessage}`
)
}
throw lastError
}
}
}
// This should never be reached, but just in case
throw lastError || new Error('Unknown error during model loading')
}
/**
@ -277,8 +197,6 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
// Add polyfills for TensorFlow.js compatibility
this.addServerCompatibilityPolyfills()
// TensorFlow.js will use its default EPSILON value
// CRITICAL: Ensure TextEncoder/TextDecoder are available before TensorFlow.js loads
try {
// Get the appropriate global object for the current environment
@ -303,7 +221,8 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
globalObj.TextEncoder = util.TextEncoder
}
if (!globalObj.TextDecoder) {
globalObj.TextDecoder = util.TextDecoder as unknown as typeof TextDecoder
globalObj.TextDecoder =
util.TextDecoder as unknown as typeof TextDecoder
}
}
} catch (utilError) {
@ -363,7 +282,9 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
} catch (error) {
this.logger('error', 'Failed to initialize TensorFlow.js:', error)
// No fallback allowed - throw error
throw new Error(`Universal Sentence Encoder initialization failed: ${error}`)
throw new Error(
`Universal Sentence Encoder initialization failed: ${error}`
)
}
// Set the backend
@ -375,13 +296,18 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
const loadFunction = findUSELoadFunction(this.use)
if (!loadFunction) {
this.logger('error', 'Could not find Universal Sentence Encoder load function')
throw new Error('Universal Sentence Encoder load function not found. Fallback mechanisms are not allowed.')
this.logger(
'error',
'Could not find Universal Sentence Encoder load function'
)
throw new Error(
'Universal Sentence Encoder load function not found. Fallback mechanisms are not allowed.'
)
}
try {
// Load the model with retry logic for network failures
this.model = await this.loadModelWithRetry(loadFunction)
// Load the model from local files first, falling back to default loading if necessary
this.model = await this.loadModelFromLocal(loadFunction)
this.initialized = true
} catch (modelError) {
this.logger(
@ -390,7 +316,9 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
modelError
)
// No fallback allowed - throw error
throw new Error(`Universal Sentence Encoder model loading failed: ${modelError}`)
throw new Error(
`Universal Sentence Encoder model loading failed: ${modelError}`
)
}
// Restore original console.warn
@ -402,7 +330,9 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
error
)
// No fallback allowed - throw error
throw new Error(`Universal Sentence Encoder initialization failed: ${error}`)
throw new Error(
`Universal Sentence Encoder initialization failed: ${error}`
)
}
}
@ -410,15 +340,6 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
* Embed text into a vector using Universal Sentence Encoder
* @param data Text to embed
*/
/**
* This method has been removed as we should always use Universal Sentence Encoder
* and never fall back to alternative vector generation methods
* @deprecated
*/
private generateFallbackVector(text: string): Vector {
throw new Error('Fallback vector generation is not allowed. Universal Sentence Encoder must be used for all embeddings.')
}
public async embed(data: string | string[]): Promise<Vector> {
if (!this.initialized) {
await this.init()
@ -430,7 +351,7 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
if (typeof data === 'string') {
// Handle empty string case
if (data.trim() === '') {
// Return a zero vector of appropriate dimension (512 is the default for USE)
// Return a zero vector of 512 dimensions (standard for Universal Sentence Encoder)
return new Array(512).fill(0)
}
textToEmbed = [data]
@ -453,11 +374,9 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
)
}
// Ensure the model is available - no fallbacks allowed
// Ensure the model is available
if (!this.model) {
throw new Error(
'Universal Sentence Encoder model is not available. Fallback mechanisms are not allowed.'
)
throw new Error('Universal Sentence Encoder model is not available')
}
// Get embeddings
@ -471,14 +390,14 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
// Get the first embedding
let embedding = embeddingArray[0]
// Ensure the embedding is exactly 512 dimensions
// Always ensure the embedding is exactly 512 dimensions
if (embedding.length !== 512) {
this.logger(
'warn',
`Embedding dimension mismatch: expected 512, got ${embedding.length}. Standardizing...`
`Embedding dimension mismatch: expected 512, got ${embedding.length}. Standardizing to 512 dimensions.`
)
// If the embedding is too short, pad with zeros
if (embedding.length < 512) {
const paddedEmbedding = new Array(512).fill(0)
@ -486,27 +405,15 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
paddedEmbedding[i] = embedding[i]
}
embedding = paddedEmbedding
}
}
// If the embedding is too long, truncate
else if (embedding.length > 512) {
// Special handling for 1536-dimensional vectors (common with newer models)
if (embedding.length === 1536) {
// Take every third value to reduce from 1536 to 512
const reducedEmbedding = new Array(512).fill(0)
for (let i = 0; i < 512; i++) {
reducedEmbedding[i] = embedding[i * 3]
}
embedding = reducedEmbedding
} else {
// For other dimensions, just truncate
embedding = embedding.slice(0, 512)
}
embedding = embedding.slice(0, 512)
}
}
return embedding
} catch (error) {
// No fallback - throw the error
this.logger(
'error',
'Failed to embed text with Universal Sentence Encoder:',
@ -543,11 +450,9 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
return dataArray.map(() => new Array(512).fill(0))
}
// Ensure the model is available - no fallbacks allowed
// Ensure the model is available
if (!this.model) {
throw new Error(
'Universal Sentence Encoder model is not available. Fallback mechanisms are not allowed.'
)
throw new Error('Universal Sentence Encoder model is not available')
}
// Get embeddings for all texts in a single batch operation
@ -564,9 +469,9 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
if (embedding.length !== 512) {
this.logger(
'warn',
`Batch embedding dimension mismatch: expected 512, got ${embedding.length}. Standardizing...`
`Batch embedding dimension mismatch: expected 512, got ${embedding.length}. Standardizing to 512 dimensions.`
)
// If the embedding is too short, pad with zeros
if (embedding.length < 512) {
const paddedEmbedding = new Array(512).fill(0)
@ -574,21 +479,10 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
paddedEmbedding[i] = embedding[i]
}
return paddedEmbedding
}
}
// If the embedding is too long, truncate
else if (embedding.length > 512) {
// Special handling for 1536-dimensional vectors (common with newer models)
if (embedding.length === 1536) {
// Take every third value to reduce from 1536 to 512
const reducedEmbedding = new Array(512).fill(0)
for (let i = 0; i < 512; i++) {
reducedEmbedding[i] = embedding[i * 3]
}
return reducedEmbedding
} else {
// For other dimensions, just truncate
return embedding.slice(0, 512)
}
return embedding.slice(0, 512)
}
}
return embedding
@ -612,13 +506,14 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
return results
} catch (error) {
// No fallback - throw the error
this.logger(
'error',
'Failed to batch embed text with Universal Sentence Encoder:',
error
)
throw new Error(`Universal Sentence Encoder batch embedding failed: ${error}`)
throw new Error(
`Universal Sentence Encoder batch embedding failed: ${error}`
)
}
}
@ -693,7 +588,8 @@ function findUSELoadFunction(
} else if (
sentenceEncoderModule.default &&
sentenceEncoderModule.default.UniversalSentenceEncoder &&
typeof sentenceEncoderModule.default.UniversalSentenceEncoder.load === 'function'
typeof sentenceEncoderModule.default.UniversalSentenceEncoder.load ===
'function'
) {
loadFunction = sentenceEncoderModule.default.UniversalSentenceEncoder.load
}
@ -734,7 +630,7 @@ function findUSELoadFunction(
if (loadFunction) {
return async () => await loadFunction()
}
return null
}
@ -805,17 +701,19 @@ let sharedModel: UniversalSentenceEncoder | null = null
let sharedModelInitialized = false
let sharedModelVerbose = true
export function createTensorFlowEmbeddingFunction(options: { verbose?: boolean } = {}): EmbeddingFunction {
export function createTensorFlowEmbeddingFunction(
options: { verbose?: boolean } = {}
): EmbeddingFunction {
// Update verbose setting if provided
if (options.verbose !== undefined) {
sharedModelVerbose = options.verbose
}
// Create the shared model if it doesn't exist yet
if (!sharedModel) {
sharedModel = new UniversalSentenceEncoder({ verbose: sharedModelVerbose })
}
return async (data: any): Promise<Vector> => {
try {
// Initialize the model if it hasn't been initialized yet
@ -832,7 +730,12 @@ export function createTensorFlowEmbeddingFunction(options: { verbose?: boolean }
return await sharedModel!.embed(data)
} catch (error) {
logIfNotTest('error', 'Failed to use Universal Sentence Encoder:', [error], sharedModelVerbose)
logIfNotTest(
'error',
'Failed to use Universal Sentence Encoder:',
[error],
sharedModelVerbose
)
// No fallback - Universal Sentence Encoder is required
throw new Error(
`Universal Sentence Encoder is required and no fallbacks are allowed: ${error}`
@ -849,7 +752,9 @@ export function createTensorFlowEmbeddingFunction(options: { verbose?: boolean }
* @param options Configuration options
* @param options.verbose Whether to log non-essential messages (default: true)
*/
export function getDefaultEmbeddingFunction(options: { verbose?: boolean } = {}): EmbeddingFunction {
export function getDefaultEmbeddingFunction(
options: { verbose?: boolean } = {}
): EmbeddingFunction {
return createTensorFlowEmbeddingFunction(options)
}
@ -859,7 +764,8 @@ export function getDefaultEmbeddingFunction(options: { verbose?: boolean } = {})
* TensorFlow.js is required for this to work
* Uses CPU for compatibility
*/
export const defaultEmbeddingFunction: EmbeddingFunction = getDefaultEmbeddingFunction()
export const defaultEmbeddingFunction: EmbeddingFunction =
getDefaultEmbeddingFunction()
/**
* Creates a batch embedding function that uses UniversalSentenceEncoder
@ -874,19 +780,21 @@ let sharedBatchModel: UniversalSentenceEncoder | null = null
let sharedBatchModelInitialized = false
let sharedBatchModelVerbose = true
export function createBatchEmbeddingFunction(options: { verbose?: boolean } = {}): (
dataArray: string[]
) => Promise<Vector[]> {
export function createBatchEmbeddingFunction(
options: { verbose?: boolean } = {}
): (dataArray: string[]) => Promise<Vector[]> {
// Update verbose setting if provided
if (options.verbose !== undefined) {
sharedBatchModelVerbose = options.verbose
}
// Create the shared model if it doesn't exist yet
if (!sharedBatchModel) {
sharedBatchModel = new UniversalSentenceEncoder({ verbose: sharedBatchModelVerbose })
sharedBatchModel = new UniversalSentenceEncoder({
verbose: sharedBatchModelVerbose
})
}
return async (dataArray: string[]): Promise<Vector[]> => {
try {
// Initialize the model if it hasn't been initialized yet
@ -903,7 +811,12 @@ export function createBatchEmbeddingFunction(options: { verbose?: boolean } = {}
return await sharedBatchModel!.embedBatch(dataArray)
} catch (error) {
logIfNotTest('error', 'Failed to use Universal Sentence Encoder batch embedding:', [error], sharedBatchModelVerbose)
logIfNotTest(
'error',
'Failed to use Universal Sentence Encoder batch embedding:',
[error],
sharedBatchModelVerbose
)
// No fallback - Universal Sentence Encoder is required
throw new Error(
`Universal Sentence Encoder is required for batch embedding and no fallbacks are allowed: ${error}`
@ -920,9 +833,9 @@ export function createBatchEmbeddingFunction(options: { verbose?: boolean } = {}
* @param options Configuration options
* @param options.verbose Whether to log non-essential messages (default: true)
*/
export function getDefaultBatchEmbeddingFunction(options: { verbose?: boolean } = {}): (
dataArray: string[]
) => Promise<Vector[]> {
export function getDefaultBatchEmbeddingFunction(
options: { verbose?: boolean } = {}
): (dataArray: string[]) => Promise<Vector[]> {
return createBatchEmbeddingFunction(options)
}

326
src/utils/robustModelLoader.ts Normal file
View file

@ -0,0 +1,326 @@
/**
* Robust Model Loader - Enhanced model loading with retry mechanisms and fallbacks
*
* This module provides a more reliable way to load TensorFlow models with:
* - Exponential backoff retry mechanisms
* - Timeout handling
* - Multiple fallback strategies
* - Better error handling and logging
* - Optional local model bundling support
*/
import { EmbeddingModel } from '../coreTypes.js'
export interface ModelLoadOptions {
/** Maximum number of retry attempts */
maxRetries?: number
/** Initial retry delay in milliseconds */
initialRetryDelay?: number
/** Maximum retry delay in milliseconds */
maxRetryDelay?: number
/** Request timeout in milliseconds */
timeout?: number
/** Whether to use exponential backoff */
useExponentialBackoff?: boolean
/** Fallback model URLs to try if primary fails */
fallbackUrls?: string[]
/** Whether to enable verbose logging */
verbose?: boolean
/** Whether to prefer local bundled model if available */
preferLocalModel?: boolean
}
export interface RetryConfig {
attempt: number
maxRetries: number
delay: number
error: Error
}
export class RobustModelLoader {
private options: Required<ModelLoadOptions>
private loadAttempts: Map<string, number> = new Map()
constructor(options: ModelLoadOptions = {}) {
this.options = {
maxRetries: options.maxRetries ?? 3,
initialRetryDelay: options.initialRetryDelay ?? 1000,
maxRetryDelay: options.maxRetryDelay ?? 30000,
timeout: options.timeout ?? 60000, // 60 seconds
useExponentialBackoff: options.useExponentialBackoff ?? true,
fallbackUrls: options.fallbackUrls ?? [],
verbose: options.verbose ?? false,
preferLocalModel: options.preferLocalModel ?? true
}
}
/**
* Load a model with robust retry and fallback mechanisms
*/
async loadModel(
primaryLoadFunction: () => Promise<EmbeddingModel>,
modelIdentifier: string = 'default'
): Promise<EmbeddingModel> {
const startTime = Date.now()
this.log(`Starting robust model loading for: ${modelIdentifier}`)
// Try local bundled model first if preferred
if (this.options.preferLocalModel) {
try {
const localModel = await this.tryLoadLocalBundledModel()
if (localModel) {
this.log(`Successfully loaded local bundled model in ${Date.now() - startTime}ms`)
return localModel
}
} catch (error) {
this.log(`Local bundled model not available: ${error}`)
}
}
// Try primary load function with retries
try {
const model = await this.loadWithRetries(
primaryLoadFunction,
`primary-${modelIdentifier}`
)
this.log(`Successfully loaded model via primary method in ${Date.now() - startTime}ms`)
return model
} catch (primaryError) {
this.log(`Primary model loading failed: ${primaryError}`)
// Try fallback URLs if available
for (let i = 0; i < this.options.fallbackUrls.length; i++) {
const fallbackUrl = this.options.fallbackUrls[i]
this.log(`Trying fallback URL ${i + 1}/${this.options.fallbackUrls.length}: ${fallbackUrl}`)
try {
const fallbackModel = await this.loadWithRetries(
() => this.loadFromUrl(fallbackUrl),
`fallback-${i}-${modelIdentifier}`
)
this.log(`Successfully loaded model via fallback ${i + 1} in ${Date.now() - startTime}ms`)
return fallbackModel
} catch (fallbackError) {
this.log(`Fallback ${i + 1} failed: ${fallbackError}`)
}
}
// All attempts failed
const totalTime = Date.now() - startTime
const errorMessage = `All model loading attempts failed after ${totalTime}ms. Primary error: ${primaryError}`
this.log(errorMessage)
throw new Error(errorMessage)
}
}
/**
* Load a model with retry logic and exponential backoff
*/
private async loadWithRetries(
loadFunction: () => Promise<EmbeddingModel>,
identifier: string
): Promise<EmbeddingModel> {
let lastError: Error
const currentAttempts = this.loadAttempts.get(identifier) || 0
for (let attempt = currentAttempts; attempt <= this.options.maxRetries; attempt++) {
this.loadAttempts.set(identifier, attempt)
try {
this.log(`Attempt ${attempt + 1}/${this.options.maxRetries + 1} for ${identifier}`)
// Apply timeout to the load function
const model = await this.withTimeout(loadFunction(), this.options.timeout)
// Success - clear attempt counter
this.loadAttempts.delete(identifier)
return model
} catch (error) {
lastError = error as Error
this.log(`Attempt ${attempt + 1} failed: ${lastError.message}`)
// Don't retry on the last attempt
if (attempt === this.options.maxRetries) {
break
}
// Calculate delay for next attempt
const delay = this.calculateRetryDelay(attempt)
this.log(`Retrying in ${delay}ms...`)
// Wait before next attempt
await this.sleep(delay)
}
}
// All retries exhausted
this.loadAttempts.delete(identifier)
throw lastError!
}
/**
* Try to load a locally bundled model
*/
private async tryLoadLocalBundledModel(): Promise<EmbeddingModel | null> {
try {
// Check if we're in Node.js environment
const isNode = typeof process !== 'undefined' &&
process.versions != null &&
process.versions.node != null
if (isNode) {
// Try to load from bundled model directory
const path = await import('path')
const fs = await import('fs')
const { fileURLToPath } = await import('url')
const __filename = fileURLToPath(import.meta.url)
const __dirname = path.dirname(__filename)
// Look for bundled model in multiple possible locations
const possiblePaths = [
path.join(__dirname, '..', '..', 'models', 'bundled', 'universal-sentence-encoder'),
path.join(__dirname, '..', '..', 'node_modules', '@soulcraft', 'brainy-models', 'universal-sentence-encoder'),
path.join(process.cwd(), 'node_modules', '@soulcraft', 'brainy-models', 'universal-sentence-encoder')
]
for (const modelPath of possiblePaths) {
const modelJsonPath = path.join(modelPath, 'model.json')
if (fs.existsSync(modelJsonPath)) {
this.log(`Found bundled model at: ${modelJsonPath}`)
// Load TensorFlow.js if not already loaded
const tf = await import('@tensorflow/tfjs')
const model = await tf.loadLayersModel(`file://${modelJsonPath}`)
// Return a wrapper that matches the Universal Sentence Encoder interface
return this.createModelWrapper(model)
}
}
}
return null
} catch (error) {
this.log(`Error checking for bundled model: ${error}`)
return null
}
}
/**
* Load model from a specific URL
*/
private async loadFromUrl(url: string): Promise<EmbeddingModel> {
// This would need to be implemented based on the specific model type
// For now, we'll throw an error indicating this needs implementation
throw new Error(`Loading from custom URL not yet implemented: ${url}`)
}
/**
* Create a model wrapper that matches the Universal Sentence Encoder interface
*/
private createModelWrapper(tfModel: any): EmbeddingModel {
return {
init: async () => {
// Model is already loaded
},
embed: async (sentences: string | string[]) => {
const input = Array.isArray(sentences) ? sentences : [sentences]
// This is a simplified implementation - would need proper preprocessing
const inputTensors = tfModel.predict(input)
return inputTensors
},
dispose: async () => {
if (tfModel && tfModel.dispose) {
tfModel.dispose()
}
}
}
}
/**
* Apply timeout to a promise
*/
private async withTimeout<T>(promise: Promise<T>, timeoutMs: number): Promise<T> {
const timeoutPromise = new Promise<never>((_, reject) => {
setTimeout(() => {
reject(new Error(`Operation timed out after ${timeoutMs}ms`))
}, timeoutMs)
})
return Promise.race([promise, timeoutPromise])
}
/**
* Calculate retry delay with exponential backoff
*/
private calculateRetryDelay(attempt: number): number {
if (!this.options.useExponentialBackoff) {
return this.options.initialRetryDelay
}
// Exponential backoff: delay = initialDelay * (2 ^ attempt) + jitter
const exponentialDelay = this.options.initialRetryDelay * Math.pow(2, attempt)
// Add jitter (random factor) to prevent thundering herd
const jitter = Math.random() * 1000
// Cap at maximum delay
const delay = Math.min(exponentialDelay + jitter, this.options.maxRetryDelay)
return Math.floor(delay)
}
/**
* Sleep for specified milliseconds
*/
private sleep(ms: number): Promise<void> {
return new Promise(resolve => setTimeout(resolve, ms))
}
/**
* Log message if verbose mode is enabled
*/
private log(message: string): void {
if (this.options.verbose) {
console.log(`[RobustModelLoader] ${message}`)
}
}
/**
* Get loading statistics
*/
getLoadingStats(): { [key: string]: number } {
const stats: { [key: string]: number } = {}
for (const [identifier, attempts] of this.loadAttempts.entries()) {
stats[identifier] = attempts
}
return stats
}
/**
* Reset loading statistics
*/
resetStats(): void {
this.loadAttempts.clear()
}
}
/**
* Create a robust model loader with sensible defaults
*/
export function createRobustModelLoader(options?: ModelLoadOptions): RobustModelLoader {
return new RobustModelLoader(options)
}
/**
* Utility function to create fallback URLs for Universal Sentence Encoder
*/
export function getUniversalSentenceEncoderFallbacks(): string[] {
return [
'https://tfhub.dev/tensorflow/tfjs-model/universal-sentence-encoder/1/default/1',
'https://storage.googleapis.com/tfjs-models/savedmodel/universal_sentence_encoder/1/model.json',
// Add more fallback URLs as they become available
]
}