feat: replace transformers.js with direct ONNX WASM for Bun compatibility
- Remove @huggingface/transformers dependency (539MB native binaries) - Add direct ONNX Runtime Web embedding engine - Bundle all-MiniLM-L6-v2-q8 model (24MB, no runtime downloads) - Works with Node.js, Bun, and bun build --compile - Air-gap compatible: fully self-contained, no internet required New WASM embedding components: - WASMEmbeddingEngine: Main integration class - WordPieceTokenizer: Pure TypeScript tokenizer - EmbeddingPostProcessor: Mean pooling + L2 normalization - ONNXInferenceEngine: Direct ONNX Runtime Web wrapper - AssetLoader: Model file loading Tests added: - 11 WASM embedding integration tests - 8 Bun compatibility tests New npm scripts: - test:wasm - Run WASM embedding tests - test:bun - Run tests with Bun - test:bun:compile - Build and run compiled binary
This commit is contained in:
parent
c1deb7a623
commit
1f59aa2013
21 changed files with 34431 additions and 3459 deletions
|
|
@ -1,24 +1,19 @@
|
|||
/**
|
||||
* Unified Embedding Manager
|
||||
*
|
||||
*
|
||||
* THE single source of truth for all embedding operations in Brainy.
|
||||
* Combines model management, precision configuration, and embedding generation
|
||||
* into one clean, maintainable class.
|
||||
*
|
||||
* Uses direct ONNX WASM inference for universal compatibility.
|
||||
*
|
||||
* Features:
|
||||
* - Singleton pattern ensures ONE model instance
|
||||
* - Automatic Q8 (default) or FP32 precision
|
||||
* - Model downloading and caching
|
||||
* - Thread-safe initialization
|
||||
* - Direct ONNX WASM (no transformers.js dependency)
|
||||
* - Bundled model (no runtime downloads)
|
||||
* - Works everywhere: Node.js, Bun, Bun --compile, browsers
|
||||
* - Memory monitoring
|
||||
*
|
||||
* This replaces: SingletonModelManager, TransformerEmbedding, ModelPrecisionManager,
|
||||
* hybridModelManager, universalMemoryManager, and more.
|
||||
*/
|
||||
|
||||
import { Vector, EmbeddingFunction } from '../coreTypes.js'
|
||||
import { pipeline, env } from '@huggingface/transformers'
|
||||
import { isNode } from '../utils/environment.js'
|
||||
import { WASMEmbeddingEngine } from './wasm/index.js'
|
||||
|
||||
// Types
|
||||
export type ModelPrecision = 'q8' | 'fp32'
|
||||
|
|
@ -38,22 +33,23 @@ let globalInitPromise: Promise<void> | null = null
|
|||
|
||||
/**
|
||||
* Unified Embedding Manager - Clean, simple, reliable
|
||||
*
|
||||
* Now powered by direct ONNX WASM for universal compatibility.
|
||||
*/
|
||||
export class EmbeddingManager {
|
||||
private model: any = null
|
||||
private precision: ModelPrecision
|
||||
private modelName = 'Xenova/all-MiniLM-L6-v2'
|
||||
private engine: WASMEmbeddingEngine
|
||||
private precision: ModelPrecision = 'q8'
|
||||
private modelName = 'all-MiniLM-L6-v2'
|
||||
private initialized = false
|
||||
private initTime: number | null = null
|
||||
private embedCount = 0
|
||||
private locked = false
|
||||
|
||||
|
||||
private constructor() {
|
||||
// Always use Q8 for optimal size/performance (99% accuracy, 75% smaller)
|
||||
this.precision = 'q8'
|
||||
console.log(`🎯 EmbeddingManager: Using Q8 precision`)
|
||||
this.engine = WASMEmbeddingEngine.getInstance()
|
||||
console.log('🎯 EmbeddingManager: Using Q8 precision (WASM)')
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* Get the singleton instance
|
||||
*/
|
||||
|
|
@ -63,16 +59,18 @@ export class EmbeddingManager {
|
|||
}
|
||||
return globalInstance
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* Initialize the model (happens once)
|
||||
*/
|
||||
async init(): Promise<void> {
|
||||
// In unit test mode, skip real model initialization
|
||||
const isTestMode = process.env.BRAINY_UNIT_TEST === 'true' || (globalThis as any).__BRAINY_UNIT_TEST__
|
||||
|
||||
const isTestMode =
|
||||
process.env.BRAINY_UNIT_TEST === 'true' ||
|
||||
(globalThis as any).__BRAINY_UNIT_TEST__
|
||||
|
||||
if (isTestMode) {
|
||||
// Production safeguard: Warn if mock mode is active but NODE_ENV is production
|
||||
// Production safeguard
|
||||
if (process.env.NODE_ENV === 'production') {
|
||||
throw new Error(
|
||||
'CRITICAL: Mock embeddings detected in production environment! ' +
|
||||
|
|
@ -80,7 +78,7 @@ export class EmbeddingManager {
|
|||
'This is a security risk. Remove test flags before deploying to production.'
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
if (!this.initialized) {
|
||||
this.initialized = true
|
||||
this.initTime = 1 // Mock init time
|
||||
|
|
@ -88,108 +86,65 @@ export class EmbeddingManager {
|
|||
}
|
||||
return
|
||||
}
|
||||
|
||||
|
||||
// Already initialized
|
||||
if (this.initialized && this.model) {
|
||||
if (this.initialized && this.engine.isInitialized()) {
|
||||
return
|
||||
}
|
||||
|
||||
|
||||
// Initialization in progress
|
||||
if (globalInitPromise) {
|
||||
await globalInitPromise
|
||||
return
|
||||
}
|
||||
|
||||
|
||||
// Start initialization
|
||||
globalInitPromise = this.performInit()
|
||||
|
||||
|
||||
try {
|
||||
await globalInitPromise
|
||||
} finally {
|
||||
globalInitPromise = null
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* Perform actual initialization
|
||||
*/
|
||||
private async performInit(): Promise<void> {
|
||||
const startTime = Date.now()
|
||||
console.log(`🚀 Initializing embedding model (${this.precision.toUpperCase()})...`)
|
||||
|
||||
|
||||
try {
|
||||
// Configure transformers.js environment
|
||||
const modelsPath = this.getModelsPath()
|
||||
env.cacheDir = modelsPath
|
||||
env.allowLocalModels = true
|
||||
env.useFSCache = true
|
||||
// Initialize WASM engine (handles all model loading)
|
||||
await this.engine.initialize()
|
||||
|
||||
// Check if models exist locally (only in Node.js)
|
||||
if (isNode()) {
|
||||
try {
|
||||
const nodeRequire = typeof require !== 'undefined' ? require : null
|
||||
if (nodeRequire) {
|
||||
const path = nodeRequire('node:path')
|
||||
const fs = nodeRequire('node:fs')
|
||||
const modelPath = path.join(modelsPath, ...this.modelName.split('/'))
|
||||
const hasLocalModels = fs.existsSync(modelPath)
|
||||
|
||||
if (hasLocalModels) {
|
||||
console.log('✅ Using cached models from:', modelPath)
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
// Silently continue if require fails
|
||||
}
|
||||
}
|
||||
|
||||
// Configure pipeline options for the selected precision
|
||||
const pipelineOptions: any = {
|
||||
cache_dir: modelsPath,
|
||||
local_files_only: false,
|
||||
// Always use Q8 precision
|
||||
dtype: 'q8',
|
||||
quantized: true,
|
||||
// Memory optimizations
|
||||
session_options: {
|
||||
enableCpuMemArena: false,
|
||||
enableMemPattern: false,
|
||||
interOpNumThreads: 1,
|
||||
intraOpNumThreads: 1,
|
||||
graphOptimizationLevel: 'disabled'
|
||||
}
|
||||
}
|
||||
|
||||
// Load the model
|
||||
this.model = await pipeline('feature-extraction', this.modelName, pipelineOptions)
|
||||
|
||||
// Lock precision after successful initialization
|
||||
this.locked = true
|
||||
this.initialized = true
|
||||
this.initTime = Date.now() - startTime
|
||||
|
||||
|
||||
// Log success
|
||||
const memoryMB = this.getMemoryUsage()
|
||||
console.log(`✅ Model loaded in ${this.initTime}ms`)
|
||||
console.log(`📊 Precision: Q8 | Memory: ${memoryMB}MB`)
|
||||
console.log(`🔒 Configuration locked`)
|
||||
|
||||
console.log('🔒 Configuration locked')
|
||||
} catch (error) {
|
||||
this.initialized = false
|
||||
this.model = null
|
||||
throw new Error(`Failed to initialize embedding model: ${error instanceof Error ? error.message : String(error)}`)
|
||||
throw new Error(
|
||||
`Failed to initialize embedding model: ${error instanceof Error ? error.message : String(error)}`
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* Generate embeddings
|
||||
*/
|
||||
async embed(text: string | string[] | Record<string, unknown>): Promise<Vector> {
|
||||
// Check for unit test environment - use mocks to prevent ONNX conflicts
|
||||
const isTestMode = process.env.BRAINY_UNIT_TEST === 'true' || (globalThis as any).__BRAINY_UNIT_TEST__
|
||||
// Check for unit test environment
|
||||
const isTestMode =
|
||||
process.env.BRAINY_UNIT_TEST === 'true' ||
|
||||
(globalThis as any).__BRAINY_UNIT_TEST__
|
||||
|
||||
if (isTestMode) {
|
||||
// Production safeguard
|
||||
if (process.env.NODE_ENV === 'production') {
|
||||
throw new Error('CRITICAL: Mock embeddings in production!')
|
||||
}
|
||||
|
|
@ -199,55 +154,40 @@ export class EmbeddingManager {
|
|||
// Ensure initialized
|
||||
await this.init()
|
||||
|
||||
if (!this.model) {
|
||||
throw new Error('Model not initialized')
|
||||
}
|
||||
|
||||
// CRITICAL FIX: Ensure input is always a string
|
||||
// Normalize input to string
|
||||
let input: string
|
||||
if (Array.isArray(text)) {
|
||||
// Join array elements, converting each to string first
|
||||
input = text.map(t => typeof t === 'string' ? t : String(t)).join(' ')
|
||||
input = text.map((t) => (typeof t === 'string' ? t : String(t))).join(' ')
|
||||
} else if (typeof text === 'string') {
|
||||
input = text
|
||||
} else if (typeof text === 'object') {
|
||||
// Convert object to string representation
|
||||
input = JSON.stringify(text)
|
||||
} else {
|
||||
// This shouldn't happen but let's be defensive
|
||||
console.warn('EmbeddingManager.embed received unexpected input type:', typeof text)
|
||||
input = String(text)
|
||||
}
|
||||
|
||||
// Generate embedding
|
||||
const output = await this.model(input, {
|
||||
pooling: 'mean',
|
||||
normalize: true
|
||||
})
|
||||
|
||||
// Extract embedding vector
|
||||
const embedding = Array.from(output.data) as number[]
|
||||
|
||||
|
||||
// Generate embedding using WASM engine
|
||||
const embedding = await this.engine.embed(input)
|
||||
|
||||
// Validate dimensions
|
||||
if (embedding.length !== 384) {
|
||||
console.warn(`Unexpected embedding dimension: ${embedding.length}`)
|
||||
// Pad or truncate
|
||||
if (embedding.length < 384) {
|
||||
return [...embedding, ...new Array(384 - embedding.length).fill(0)]
|
||||
} else {
|
||||
return embedding.slice(0, 384)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
this.embedCount++
|
||||
return embedding
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* Generate mock embeddings for unit tests
|
||||
*/
|
||||
private getMockEmbedding(text: string | string[] | Record<string, unknown>): Vector {
|
||||
// Use the same mock logic as setup-unit.ts for consistency
|
||||
const input = Array.isArray(text) ? text.join(' ') : text
|
||||
const str = typeof input === 'string' ? input : JSON.stringify(input)
|
||||
const vector = new Array(384).fill(0)
|
||||
|
|
@ -262,11 +202,10 @@ export class EmbeddingManager {
|
|||
vector[i] += Math.sin(i * 0.1 + str.length) * 0.1
|
||||
}
|
||||
|
||||
// Track mock embedding count
|
||||
this.embedCount++
|
||||
return vector
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* Get embedding function for compatibility
|
||||
*/
|
||||
|
|
@ -275,66 +214,7 @@ export class EmbeddingManager {
|
|||
return await this.embed(data)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* Get models directory path
|
||||
* Note: In browser environments, returns a simple default path
|
||||
* In Node.js, checks multiple locations for the models directory
|
||||
*/
|
||||
private getModelsPath(): string {
|
||||
// In browser environments, use a default path
|
||||
if (!isNode()) {
|
||||
return './models'
|
||||
}
|
||||
|
||||
// Node.js-specific model path resolution
|
||||
// Cache the result for performance
|
||||
if (!this.modelsPathCache) {
|
||||
this.modelsPathCache = this.resolveModelsPathSync()
|
||||
}
|
||||
return this.modelsPathCache
|
||||
}
|
||||
|
||||
private modelsPathCache: string | null = null
|
||||
|
||||
private resolveModelsPathSync(): string {
|
||||
// For Node.js environments, we can safely assume these modules exist
|
||||
// TypeScript will handle the imports at build time
|
||||
// At runtime, these will only be called if isNode() is true
|
||||
|
||||
// Default fallback path
|
||||
const defaultPath = './models'
|
||||
|
||||
try {
|
||||
// Create a conditional require function that only works in Node
|
||||
const nodeRequire = typeof require !== 'undefined' ? require : null
|
||||
if (!nodeRequire) return defaultPath
|
||||
|
||||
const fs = nodeRequire('node:fs')
|
||||
const path = nodeRequire('node:path')
|
||||
|
||||
const paths = [
|
||||
process.env.BRAINY_MODELS_PATH,
|
||||
'./models',
|
||||
path.join(process.cwd(), 'models'),
|
||||
path.join(process.env.HOME || '', '.brainy', 'models')
|
||||
]
|
||||
|
||||
for (const p of paths) {
|
||||
if (p && fs.existsSync(p)) {
|
||||
return p
|
||||
}
|
||||
}
|
||||
|
||||
// Default Node.js path
|
||||
return path.join(process.cwd(), 'models')
|
||||
} catch {
|
||||
// Fallback if require fails
|
||||
return defaultPath
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get memory usage in MB
|
||||
*/
|
||||
|
|
@ -345,35 +225,36 @@ export class EmbeddingManager {
|
|||
}
|
||||
return null
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* Get current statistics
|
||||
*/
|
||||
getStats(): EmbeddingStats {
|
||||
const engineStats = this.engine.getStats()
|
||||
return {
|
||||
initialized: this.initialized,
|
||||
precision: this.precision,
|
||||
modelName: this.modelName,
|
||||
embedCount: this.embedCount,
|
||||
embedCount: this.embedCount + engineStats.embedCount,
|
||||
initTime: this.initTime,
|
||||
memoryMB: this.getMemoryUsage()
|
||||
memoryMB: this.getMemoryUsage(),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* Check if initialized
|
||||
*/
|
||||
isInitialized(): boolean {
|
||||
return this.initialized
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* Get current precision
|
||||
*/
|
||||
getPrecision(): ModelPrecision {
|
||||
return this.precision
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* Validate precision matches expected
|
||||
*/
|
||||
|
|
@ -393,7 +274,9 @@ export const embeddingManager = EmbeddingManager.getInstance()
|
|||
/**
|
||||
* Direct embed function
|
||||
*/
|
||||
export async function embed(text: string | string[] | Record<string, unknown>): Promise<Vector> {
|
||||
export async function embed(
|
||||
text: string | string[] | Record<string, unknown>
|
||||
): Promise<Vector> {
|
||||
return await embeddingManager.embed(text)
|
||||
}
|
||||
|
||||
|
|
@ -409,4 +292,4 @@ export function getEmbeddingFunction(): EmbeddingFunction {
|
|||
*/
|
||||
export function getEmbeddingStats(): EmbeddingStats {
|
||||
return embeddingManager.getStats()
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue