feat: implement clean embedding architecture with Q8/FP32 precision control
- Unified embedding system with single EmbeddingManager - Q8 model support with 75% smaller footprint (23MB vs 90MB) - Intelligent precision auto-selection based on environment - Clean cached embeddings with TTL and memory management - Zero-config setup with smart defaults - Complete storage structure documentation - Removed legacy worker and hybrid managers - Streamlined model configuration and precision management
This commit is contained in:
parent
3227ad907c
commit
184d5dcf34
23 changed files with 1575 additions and 1369 deletions
354
src/embeddings/EmbeddingManager.ts
Normal file
354
src/embeddings/EmbeddingManager.ts
Normal file
|
|
@ -0,0 +1,354 @@
|
|||
/**
|
||||
* 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.
|
||||
*
|
||||
* Features:
|
||||
* - Singleton pattern ensures ONE model instance
|
||||
* - Automatic Q8 (default) or FP32 precision
|
||||
* - Model downloading and caching
|
||||
* - Thread-safe initialization
|
||||
* - Memory monitoring
|
||||
*
|
||||
* This replaces: SingletonModelManager, TransformerEmbedding, ModelPrecisionManager,
|
||||
* hybridModelManager, universalMemoryManager, and more.
|
||||
*/
|
||||
|
||||
import { Vector, EmbeddingFunction } from '../coreTypes.js'
|
||||
import { pipeline, env } from '@huggingface/transformers'
|
||||
import { existsSync } from 'fs'
|
||||
import { join } from 'path'
|
||||
|
||||
// Types
|
||||
export type ModelPrecision = 'q8' | 'fp32'
|
||||
|
||||
interface EmbeddingStats {
|
||||
initialized: boolean
|
||||
precision: ModelPrecision
|
||||
modelName: string
|
||||
embedCount: number
|
||||
initTime: number | null
|
||||
memoryMB: number | null
|
||||
}
|
||||
|
||||
// Global state for true singleton across entire process
|
||||
let globalInstance: EmbeddingManager | null = null
|
||||
let globalInitPromise: Promise<void> | null = null
|
||||
|
||||
/**
|
||||
* Unified Embedding Manager - Clean, simple, reliable
|
||||
*/
|
||||
export class EmbeddingManager {
|
||||
private model: any = null
|
||||
private precision: ModelPrecision
|
||||
private modelName = 'Xenova/all-MiniLM-L6-v2'
|
||||
private initialized = false
|
||||
private initTime: number | null = null
|
||||
private embedCount = 0
|
||||
private locked = false
|
||||
|
||||
private constructor() {
|
||||
// Determine precision - Q8 by default
|
||||
this.precision = this.determinePrecision()
|
||||
console.log(`🎯 EmbeddingManager: Using ${this.precision.toUpperCase()} precision`)
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the singleton instance
|
||||
*/
|
||||
static getInstance(): EmbeddingManager {
|
||||
if (!globalInstance) {
|
||||
globalInstance = new EmbeddingManager()
|
||||
}
|
||||
return globalInstance
|
||||
}
|
||||
|
||||
/**
|
||||
* Initialize the model (happens once)
|
||||
*/
|
||||
async init(): Promise<void> {
|
||||
// In unit test mode, skip real model initialization
|
||||
if (process.env.BRAINY_UNIT_TEST === 'true' || (globalThis as any).__BRAINY_UNIT_TEST__) {
|
||||
if (!this.initialized) {
|
||||
this.initialized = true
|
||||
this.initTime = 1 // Mock init time
|
||||
console.log('🧪 EmbeddingManager: Using mocked embeddings for unit tests')
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// Already initialized
|
||||
if (this.initialized && this.model) {
|
||||
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
|
||||
|
||||
// Check if models exist locally
|
||||
const modelPath = join(modelsPath, ...this.modelName.split('/'))
|
||||
const hasLocalModels = existsSync(modelPath)
|
||||
|
||||
if (hasLocalModels) {
|
||||
console.log('✅ Using cached models from:', modelPath)
|
||||
}
|
||||
|
||||
// Configure pipeline options for the selected precision
|
||||
const pipelineOptions: any = {
|
||||
cache_dir: modelsPath,
|
||||
local_files_only: false,
|
||||
// Specify precision
|
||||
dtype: this.precision,
|
||||
quantized: this.precision === 'q8',
|
||||
// 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: ${this.precision.toUpperCase()} | Memory: ${memoryMB}MB`)
|
||||
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)}`)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate embeddings
|
||||
*/
|
||||
async embed(text: string | string[]): Promise<Vector> {
|
||||
// Check for unit test environment - use mocks to prevent ONNX conflicts
|
||||
if (process.env.BRAINY_UNIT_TEST === 'true' || (globalThis as any).__BRAINY_UNIT_TEST__) {
|
||||
return this.getMockEmbedding(text)
|
||||
}
|
||||
|
||||
// Ensure initialized
|
||||
await this.init()
|
||||
|
||||
if (!this.model) {
|
||||
throw new Error('Model not initialized')
|
||||
}
|
||||
|
||||
// Handle array input
|
||||
const input = Array.isArray(text) ? text.join(' ') : text
|
||||
|
||||
// Generate embedding
|
||||
const output = await this.model(input, {
|
||||
pooling: 'mean',
|
||||
normalize: true
|
||||
})
|
||||
|
||||
// Extract embedding vector
|
||||
const embedding = Array.from(output.data) as number[]
|
||||
|
||||
// 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[]): 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)
|
||||
|
||||
// Create semi-realistic embeddings based on text content
|
||||
for (let i = 0; i < Math.min(str.length, 384); i++) {
|
||||
vector[i] = (str.charCodeAt(i % str.length) % 256) / 256
|
||||
}
|
||||
|
||||
// Add position-based variation
|
||||
for (let i = 0; i < 384; i++) {
|
||||
vector[i] += Math.sin(i * 0.1 + str.length) * 0.1
|
||||
}
|
||||
|
||||
// Track mock embedding count
|
||||
this.embedCount++
|
||||
return vector
|
||||
}
|
||||
|
||||
/**
|
||||
* Get embedding function for compatibility
|
||||
*/
|
||||
getEmbeddingFunction(): EmbeddingFunction {
|
||||
return async (data: string | string[]): Promise<Vector> => {
|
||||
return await this.embed(data)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Determine model precision
|
||||
*/
|
||||
private determinePrecision(): ModelPrecision {
|
||||
// Check environment variable overrides
|
||||
if (process.env.BRAINY_MODEL_PRECISION === 'fp32') {
|
||||
return 'fp32'
|
||||
}
|
||||
if (process.env.BRAINY_MODEL_PRECISION === 'q8') {
|
||||
return 'q8'
|
||||
}
|
||||
if (process.env.BRAINY_FORCE_FP32 === 'true') {
|
||||
return 'fp32'
|
||||
}
|
||||
|
||||
// Default to Q8 - optimal for most use cases
|
||||
return 'q8'
|
||||
}
|
||||
|
||||
/**
|
||||
* Get models directory path
|
||||
*/
|
||||
private getModelsPath(): string {
|
||||
// Check various possible locations
|
||||
const paths = [
|
||||
process.env.BRAINY_MODELS_PATH,
|
||||
'./models',
|
||||
join(process.cwd(), 'models'),
|
||||
join(process.env.HOME || '', '.brainy', 'models')
|
||||
]
|
||||
|
||||
for (const path of paths) {
|
||||
if (path && existsSync(path)) {
|
||||
return path
|
||||
}
|
||||
}
|
||||
|
||||
// Default
|
||||
return join(process.cwd(), 'models')
|
||||
}
|
||||
|
||||
/**
|
||||
* Get memory usage in MB
|
||||
*/
|
||||
private getMemoryUsage(): number | null {
|
||||
if (typeof process !== 'undefined' && process.memoryUsage) {
|
||||
const usage = process.memoryUsage()
|
||||
return Math.round(usage.heapUsed / 1024 / 1024)
|
||||
}
|
||||
return null
|
||||
}
|
||||
|
||||
/**
|
||||
* Get current statistics
|
||||
*/
|
||||
getStats(): EmbeddingStats {
|
||||
return {
|
||||
initialized: this.initialized,
|
||||
precision: this.precision,
|
||||
modelName: this.modelName,
|
||||
embedCount: this.embedCount,
|
||||
initTime: this.initTime,
|
||||
memoryMB: this.getMemoryUsage()
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if initialized
|
||||
*/
|
||||
isInitialized(): boolean {
|
||||
return this.initialized
|
||||
}
|
||||
|
||||
/**
|
||||
* Get current precision
|
||||
*/
|
||||
getPrecision(): ModelPrecision {
|
||||
return this.precision
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate precision matches expected
|
||||
*/
|
||||
validatePrecision(expected: ModelPrecision): void {
|
||||
if (this.locked && expected !== this.precision) {
|
||||
throw new Error(
|
||||
`Precision mismatch! System using ${this.precision.toUpperCase()} ` +
|
||||
`but ${expected.toUpperCase()} was requested. Cannot mix precisions.`
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Export singleton instance and convenience functions
|
||||
export const embeddingManager = EmbeddingManager.getInstance()
|
||||
|
||||
/**
|
||||
* Direct embed function
|
||||
*/
|
||||
export async function embed(text: string | string[]): Promise<Vector> {
|
||||
return await embeddingManager.embed(text)
|
||||
}
|
||||
|
||||
/**
|
||||
* Get embedding function for compatibility
|
||||
*/
|
||||
export function getEmbeddingFunction(): EmbeddingFunction {
|
||||
return embeddingManager.getEmbeddingFunction()
|
||||
}
|
||||
|
||||
/**
|
||||
* Get statistics
|
||||
*/
|
||||
export function getEmbeddingStats(): EmbeddingStats {
|
||||
return embeddingManager.getStats()
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue