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

refactor: simplify build system and improve model loading flexibility

Browse source 
- Remove Rollup bundling in favor of direct TypeScript compilation
- Move from bundled models to dynamic model loading with configurable paths
- Add Docker deployment examples and documentation
- Implement robust model loader with fallback mechanisms
- Update storage adapters for better cross-environment compatibility
- Add comprehensive tests for model loading and package installation
- Simplify package.json scripts and remove complex build configurations
- Clean up deprecated demo files and old bundling scripts

BREAKING CHANGE: Models are no longer bundled with the package. They are now loaded dynamically from CDN or custom paths.
This commit is contained in:
David Snelling 2025-08-05 16:09:30 -07:00
parent 89413ebec2
commit 52a43d51d4
51 changed files with 4835 additions and 8007 deletions
Download patch file Download diff file Expand all files Collapse all files

2
src/utils/distance.ts
View file

@ -152,7 +152,7 @@ export async function calculateDistancesBatch(
// Ensure TextEncoder/TextDecoder are globally available in Node.js
const util = await import('util')
if (typeof global.TextEncoder === 'undefined') {
global.TextEncoder = util.TextEncoder
global.TextEncoder = util.TextEncoder as unknown as typeof TextEncoder
}
if (typeof global.TextDecoder === 'undefined') {
global.TextDecoder = util.TextDecoder as unknown as typeof TextDecoder

50
src/utils/embedding.ts
View file

@ -250,7 +250,7 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
) {
const util = await import('util')
if (!globalObj.TextEncoder) {
globalObj.TextEncoder = util.TextEncoder
globalObj.TextEncoder = util.TextEncoder as unknown as typeof TextEncoder
}
if (!globalObj.TextDecoder) {
globalObj.TextDecoder =
@ -309,8 +309,8 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
this.backend = 'cpu'
}
// Load Universal Sentence Encoder using dynamic import
this.use = await import('@tensorflow-models/universal-sentence-encoder')
// Note: @tensorflow-models/universal-sentence-encoder is no longer used
// Model loading is handled entirely by robustLoader
} catch (error) {
this.logger('error', 'Failed to initialize TensorFlow.js:', error)
// No fallback allowed - throw error
@ -324,30 +324,35 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
await this.tf.setBackend(this.backend)
}
// Try to find the load function in different possible module structures
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.'
)
}
// Load model using robustLoader which handles all loading strategies:
// 1. @soulcraft/brainy-models package if available (offline mode)
// 2. Direct TensorFlow.js URL loading as fallback
try {
// Load the model from local files first, falling back to default loading if necessary
this.model = await this.loadModelFromLocal(loadFunction)
this.model = await this.robustLoader.loadModelWithFallbacks()
this.initialized = true
// If the model doesn't have an embed method but has embedToArrays, wrap it
if (!this.model.embed && this.model.embedToArrays) {
const originalModel = this.model
this.model = {
embed: async (sentences: string | string[]) => {
const input = Array.isArray(sentences) ? sentences : [sentences]
const embeddings = await originalModel.embedToArrays(input)
// Return TensorFlow tensor-like object
return {
array: async () => embeddings,
arraySync: () => embeddings
}
},
dispose: () => originalModel.dispose ? originalModel.dispose() : undefined
}
}
} catch (modelError) {
this.logger(
'error',
'Failed to load Universal Sentence Encoder model:',
modelError
)
// No fallback allowed - throw error
throw new Error(
`Universal Sentence Encoder model loading failed: ${modelError}`
)
@ -572,10 +577,9 @@ 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
* Helper function - NO LONGER USED
* Kept for compatibility but will be removed in next major version
* @deprecated Since we removed @tensorflow-models/universal-sentence-encoder dependency
*/
function findUSELoadFunction(
sentenceEncoderModule: any

290
src/utils/robustModelLoader.ts
View file

@ -32,6 +32,8 @@ export interface ModelLoadOptions {
verbose?: boolean
/** Whether to prefer local bundled model if available */
preferLocalModel?: boolean
/** Custom directory path where models are stored (for Docker deployments) */
customModelsPath?: string
}
export interface RetryConfig {
@ -42,10 +44,16 @@ export interface RetryConfig {
}
export class RobustModelLoader {
private options: Required<ModelLoadOptions>
private options: Required<Omit<ModelLoadOptions, 'customModelsPath'>> & { customModelsPath?: string }
private loadAttempts: Map<string, number> = new Map()
constructor(options: ModelLoadOptions = {}) {
// Check for environment variables
const envModelsPath = process.env.BRAINY_MODELS_PATH || process.env.MODELS_PATH
// Auto-detect if we need to use an auto-extracted models directory
const autoDetectedPath = this.autoDetectModelsPath()
this.options = {
maxRetries: options.maxRetries ?? 3,
initialRetryDelay: options.initialRetryDelay ?? 1000,
@ -54,10 +62,126 @@ export class RobustModelLoader {
useExponentialBackoff: options.useExponentialBackoff ?? true,
fallbackUrls: options.fallbackUrls ?? [],
verbose: options.verbose ?? false,
preferLocalModel: options.preferLocalModel ?? true
preferLocalModel: options.preferLocalModel ?? true,
customModelsPath: options.customModelsPath ?? envModelsPath ?? autoDetectedPath
}
}
/**
* Auto-detect extracted models directory
*/
private autoDetectModelsPath(): string | undefined {
try {
// Check if we're in Node.js environment
const isNode = typeof process !== 'undefined' &&
process.versions != null &&
process.versions.node != null
if (!isNode) {
return undefined
}
// Try to detect extracted models directory
const possiblePaths = [
// Standard extraction location
'./models',
'../models',
'/app/models',
// Project root relative paths
process.cwd() + '/models',
// Docker/container standard paths
'/usr/src/app/models',
'/home/app/models'
]
// Use require to access fs and path synchronously (only in Node.js)
let fs: any, path: any
try {
fs = require('fs')
path = require('path')
} catch (error) {
// If require fails, we're probably in a browser environment
return undefined
}
for (const modelPath of possiblePaths) {
try {
// Check for marker file that indicates successful extraction
const markerFile = path.join(modelPath, '.brainy-models-extracted')
if (fs.existsSync(markerFile)) {
console.log(`🎯 Auto-detected extracted models at: ${modelPath}`)
return modelPath
}
// Fallback: check for universal-sentence-encoder directory
const useDir = path.join(modelPath, 'universal-sentence-encoder')
const modelJson = path.join(useDir, 'model.json')
if (fs.existsSync(modelJson)) {
console.log(`🎯 Auto-detected models directory at: ${modelPath}`)
return modelPath
}
} catch (error) {
continue
}
}
return undefined
} catch (error) {
return undefined
}
}
/**
* Load model with all available fallback strategies
*/
async loadModelWithFallbacks(): Promise<EmbeddingModel> {
const startTime = Date.now()
this.log('Starting model loading with all fallback strategies')
// Try local bundled model first (from @soulcraft/brainy-models if available)
const localModel = await this.tryLoadLocalBundledModel()
if (localModel) {
const loadTime = Date.now() - startTime
this.log(`✅ Model loaded successfully from local bundle in ${loadTime}ms`)
return localModel
}
// Fallback to loading from URLs
console.warn('⚠️ Local model not found. Falling back to remote model loading.')
console.warn(' For best performance and reliability:')
console.warn(' 1. Install @soulcraft/brainy-models: npm install @soulcraft/brainy-models')
console.warn(' 2. Or set BRAINY_MODELS_PATH environment variable for Docker deployments')
console.warn(' 3. Or use customModelsPath option in RobustModelLoader')
const fallbackUrls = getUniversalSentenceEncoderFallbacks()
for (const url of fallbackUrls) {
try {
this.log(`Attempting to load model from: ${url}`)
const model = await this.withTimeout(this.loadFromUrl(url), this.options.timeout)
const loadTime = Date.now() - startTime
// Verify it's the correct model by checking if it can embed text
try {
const testEmbedding = await model.embed('test')
if (!testEmbedding || (Array.isArray(testEmbedding) && testEmbedding.length !== 512)) {
throw new Error(`Model verification failed: incorrect embedding dimensions (expected 512, got ${Array.isArray(testEmbedding) ? testEmbedding.length : 'invalid'})`)
}
} catch (verifyError) {
console.warn(`⚠️ Model verification failed for ${url}: ${verifyError}`)
continue
}
console.warn(`✅ Successfully loaded Universal Sentence Encoder from remote URL: ${url}`)
console.warn(` Load time: ${loadTime}ms`)
return model
} catch (error) {
this.log(`Failed to load from ${url}: ${error}`)
}
}
throw new Error('Failed to load model from all available sources')
}
/**
* Load a model with robust retry and fallback mechanisms
*/
@ -168,15 +292,27 @@ export class RobustModelLoader {
*/
private async tryLoadLocalBundledModel(): Promise<EmbeddingModel | null> {
try {
// First, try to use @soulcraft/brainy-models package if available
// First, try custom models directory if specified (for Docker deployments)
if (this.options.customModelsPath) {
console.log(`Checking custom models directory: ${this.options.customModelsPath}`)
const customModel = await this.tryLoadFromCustomPath(this.options.customModelsPath)
if (customModel) {
console.log('✅ Successfully loaded model from custom directory')
console.log(' Using custom model path for Docker/production deployment')
return customModel
}
}
// Second, try to use @soulcraft/brainy-models package if available
try {
this.log('Checking for @soulcraft/brainy-models package...')
console.log('Checking for @soulcraft/brainy-models package...')
// Use dynamic import with string literal to avoid TypeScript compilation errors for optional dependency
const packageName = '@soulcraft/brainy-models'
const brainyModels = await import(packageName).catch(() => null)
if (brainyModels?.BundledUniversalSentenceEncoder) {
this.log('✅ Found @soulcraft/brainy-models package, using bundled model for maximum reliability')
console.log('✅ Found @soulcraft/brainy-models package installed')
console.log(' Using local bundled model for maximum performance and reliability')
const encoder = new brainyModels.BundledUniversalSentenceEncoder({
verbose: this.options.verbose,
@ -184,6 +320,7 @@ export class RobustModelLoader {
})
await encoder.load()
console.log('✅ Local Universal Sentence Encoder model loaded successfully')
// Return a wrapper that matches the Universal Sentence Encoder interface
return {
@ -212,10 +349,16 @@ export class RobustModelLoader {
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')
try {
// Try to load from bundled model directory
// Use dynamic import with a non-literal string to prevent Rollup from bundling these
const pathModule = 'path'
const fsModule = 'fs'
const urlModule = 'url'
const path = await import(/* @vite-ignore */ pathModule)
const fs = await import(/* @vite-ignore */ fsModule)
const { fileURLToPath } = await import(/* @vite-ignore */ urlModule)
const __filename = fileURLToPath(import.meta.url)
const __dirname = path.dirname(__filename)
@ -267,6 +410,9 @@ export class RobustModelLoader {
return this.createModelWrapper(model)
}
}
} catch (nodeImportError) {
this.log(`Could not load Node.js modules in browser: ${nodeImportError}`)
}
}
return null
@ -276,13 +422,109 @@ export class RobustModelLoader {
}
}
/**
* Try to load model from a custom directory path
*/
private async tryLoadFromCustomPath(customPath: string): 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) {
console.log('Custom model path only supported in Node.js environment')
return null
}
// Dynamic imports to avoid bundling issues
const pathModule = 'path'
const fsModule = 'fs'
const path = await import(/* @vite-ignore */ pathModule)
const fs = await import(/* @vite-ignore */ fsModule)
// Look for models in standard subdirectories
const possibleModelPaths = [
// Direct path to universal-sentence-encoder
path.join(customPath, 'universal-sentence-encoder'),
// Mirroring @soulcraft/brainy-models structure
path.join(customPath, 'models', 'universal-sentence-encoder'),
// TensorFlow hub model structure
path.join(customPath, 'tfhub', 'universal-sentence-encoder'),
// Simple models directory
path.join(customPath, 'use'),
// Check if customPath itself contains model.json
customPath
]
for (const modelPath of possibleModelPaths) {
const modelJsonPath = path.join(modelPath, 'model.json')
if (fs.existsSync(modelJsonPath)) {
console.log(`Found model at custom path: ${modelJsonPath}`)
// Load TensorFlow.js if not already loaded
const tf = await import('@tensorflow/tfjs')
// Read and validate the model.json
const modelJsonContent = JSON.parse(fs.readFileSync(modelJsonPath, 'utf8'))
// Ensure the format field exists for TensorFlow.js compatibility
if (!modelJsonContent.format) {
modelJsonContent.format = 'tfjs-graph-model'
try {
fs.writeFileSync(modelJsonPath, JSON.stringify(modelJsonContent, null, 2))
console.log(`✅ Added missing "format" field to model.json for TensorFlow.js compatibility`)
} catch (writeError) {
console.log(`⚠️ Could not write format field to model.json: ${writeError}`)
}
}
const modelFormat = modelJsonContent.format || 'tfjs-graph-model'
let model
if (modelFormat === 'tfjs-graph-model') {
// Use loadGraphModel for graph models
model = await tf.loadGraphModel(`file://${modelJsonPath}`)
} else {
// Use loadLayersModel for layers models (default)
model = await tf.loadLayersModel(`file://${modelJsonPath}`)
}
// Return a wrapper that matches the Universal Sentence Encoder interface
return this.createModelWrapper(model)
}
}
console.log(`No model found in custom path: ${customPath}`)
return null
} catch (error) {
console.log(`Error loading from custom path ${customPath}: ${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}`)
try {
this.log(`Loading model from URL: ${url}`)
// Import TensorFlow.js
const tf = await import('@tensorflow/tfjs')
// Load the model as a graph model
const model = await tf.loadGraphModel(url)
this.log(`✅ Successfully loaded model from: ${url}`)
// Return a wrapper that matches the Universal Sentence Encoder interface
return this.createModelWrapper(model)
} catch (error) {
throw new Error(`Failed to load model from ${url}: ${error}`)
}
}
/**
@ -294,11 +536,29 @@ export class RobustModelLoader {
// Model is already loaded
},
embed: async (sentences: string | string[]) => {
const tf = await import('@tensorflow/tfjs')
const input = Array.isArray(sentences) ? sentences : [sentences]
// This is a simplified implementation - would need proper preprocessing
const inputTensors = tfModel.predict(input)
return inputTensors
// Universal Sentence Encoder expects tokenized input
// For the tfhub model, we need to handle text preprocessing
// The model expects a tensor of strings
const inputTensor = tf.tensor(input)
try {
// Run the model prediction
const embeddings = await tfModel.predict(inputTensor)
// Convert to array and clean up
const result = await embeddings.array()
embeddings.dispose()
inputTensor.dispose()
// Return first embedding if single input, otherwise return all
return Array.isArray(sentences) ? result : (result[0] || [])
} catch (error) {
inputTensor.dispose()
throw new Error(`Failed to generate embeddings: ${error}`)
}
},
dispose: async () => {
if (tfModel && tfModel.dispose) {