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

4
src/augmentations/conduitAugmentations.ts
View file

@ -745,7 +745,7 @@ export class WebRTCConduitAugmentation extends BaseConduitAugmentation implement
} else if (data instanceof ArrayBuffer) {
dc.send(new Uint8Array(data))
} else if (ArrayBuffer.isView(data)) {
dc.send(data)
dc.send(data as ArrayBufferView<ArrayBuffer>)
} else {
// Convert to JSON string
dc.send(JSON.stringify(data))
@ -923,7 +923,7 @@ export class WebRTCConduitAugmentation extends BaseConduitAugmentation implement
} else if (data instanceof ArrayBuffer) {
dataChannel.send(new Uint8Array(data))
} else if (ArrayBuffer.isView(data)) {
dataChannel.send(data)
dataChannel.send(data as ArrayBufferView<ArrayBuffer>)
} else {
// Convert to JSON string
dataChannel.send(JSON.stringify(data))

20
src/augmentations/memoryAugmentations.ts
View file

@ -4,7 +4,8 @@ import {
AugmentationResponse
} from '../types/augmentations.js'
import {StorageAdapter, Vector} from '../coreTypes.js'
import {MemoryStorage, FileSystemStorage, OPFSStorage} from '../storage/storageFactory.js'
import {MemoryStorage, OPFSStorage} from '../storage/storageFactory.js'
// FileSystemStorage will be dynamically imported when needed to avoid fs imports in browser
import {cosineDistance} from '../utils/distance.js'
/**
@ -253,9 +254,24 @@ export class MemoryStorageAugmentation extends BaseMemoryAugmentation {
export class FileSystemStorageAugmentation extends BaseMemoryAugmentation {
readonly description = 'Memory augmentation that stores data in the file system'
enabled = true
private rootDirectory: string
constructor(name: string, rootDirectory?: string) {
super(name, new FileSystemStorage(rootDirectory || '.'))
// Temporarily use MemoryStorage, will be replaced in initialize()
super(name, new MemoryStorage())
this.rootDirectory = rootDirectory || '.'
}
async initialize(): Promise<void> {
try {
// Dynamically import FileSystemStorage
const { FileSystemStorage } = await import('../storage/adapters/fileSystemStorage.js')
this.storage = new FileSystemStorage(this.rootDirectory)
await super.initialize()
} catch (error) {
console.error('Failed to load FileSystemStorage:', error)
throw new Error(`Failed to initialize FileSystemStorage: ${error}`)
}
}
getType(): AugmentationType {

37
src/browserFramework.ts Normal file
View file

@ -0,0 +1,37 @@
/**
* Browser Framework Entry Point for Brainy
* Optimized for modern frameworks like Angular, React, Vue, etc.
* Auto-detects environment and uses optimal storage (OPFS in browsers)
*/
import { BrainyData, BrainyDataConfig } from './brainyData.js'
import { VerbType, NounType } from './types/graphTypes.js'
/**
* Create a BrainyData instance optimized for browser frameworks
* Auto-detects environment and selects optimal storage and settings
*/
export async function createBrowserBrainyData(config: Partial<BrainyDataConfig> = {}): Promise<BrainyData> {
// BrainyData already has environment detection and will automatically:
// - Use OPFS storage in browsers with fallback to Memory
// - Use FileSystem storage in Node.js
// - Request persistent storage when appropriate
const browserConfig: BrainyDataConfig = {
storage: {
requestPersistentStorage: true // Request persistent storage for better performance
},
...config
}
const brainyData = new BrainyData(browserConfig)
await brainyData.init()
return brainyData
}
// Re-export types and constants for framework use
export { VerbType, NounType, BrainyData }
export type { BrainyDataConfig }
// Default export for easy importing
export default createBrowserBrainyData

252
src/demo.ts Normal file
View file

@ -0,0 +1,252 @@
/**
* Demo-specific entry point for browser environments
* This excludes all Node.js-specific functionality to avoid import issues
*/
// Import only browser-compatible modules
import { MemoryStorage } from './storage/adapters/memoryStorage.js'
import { OPFSStorage } from './storage/adapters/opfsStorage.js'
import { UniversalSentenceEncoder } from './utils/embedding.js'
import { cosineDistance, euclideanDistance } from './utils/distance.js'
import { isBrowser } from './utils/environment.js'
// Core types we need for the demo
export interface Vector extends Array<number> {}
export interface SearchResult {
id: string
score: number
metadata: any
text?: string
}
export interface VerbData {
id: string
source: string
target: string
verb: string
metadata: any
timestamp: number
}
/**
* Simplified BrainyData class for demo purposes
* Only includes browser-compatible functionality
*/
export class DemoBrainyData {
private storage: MemoryStorage | OPFSStorage
private embedder: UniversalSentenceEncoder | null = null
private initialized = false
private vectors = new Map<string, Vector>()
private metadata = new Map<string, any>()
private verbs = new Map<string, VerbData[]>()
constructor() {
// Always use memory storage for demo simplicity
this.storage = new MemoryStorage()
}
/**
* Initialize the database
*/
async init(): Promise<void> {
if (this.initialized) return
try {
await this.storage.init()
// Initialize the embedder
this.embedder = new UniversalSentenceEncoder({ verbose: false })
await this.embedder.init()
this.initialized = true
console.log('✅ Demo BrainyData initialized successfully')
} catch (error) {
console.error('Failed to initialize demo BrainyData:', error)
throw error
}
}
/**
* Add a document to the database
*/
async add(text: string, metadata: any = {}): Promise<string> {
if (!this.initialized || !this.embedder) {
throw new Error('Database not initialized')
}
const id = this.generateId()
try {
// Generate embedding
const vector = await this.embedder.embed(text)
// Store data
this.vectors.set(id, vector)
this.metadata.set(id, { text, ...metadata, timestamp: Date.now() })
return id
} catch (error) {
console.error('Failed to add document:', error)
throw error
}
}
/**
* Search for similar documents
*/
async searchText(query: string, limit: number = 10): Promise<SearchResult[]> {
if (!this.initialized || !this.embedder) {
throw new Error('Database not initialized')
}
try {
// Generate query embedding
const queryVector = await this.embedder.embed(query)
// Calculate similarities
const results: SearchResult[] = []
for (const [id, vector] of this.vectors.entries()) {
const score = 1 - cosineDistance(queryVector, vector) // Convert distance to similarity
const metadata = this.metadata.get(id)
results.push({
id,
score,
metadata,
text: metadata?.text
})
}
// Sort by score (highest first) and limit
return results
.sort((a, b) => b.score - a.score)
.slice(0, limit)
} catch (error) {
console.error('Search failed:', error)
throw error
}
}
/**
* Add a relationship between two documents
*/
async addVerb(sourceId: string, targetId: string, verb: string, metadata: any = {}): Promise<string> {
const verbId = this.generateId()
const verbData: VerbData = {
id: verbId,
source: sourceId,
target: targetId,
verb,
metadata,
timestamp: Date.now()
}
if (!this.verbs.has(sourceId)) {
this.verbs.set(sourceId, [])
}
this.verbs.get(sourceId)!.push(verbData)
return verbId
}
/**
* Get relationships from a source document
*/
async getVerbsBySource(sourceId: string): Promise<VerbData[]> {
return this.verbs.get(sourceId) || []
}
/**
* Get a document by ID
*/
async get(id: string): Promise<any | null> {
const metadata = this.metadata.get(id)
const vector = this.vectors.get(id)
if (!metadata || !vector) return null
return {
id,
vector,
...metadata
}
}
/**
* Delete a document
*/
async delete(id: string): Promise<boolean> {
const deleted = this.vectors.delete(id) && this.metadata.delete(id)
this.verbs.delete(id)
return deleted
}
/**
* Update document metadata
*/
async updateMetadata(id: string, newMetadata: any): Promise<boolean> {
const metadata = this.metadata.get(id)
if (!metadata) return false
this.metadata.set(id, { ...metadata, ...newMetadata })
return true
}
/**
* Get the number of documents
*/
size(): number {
return this.vectors.size
}
/**
* Generate a random ID
*/
private generateId(): string {
return 'id-' + Math.random().toString(36).substr(2, 9) + '-' + Date.now()
}
/**
* Get storage info
*/
getStorage(): MemoryStorage | OPFSStorage {
return this.storage
}
}
// Export noun and verb types for compatibility
export const NounType = {
Person: 'Person',
Organization: 'Organization',
Location: 'Location',
Thing: 'Thing',
Concept: 'Concept',
Event: 'Event',
Document: 'Document',
Media: 'Media',
File: 'File',
Message: 'Message',
Content: 'Content'
} as const
export const VerbType = {
RelatedTo: 'related_to',
Contains: 'contains',
PartOf: 'part_of',
LocatedAt: 'located_at',
References: 'references',
Owns: 'owns',
CreatedBy: 'created_by',
BelongsTo: 'belongs_to',
Likes: 'likes',
Follows: 'follows'
} as const
// Export the main class as BrainyData for compatibility
export { DemoBrainyData as BrainyData }
// Default export
export default DemoBrainyData

5
src/index.ts
View file

@ -99,7 +99,6 @@ export {
import {
OPFSStorage,
MemoryStorage,
FileSystemStorage,
R2Storage,
S3CompatibleStorage,
createStorage
@ -108,12 +107,14 @@ import {
export {
OPFSStorage,
MemoryStorage,
FileSystemStorage,
R2Storage,
S3CompatibleStorage,
createStorage
}
// FileSystemStorage is exported separately to avoid browser build issues
export { FileSystemStorage } from './storage/adapters/fileSystemStorage.js'
// Export unified pipeline
import {
Pipeline,

2340
src/storage/adapters/opfsStorage.ts
View file

File diff suppressed because it is too large Load diff

3
src/storage/storageFactory.ts
View file

@ -498,4 +498,5 @@ export {
}
// Export FileSystemStorage conditionally
export { FileSystemStorage } from './adapters/fileSystemStorage.js'
// NOTE: FileSystemStorage is now only imported dynamically to avoid fs imports in browser builds
// export { FileSystemStorage } from './adapters/fileSystemStorage.js'

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) {