fix: improve local model loading with USE-lite tokenizer support
- Update RobustModelLoader to properly handle @tensorflow-models/universal-sentence-encoder - Add support for loading USE-lite model with tokenizer from local files - Fix file:// URL handling issues in Node.js environment - Improve fallback mechanism for model loading - Add better error messages and logging for debugging
This commit is contained in:
parent
2a1db55a1d
commit
daf3b6243a
1 changed files with 50 additions and 48 deletions
|
|
@ -303,66 +303,68 @@ export class RobustModelLoader {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Second, try to use @soulcraft/brainy-models package if available
|
// Second, try to use @tensorflow-models/universal-sentence-encoder if available
|
||||||
|
// This package includes the tokenizer which is required for the model to work
|
||||||
try {
|
try {
|
||||||
console.log('Checking for @soulcraft/brainy-models package...')
|
console.log('Checking for @tensorflow-models/universal-sentence-encoder package...')
|
||||||
// Use dynamic import with string literal to avoid TypeScript compilation errors for optional dependency
|
const usePackageName = '@tensorflow-models/universal-sentence-encoder'
|
||||||
const packageName = '@soulcraft/brainy-models'
|
const use = await import(usePackageName).catch(() => null)
|
||||||
const brainyModels = await import(packageName).catch(() => null)
|
|
||||||
|
|
||||||
if (brainyModels?.BundledUniversalSentenceEncoder) {
|
if (use && use.load) {
|
||||||
console.log('✅ Found @soulcraft/brainy-models package installed')
|
console.log('✅ Found @tensorflow-models/universal-sentence-encoder package')
|
||||||
console.log(' Using local bundled model for maximum performance and reliability')
|
|
||||||
|
// Check if we have local model files from @soulcraft/brainy-models
|
||||||
|
let modelUrl: string | undefined = undefined
|
||||||
|
let vocabUrl: string | undefined = undefined
|
||||||
|
|
||||||
try {
|
try {
|
||||||
const encoder = new brainyModels.BundledUniversalSentenceEncoder({
|
// Try to find local model files
|
||||||
verbose: this.options.verbose,
|
const pathModule = 'path'
|
||||||
preferCompressed: false
|
const fsModule = 'fs'
|
||||||
})
|
const path = await import(/* @vite-ignore */ pathModule)
|
||||||
|
const fs = await import(/* @vite-ignore */ fsModule)
|
||||||
|
|
||||||
// Log metadata if available
|
// Check for brainy-models package
|
||||||
if (encoder.metadata) {
|
try {
|
||||||
console.log('📋 Model metadata:', encoder.metadata)
|
const brainyModelsPath = require.resolve('@soulcraft/brainy-models/package.json')
|
||||||
}
|
const brainyModelsDir = path.dirname(brainyModelsPath)
|
||||||
|
const modelDir = path.join(brainyModelsDir, 'models', 'universal-sentence-encoder')
|
||||||
await encoder.load()
|
const modelJsonPath = path.join(modelDir, 'model.json')
|
||||||
console.log('✅ Local Universal Sentence Encoder model loaded successfully')
|
|
||||||
|
if (fs.existsSync(modelJsonPath)) {
|
||||||
// Return a wrapper that matches the Universal Sentence Encoder interface
|
// Use file:// URL for local model
|
||||||
return {
|
modelUrl = `file://${modelJsonPath}`
|
||||||
init: async () => {
|
console.log(`📁 Using local model files from: ${modelDir}`)
|
||||||
// Already initialized
|
|
||||||
},
|
|
||||||
embed: async (sentences: string | string[]) => {
|
|
||||||
const input = Array.isArray(sentences) ? sentences : [sentences]
|
|
||||||
const embeddings = await encoder.embedToArrays(input)
|
|
||||||
|
|
||||||
// Return the first embedding as a Vector (number[])
|
// Check for vocab file
|
||||||
return embeddings[0] || []
|
const vocabPath = path.join(brainyModelsDir, 'models', 'vocab.json')
|
||||||
},
|
if (fs.existsSync(vocabPath)) {
|
||||||
dispose: async () => {
|
vocabUrl = `file://${vocabPath}`
|
||||||
encoder.dispose()
|
console.log(`📁 Using local vocab file from: ${vocabPath}`)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
} catch (e) {
|
||||||
|
// brainy-models not found, will use remote models
|
||||||
}
|
}
|
||||||
|
} catch (e) {
|
||||||
|
// Not in Node.js environment or files not found
|
||||||
|
}
|
||||||
|
|
||||||
|
try {
|
||||||
|
// Load the USE model with tokenizer
|
||||||
|
console.log('Loading Universal Sentence Encoder with tokenizer...')
|
||||||
|
const model = await use.load({ modelUrl, vocabUrl })
|
||||||
|
console.log('✅ Universal Sentence Encoder loaded successfully with tokenizer')
|
||||||
|
|
||||||
|
// The loaded model already has the correct interface
|
||||||
|
return model
|
||||||
} catch (loadError) {
|
} catch (loadError) {
|
||||||
console.error('Failed to load bundled model:', loadError)
|
console.error('Failed to load USE with tokenizer:', loadError)
|
||||||
// Try alternative loading method if available
|
// Fall through to try other methods
|
||||||
if (brainyModels.loadUniversalSentenceEncoder) {
|
|
||||||
console.log('Trying alternative loading method...')
|
|
||||||
try {
|
|
||||||
const model = await brainyModels.loadUniversalSentenceEncoder({
|
|
||||||
verbose: this.options.verbose
|
|
||||||
})
|
|
||||||
console.log('✅ Loaded via alternative method')
|
|
||||||
return model
|
|
||||||
} catch (altError) {
|
|
||||||
console.error('Alternative loading failed:', altError)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
} catch (importError) {
|
} catch (importError) {
|
||||||
this.log(`@soulcraft/brainy-models not available: ${importError}`)
|
this.log(`@tensorflow-models/universal-sentence-encoder not available: ${importError}`)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Check if we're in Node.js environment
|
// Check if we're in Node.js environment
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue