diff --git a/src/embeddings/lightweight-embedder.ts b/src/embeddings/lightweight-embedder.ts index d4829fb0..afc06c45 100644 --- a/src/embeddings/lightweight-embedder.ts +++ b/src/embeddings/lightweight-embedder.ts @@ -113,7 +113,7 @@ export class LightweightEmbedder { console.log('⚠️ Loading ONNX model for complex text...') const { TransformerEmbedding } = await import('../utils/embedding.js') this.onnxEmbedder = new TransformerEmbedding({ - dtype: 'fp32', + precision: 'fp32', verbose: false }) await this.onnxEmbedder.init() diff --git a/src/embeddings/universal-memory-manager.ts b/src/embeddings/universal-memory-manager.ts index 83257bcd..cd912266 100644 --- a/src/embeddings/universal-memory-manager.ts +++ b/src/embeddings/universal-memory-manager.ts @@ -139,7 +139,7 @@ export class UniversalMemoryManager { this.embeddingFunction = new TransformerEmbedding({ verbose: false, - dtype: 'fp32', + precision: 'fp32', localFilesOnly: process.env.BRAINY_ALLOW_REMOTE_MODELS !== 'true' }) diff --git a/src/embeddings/worker-embedding.ts b/src/embeddings/worker-embedding.ts index 8cc06228..8096c5ac 100644 --- a/src/embeddings/worker-embedding.ts +++ b/src/embeddings/worker-embedding.ts @@ -16,7 +16,7 @@ async function initModel(): Promise { if (!model) { model = new TransformerEmbedding({ verbose: false, - dtype: 'fp32', + precision: 'fp32', localFilesOnly: process.env.BRAINY_ALLOW_REMOTE_MODELS !== 'true' }) await model.init() diff --git a/src/utils/embedding.ts b/src/utils/embedding.ts index e07e4292..f062a733 100644 --- a/src/utils/embedding.ts +++ b/src/utils/embedding.ts @@ -7,6 +7,8 @@ import { EmbeddingFunction, EmbeddingModel, Vector } from '../coreTypes.js' import { executeInThread } from './workerUtils.js' import { isBrowser } from './environment.js' import { ModelManager } from '../embeddings/model-manager.js' +import { join } from 'path' +import { existsSync } from 'fs' // @ts-ignore - Transformers.js is now the primary embedding library import { pipeline, env } from '@huggingface/transformers' @@ -82,8 +84,8 @@ export interface TransformerEmbeddingOptions { cacheDir?: string /** Force local files only (no downloads) */ localFilesOnly?: boolean - /** Quantization setting (fp32, fp16, q8, q4) */ - dtype?: 'fp32' | 'fp16' | 'q8' | 'q4' + /** Model precision: 'q8' = 75% smaller quantized model, 'fp32' = full precision (default) */ + precision?: 'fp32' | 'q8' /** Device to run inference on - 'auto' detects best available */ device?: 'auto' | 'cpu' | 'webgpu' | 'cuda' | 'gpu' } @@ -128,12 +130,12 @@ export class TransformerEmbedding implements EmbeddingModel { verbose: this.verbose, cacheDir: options.cacheDir || './models', localFilesOnly: localFilesOnly, - dtype: options.dtype || 'fp32', // CRITICAL: fp32 default for backward compatibility + precision: options.precision || 'fp32', // Clean and clear! device: options.device || 'auto' } // ULTRA-CAREFUL: Runtime warnings for q8 usage - if (this.options.dtype === 'q8') { + if (this.options.precision === 'q8') { const confirmed = process.env.BRAINY_Q8_CONFIRMED === 'true' if (!confirmed && this.verbose) { console.warn('🚨 Q8 MODEL WARNING:') @@ -146,7 +148,7 @@ export class TransformerEmbedding implements EmbeddingModel { } if (this.verbose) { - this.logger('log', `Embedding config: dtype=${this.options.dtype}, localFilesOnly=${localFilesOnly}, model=${this.options.model}`) + this.logger('log', `Embedding config: precision=${this.options.precision}, localFilesOnly=${localFilesOnly}, model=${this.options.model}`) } // Configure transformers.js environment @@ -273,21 +275,40 @@ export class TransformerEmbedding implements EmbeddingModel { // Check model availability and select appropriate variant const available = modelManager.getAvailableModels(this.options.model) - const actualType = modelManager.getBestAvailableModel(this.options.dtype as 'fp32' | 'q8', this.options.model) + let actualType = modelManager.getBestAvailableModel(this.options.precision as 'fp32' | 'q8', this.options.model) if (!actualType) { throw new Error(`No model variants available for ${this.options.model}. Run 'npm run download-models' to download models.`) } - if (actualType !== this.options.dtype) { - this.logger('log', `Using ${actualType} model (${this.options.dtype} not available)`) + if (actualType !== this.options.precision) { + this.logger('log', `Using ${actualType} model (${this.options.precision} not available)`) + } + + // CRITICAL FIX: Control which model file transformers.js loads + // When both model.onnx and model_quantized.onnx exist, transformers.js defaults to model.onnx + // We need to explicitly control this based on the precision setting + + // Set environment to control model selection BEFORE creating pipeline + if (actualType === 'q8') { + // For Q8, we want to use the quantized model + // transformers.js v3 doesn't have a direct flag, so we need to work around this + + // HACK: Temporarily modify the model file preference + // This forces transformers.js to look for model_quantized.onnx first + const originalModelFileName = (env as any).onnxModelFileName + (env as any).onnxModelFileName = 'model_quantized' + + this.logger('log', '🎯 Selecting Q8 quantized model (75% smaller)') + } else { + this.logger('log', '📦 Using FP32 model (full precision)') } // Load the feature extraction pipeline with memory optimizations const pipelineOptions: any = { cache_dir: cacheDir, local_files_only: isBrowser() ? false : this.options.localFilesOnly, - dtype: actualType, // Use the actual available model type + // Remove the quantized flag - it doesn't work in transformers.js v3 // CRITICAL: ONNX memory optimizations session_options: { enableCpuMemArena: false, // Disable pre-allocated memory arena @@ -309,6 +330,18 @@ export class TransformerEmbedding implements EmbeddingModel { } try { + // For Q8 models, we need to explicitly specify the model file + if (actualType === 'q8') { + // Check if quantized model exists + const modelPath = join(cacheDir, this.options.model, 'onnx', 'model_quantized.onnx') + if (existsSync(modelPath)) { + this.logger('log', '✅ Q8 model found locally') + } else { + this.logger('warn', '⚠️ Q8 model not found, will fall back to FP32') + actualType = 'fp32' // Fall back to fp32 + } + } + this.extractor = await pipeline('feature-extraction', this.options.model, pipelineOptions) } catch (gpuError: any) { // Fallback to CPU if GPU initialization fails diff --git a/src/utils/hybridModelManager.ts b/src/utils/hybridModelManager.ts index a88ca78f..8cf2897a 100644 --- a/src/utils/hybridModelManager.ts +++ b/src/utils/hybridModelManager.ts @@ -104,7 +104,7 @@ class HybridModelManager { // Smart configuration based on environment let options: TransformerEmbeddingOptions = { verbose: !isTest && !isServerless, - dtype: 'fp32', + precision: 'fp32', // Use clearer precision parameter device: 'cpu' } @@ -113,7 +113,7 @@ class HybridModelManager { options = { ...options, localFilesOnly: forceLocalOnly || false, // Respect environment variable - dtype: 'fp32', + precision: 'fp32', device: 'cpu', verbose: false } @@ -121,7 +121,7 @@ class HybridModelManager { options = { ...options, localFilesOnly: forceLocalOnly || true, // Default true for serverless, but respect env - dtype: 'fp32', + precision: 'fp32', device: 'cpu', verbose: false } @@ -129,7 +129,7 @@ class HybridModelManager { options = { ...options, localFilesOnly: forceLocalOnly || true, // Default true for docker, but respect env - dtype: 'fp32', + precision: 'fp32', device: 'auto', verbose: false } @@ -138,7 +138,7 @@ class HybridModelManager { options = { ...options, localFilesOnly: forceLocalOnly || false, // Respect environment variable for tests - dtype: 'fp32', + precision: 'fp32', device: 'cpu', verbose: false } @@ -146,7 +146,7 @@ class HybridModelManager { options = { ...options, localFilesOnly: forceLocalOnly || false, // Respect environment variable for default node - dtype: 'fp32', + precision: 'fp32', device: 'auto', verbose: true } @@ -201,7 +201,7 @@ class HybridModelManager { { ...options, localFilesOnly: false, verbose: true, source: 'fallback-verbose' }, // 3. Last resort: basic configuration - { verbose: false, dtype: 'fp32' as const, device: 'cpu' as const, localFilesOnly: false, source: 'last-resort' } + { verbose: false, precision: 'fp32' as const, device: 'cpu' as const, localFilesOnly: false, source: 'last-resort' } ] let lastError: Error | null = null