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

**feat(models): enhance loader reliability and compatibility**

Browse source 
- **Compatibility Enhancements**:
  - Added support to detect and inject missing `"format"` field in `model.json` files for TensorFlow.js compatibility.
  - Modified model loading logic to handle both `tfjs-graph-model` and `tfjs-layers-model` formats.

- **New Features**:
  - Introduced additional fallback paths for locating models to increase reliability in varying environments.
  - Added support for mock implementations of the Universal Sentence Encoder in test environments.

- **Bug Fixes**:
  - Fixed module loading resolution in `FileSystemStorage` with improved initialization and error handling for Node.js environments.
  - Resolved issues with test assertions to improve validation logic in core tests.

**Purpose**: Improve model loading reliability, expand compatibility with TensorFlow.js models, and enhance test environment support.
This commit is contained in:
David Snelling 2025-08-01 18:31:37 -07:00
parent ca737bc1aa
commit d71b939234
6 changed files with 566 additions and 384 deletions
Download patch file Download diff file Expand all files Collapse all files

33
src/utils/robustModelLoader.ts
View file

@ -11,6 +11,10 @@
import { EmbeddingModel } from '../coreTypes.js'
// Import the findUSELoadFunction from embedding.ts
// We need to access it directly since it's not exported
// For now, we'll implement a similar function locally
export interface ModelLoadOptions {
/** Maximum number of retry attempts */
maxRetries?: number
@ -181,6 +185,9 @@ export class RobustModelLoader {
// Look for bundled model in multiple possible locations
const possiblePaths = [
path.join(__dirname, '..', '..', 'models', 'bundled', 'universal-sentence-encoder'),
path.join(__dirname, '..', '..', 'brainy-models-package', 'models', 'universal-sentence-encoder'),
path.join(process.cwd(), 'brainy-models-package', 'models', 'universal-sentence-encoder'),
path.join(__dirname, '..', '..', 'node_modules', '@soulcraft', 'brainy-models', 'models', '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')
]
@ -192,7 +199,31 @@ export class RobustModelLoader {
// Load TensorFlow.js if not already loaded
const tf = await import('@tensorflow/tfjs')
const model = await tf.loadLayersModel(`file://${modelJsonPath}`)
// Read the model.json to check the format
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))
this.log(`✅ Added missing "format" field to model.json for TensorFlow.js compatibility`)
} catch (writeError) {
this.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)