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
3
.gitignore
vendored
3
.gitignore
vendored
|
|
@ -73,6 +73,9 @@ DISTRIBUTED_*.md
|
|||
models/
|
||||
models-cache/
|
||||
|
||||
# But include bundled WASM model assets
|
||||
!assets/models/
|
||||
|
||||
# Development planning files (not for commit)
|
||||
PLAN.md
|
||||
CLAUDE.md
|
||||
|
|
|
|||
25
assets/models/all-MiniLM-L6-v2-q8/config.json
Normal file
25
assets/models/all-MiniLM-L6-v2-q8/config.json
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
{
|
||||
"_name_or_path": "sentence-transformers/all-MiniLM-L6-v2",
|
||||
"architectures": [
|
||||
"BertModel"
|
||||
],
|
||||
"attention_probs_dropout_prob": 0.1,
|
||||
"classifier_dropout": null,
|
||||
"gradient_checkpointing": false,
|
||||
"hidden_act": "gelu",
|
||||
"hidden_dropout_prob": 0.1,
|
||||
"hidden_size": 384,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 1536,
|
||||
"layer_norm_eps": 1e-12,
|
||||
"max_position_embeddings": 512,
|
||||
"model_type": "bert",
|
||||
"num_attention_heads": 12,
|
||||
"num_hidden_layers": 6,
|
||||
"pad_token_id": 0,
|
||||
"position_embedding_type": "absolute",
|
||||
"transformers_version": "4.29.2",
|
||||
"type_vocab_size": 2,
|
||||
"use_cache": true,
|
||||
"vocab_size": 30522
|
||||
}
|
||||
BIN
assets/models/all-MiniLM-L6-v2-q8/model.onnx
Normal file
BIN
assets/models/all-MiniLM-L6-v2-q8/model.onnx
Normal file
Binary file not shown.
30686
assets/models/all-MiniLM-L6-v2-q8/tokenizer.json
Normal file
30686
assets/models/all-MiniLM-L6-v2-q8/tokenizer.json
Normal file
File diff suppressed because it is too large
Load diff
1
assets/models/all-MiniLM-L6-v2-q8/vocab.json
Normal file
1
assets/models/all-MiniLM-L6-v2-q8/vocab.json
Normal file
File diff suppressed because one or more lines are too long
4115
package-lock.json
generated
4115
package-lock.json
generated
File diff suppressed because it is too large
Load diff
11
package.json
11
package.json
|
|
@ -64,7 +64,8 @@
|
|||
"./dist/cortex/backupRestore.js": false
|
||||
},
|
||||
"engines": {
|
||||
"node": "22.x"
|
||||
"node": "22.x",
|
||||
"bun": ">=1.0.0"
|
||||
},
|
||||
"scripts": {
|
||||
"build": "npm run build:types:if-needed && npm run build:patterns:if-needed && npm run build:keywords:if-needed && tsc && tsc -p tsconfig.cli.json",
|
||||
|
|
@ -90,6 +91,10 @@
|
|||
"test:ci-unit": "CI=true vitest run --config tests/configs/vitest.unit.config.ts",
|
||||
"test:ci-integration": "NODE_OPTIONS='--max-old-space-size=16384' CI=true vitest run --config tests/configs/vitest.integration.config.ts",
|
||||
"test:ci": "npm run test:ci-unit",
|
||||
"test:bun": "bun tests/integration/bun-compile-test.ts",
|
||||
"test:bun:compile": "bun build tests/integration/bun-compile-test.ts --compile --outfile /tmp/brainy-bun-test && /tmp/brainy-bun-test",
|
||||
"test:wasm": "npx vitest run tests/integration/wasm-embeddings.test.ts",
|
||||
"download-model": "node scripts/download-model.cjs",
|
||||
"download-models": "node scripts/download-models.cjs",
|
||||
"download-models:q8": "node scripts/download-models.cjs",
|
||||
"models:verify": "node scripts/ensure-models.js",
|
||||
|
|
@ -142,7 +147,9 @@
|
|||
"dist/**/*.js",
|
||||
"dist/**/*.d.ts",
|
||||
"bin/",
|
||||
"assets/models/**/*",
|
||||
"scripts/download-models.cjs",
|
||||
"scripts/download-model.cjs",
|
||||
"scripts/ensure-models.js",
|
||||
"scripts/prepare-models.js",
|
||||
"brainy.png",
|
||||
|
|
@ -180,7 +187,7 @@
|
|||
"@azure/identity": "^4.0.0",
|
||||
"@azure/storage-blob": "^12.17.0",
|
||||
"@google-cloud/storage": "^7.14.0",
|
||||
"@huggingface/transformers": "^3.7.2",
|
||||
"onnxruntime-web": "^1.22.0",
|
||||
"@msgpack/msgpack": "^3.1.2",
|
||||
"@types/js-yaml": "^4.0.9",
|
||||
"boxen": "^8.0.1",
|
||||
|
|
|
|||
175
scripts/download-model.cjs
Normal file
175
scripts/download-model.cjs
Normal file
|
|
@ -0,0 +1,175 @@
|
|||
#!/usr/bin/env node
|
||||
/**
|
||||
* Download Model Assets
|
||||
*
|
||||
* Downloads the all-MiniLM-L6-v2 Q8 model from Hugging Face.
|
||||
* Run: node scripts/download-model.cjs
|
||||
*/
|
||||
|
||||
const fs = require('node:fs')
|
||||
const path = require('node:path')
|
||||
const https = require('node:https')
|
||||
|
||||
const MODEL_DIR = path.join(__dirname, '..', 'assets', 'models', 'all-MiniLM-L6-v2-q8')
|
||||
const BASE_URL = 'https://huggingface.co/Xenova/all-MiniLM-L6-v2/resolve/main/onnx'
|
||||
|
||||
const FILES = [
|
||||
{
|
||||
name: 'model_quantized.onnx',
|
||||
url: `${BASE_URL}/model_quantized.onnx`,
|
||||
dest: 'model.onnx',
|
||||
},
|
||||
{
|
||||
name: 'tokenizer.json',
|
||||
url: 'https://huggingface.co/Xenova/all-MiniLM-L6-v2/resolve/main/tokenizer.json',
|
||||
dest: 'tokenizer.json',
|
||||
},
|
||||
{
|
||||
name: 'config.json',
|
||||
url: 'https://huggingface.co/Xenova/all-MiniLM-L6-v2/resolve/main/config.json',
|
||||
dest: 'config.json',
|
||||
},
|
||||
{
|
||||
name: 'vocab.txt',
|
||||
url: 'https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/vocab.txt',
|
||||
dest: 'vocab.txt',
|
||||
},
|
||||
]
|
||||
|
||||
/**
|
||||
* Follow redirects and download file
|
||||
*/
|
||||
function downloadFile(url, destPath, maxRedirects = 5) {
|
||||
return new Promise((resolve, reject) => {
|
||||
if (maxRedirects === 0) {
|
||||
reject(new Error('Too many redirects'))
|
||||
return
|
||||
}
|
||||
|
||||
const doRequest = (reqUrl) => {
|
||||
const parsedUrl = new URL(reqUrl)
|
||||
const options = {
|
||||
hostname: parsedUrl.hostname,
|
||||
path: parsedUrl.pathname + parsedUrl.search,
|
||||
headers: {
|
||||
'User-Agent': 'Brainy-Model-Downloader/1.0',
|
||||
},
|
||||
}
|
||||
|
||||
https.get(options, (response) => {
|
||||
// Handle redirects
|
||||
if (response.statusCode >= 300 && response.statusCode < 400 && response.headers.location) {
|
||||
response.resume() // Consume response data to free memory
|
||||
const redirectUrl = response.headers.location.startsWith('http')
|
||||
? response.headers.location
|
||||
: new URL(response.headers.location, reqUrl).toString()
|
||||
console.log(` ↳ Redirecting to: ${redirectUrl.slice(0, 80)}...`)
|
||||
downloadFile(redirectUrl, destPath, maxRedirects - 1)
|
||||
.then(resolve)
|
||||
.catch(reject)
|
||||
return
|
||||
}
|
||||
|
||||
if (response.statusCode !== 200) {
|
||||
reject(new Error(`HTTP ${response.statusCode}`))
|
||||
return
|
||||
}
|
||||
|
||||
const fileStream = fs.createWriteStream(destPath)
|
||||
let downloadedBytes = 0
|
||||
const totalBytes = parseInt(response.headers['content-length'] || '0', 10)
|
||||
|
||||
response.on('data', (chunk) => {
|
||||
downloadedBytes += chunk.length
|
||||
if (totalBytes > 0) {
|
||||
const percent = Math.round((downloadedBytes / totalBytes) * 100)
|
||||
process.stdout.write(`\r Progress: ${percent}% (${Math.round(downloadedBytes / 1024 / 1024)}MB)`)
|
||||
}
|
||||
})
|
||||
|
||||
response.pipe(fileStream)
|
||||
|
||||
fileStream.on('finish', () => {
|
||||
fileStream.close()
|
||||
console.log(`\n ✅ Downloaded: ${path.basename(destPath)} (${Math.round(downloadedBytes / 1024 / 1024)}MB)`)
|
||||
resolve()
|
||||
})
|
||||
|
||||
fileStream.on('error', (err) => {
|
||||
fs.unlink(destPath, () => {}) // Delete partial file
|
||||
reject(err)
|
||||
})
|
||||
}).on('error', reject)
|
||||
}
|
||||
|
||||
doRequest(url)
|
||||
})
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert vocab.txt to vocab.json
|
||||
*/
|
||||
function convertVocabToJson(vocabTxtPath, vocabJsonPath) {
|
||||
console.log('📝 Converting vocab.txt to vocab.json...')
|
||||
const content = fs.readFileSync(vocabTxtPath, 'utf-8')
|
||||
const lines = content.split('\n').filter(line => line.trim())
|
||||
|
||||
const vocab = {}
|
||||
for (let i = 0; i < lines.length; i++) {
|
||||
vocab[lines[i]] = i
|
||||
}
|
||||
|
||||
fs.writeFileSync(vocabJsonPath, JSON.stringify(vocab))
|
||||
console.log(` ✅ Created vocab.json with ${Object.keys(vocab).length} tokens`)
|
||||
|
||||
// Remove vocab.txt since we have vocab.json
|
||||
fs.unlinkSync(vocabTxtPath)
|
||||
}
|
||||
|
||||
async function main() {
|
||||
console.log('🔽 Downloading all-MiniLM-L6-v2 Q8 model assets...\n')
|
||||
|
||||
// Create model directory
|
||||
fs.mkdirSync(MODEL_DIR, { recursive: true })
|
||||
console.log(`📁 Model directory: ${MODEL_DIR}\n`)
|
||||
|
||||
// Download each file
|
||||
for (const file of FILES) {
|
||||
const destPath = path.join(MODEL_DIR, file.dest)
|
||||
|
||||
// Check if already exists
|
||||
if (fs.existsSync(destPath)) {
|
||||
const stats = fs.statSync(destPath)
|
||||
if (stats.size > 0) {
|
||||
console.log(`⏭️ Skipping ${file.name} (already exists)`)
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
console.log(`📥 Downloading ${file.name}...`)
|
||||
try {
|
||||
await downloadFile(file.url, destPath)
|
||||
} catch (error) {
|
||||
console.error(` ❌ Failed to download ${file.name}: ${error.message}`)
|
||||
process.exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
// Convert vocab.txt to vocab.json
|
||||
const vocabTxtPath = path.join(MODEL_DIR, 'vocab.txt')
|
||||
const vocabJsonPath = path.join(MODEL_DIR, 'vocab.json')
|
||||
if (fs.existsSync(vocabTxtPath) && !fs.existsSync(vocabJsonPath)) {
|
||||
convertVocabToJson(vocabTxtPath, vocabJsonPath)
|
||||
}
|
||||
|
||||
console.log('\n✅ All model assets downloaded successfully!')
|
||||
console.log('\nModel files:')
|
||||
const files = fs.readdirSync(MODEL_DIR)
|
||||
for (const file of files) {
|
||||
const stats = fs.statSync(path.join(MODEL_DIR, file))
|
||||
const sizeMB = (stats.size / 1024 / 1024).toFixed(2)
|
||||
console.log(` - ${file}: ${sizeMB}MB`)
|
||||
}
|
||||
}
|
||||
|
||||
main().catch(console.error)
|
||||
|
|
@ -1,330 +1,131 @@
|
|||
/**
|
||||
* MODEL GUARDIAN - CRITICAL PATH
|
||||
*
|
||||
*
|
||||
* THIS IS THE MOST CRITICAL COMPONENT OF BRAINY
|
||||
* Without the exact model, users CANNOT access their data
|
||||
*
|
||||
*
|
||||
* Requirements:
|
||||
* 1. Model MUST be Xenova/all-MiniLM-L6-v2 (never changes)
|
||||
* 2. Model MUST be available at runtime
|
||||
* 1. Model MUST be all-MiniLM-L6-v2-q8 (bundled in package)
|
||||
* 2. Model MUST be available at runtime (embedded in npm package)
|
||||
* 3. Model MUST produce consistent 384-dim embeddings
|
||||
* 4. System MUST fail fast if model unavailable in production
|
||||
*/
|
||||
|
||||
import { env } from '@huggingface/transformers'
|
||||
import { createHash } from '../universal/crypto.js'
|
||||
import { WASMEmbeddingEngine } from '../embeddings/wasm/index.js'
|
||||
|
||||
// CRITICAL: These values MUST NEVER CHANGE
|
||||
const CRITICAL_MODEL_CONFIG = {
|
||||
modelName: 'Xenova/all-MiniLM-L6-v2',
|
||||
modelHash: {
|
||||
// SHA256 of model.onnx - computed from actual model
|
||||
'onnx/model.onnx': 'add_actual_hash_here',
|
||||
'tokenizer.json': 'add_actual_hash_here'
|
||||
} as Record<string, string>,
|
||||
modelSize: {
|
||||
'onnx/model.onnx': 90387606, // Exact size in bytes (updated to match actual file)
|
||||
'tokenizer.json': 711661
|
||||
} as Record<string, number>,
|
||||
modelName: 'all-MiniLM-L6-v2-q8',
|
||||
embeddingDimensions: 384,
|
||||
fallbackSources: [
|
||||
// Primary: Our Google Cloud Storage CDN (we control this, fastest)
|
||||
{
|
||||
name: 'Soulcraft CDN (Primary)',
|
||||
url: 'https://models.soulcraft.com/models/all-MiniLM-L6-v2.tar.gz',
|
||||
type: 'tarball'
|
||||
},
|
||||
// Secondary: GitHub releases backup
|
||||
{
|
||||
name: 'GitHub Backup',
|
||||
url: 'https://github.com/soulcraftlabs/brainy-models/releases/download/v1.0.0/all-MiniLM-L6-v2.tar.gz',
|
||||
type: 'tarball'
|
||||
},
|
||||
// Tertiary: Hugging Face (original source)
|
||||
{
|
||||
name: 'Hugging Face',
|
||||
url: 'huggingface',
|
||||
type: 'transformers'
|
||||
}
|
||||
]
|
||||
// Model is bundled in package - no external downloads needed
|
||||
bundled: true
|
||||
}
|
||||
|
||||
export class ModelGuardian {
|
||||
private static instance: ModelGuardian
|
||||
private isVerified = false
|
||||
private modelPath: string
|
||||
private lastVerification: Date | null = null
|
||||
|
||||
|
||||
private constructor() {
|
||||
this.modelPath = this.detectModelPath()
|
||||
// Model is bundled - no path detection needed
|
||||
}
|
||||
|
||||
|
||||
static getInstance(): ModelGuardian {
|
||||
if (!ModelGuardian.instance) {
|
||||
ModelGuardian.instance = new ModelGuardian()
|
||||
}
|
||||
return ModelGuardian.instance
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* CRITICAL: Verify model availability and integrity
|
||||
* This MUST be called before any embedding operations
|
||||
*/
|
||||
async ensureCriticalModel(): Promise<void> {
|
||||
console.log('DEBUG: ensureCriticalModel called')
|
||||
console.log('🛡️ MODEL GUARDIAN: Verifying critical model availability...')
|
||||
console.log(`🚀 Debug: Model path: ${this.modelPath}`)
|
||||
console.log(`🚀 Debug: Already verified: ${this.isVerified}`)
|
||||
|
||||
// Check if already verified in this session
|
||||
if (this.isVerified && this.lastVerification) {
|
||||
const hoursSinceVerification =
|
||||
const hoursSinceVerification =
|
||||
(Date.now() - this.lastVerification.getTime()) / (1000 * 60 * 60)
|
||||
|
||||
|
||||
if (hoursSinceVerification < 24) {
|
||||
console.log('✅ Model previously verified in this session')
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Step 1: Check if model exists locally
|
||||
console.log('🔍 Debug: Calling verifyLocalModel()')
|
||||
const modelExists = await this.verifyLocalModel()
|
||||
|
||||
if (modelExists) {
|
||||
console.log('✅ Critical model verified locally')
|
||||
|
||||
// Verify the bundled WASM model works
|
||||
const modelWorks = await this.verifyBundledModel()
|
||||
|
||||
if (modelWorks) {
|
||||
this.isVerified = true
|
||||
this.lastVerification = new Date()
|
||||
this.configureTransformers()
|
||||
return
|
||||
}
|
||||
|
||||
// Step 2: In production, FAIL FAST (Node.js only)
|
||||
if (typeof window === 'undefined' && process.env.NODE_ENV === 'production' && !process.env.BRAINY_ALLOW_RUNTIME_DOWNLOAD) {
|
||||
throw new Error(
|
||||
'🚨 CRITICAL FAILURE: Transformer model not found in production!\n' +
|
||||
'The model is REQUIRED for Brainy to function.\n' +
|
||||
'Users CANNOT access their data without it.\n' +
|
||||
'Solution: Run "npm run download-models" during build stage.'
|
||||
)
|
||||
}
|
||||
|
||||
// Step 3: Attempt to download from fallback sources
|
||||
console.warn('⚠️ Model not found locally, attempting download...')
|
||||
|
||||
for (const source of CRITICAL_MODEL_CONFIG.fallbackSources) {
|
||||
try {
|
||||
console.log(`📥 Trying ${source.name}...`)
|
||||
await this.downloadFromSource(source)
|
||||
|
||||
// Verify the download
|
||||
if (await this.verifyLocalModel()) {
|
||||
console.log(`✅ Successfully downloaded from ${source.name}`)
|
||||
this.isVerified = true
|
||||
this.lastVerification = new Date()
|
||||
this.configureTransformers()
|
||||
return
|
||||
}
|
||||
} catch (error) {
|
||||
console.warn(`❌ ${source.name} failed:`, (error as Error).message)
|
||||
}
|
||||
}
|
||||
|
||||
// Step 4: CRITICAL FAILURE
|
||||
|
||||
// CRITICAL FAILURE
|
||||
throw new Error(
|
||||
'🚨 CRITICAL FAILURE: Cannot obtain transformer model!\n' +
|
||||
'Tried all fallback sources.\n' +
|
||||
'Brainy CANNOT function without the model.\n' +
|
||||
'Users CANNOT access their data.\n' +
|
||||
'Please check network connectivity or pre-download models.'
|
||||
'🚨 CRITICAL FAILURE: Bundled transformer model not working!\n' +
|
||||
'The model is REQUIRED for Brainy to function.\n' +
|
||||
'Users CANNOT access their data without it.\n' +
|
||||
'This indicates a package installation issue.'
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* Verify the local model files exist and are correct
|
||||
* Verify the bundled WASM model works correctly
|
||||
*/
|
||||
private async verifyLocalModel(): Promise<boolean> {
|
||||
// Browser doesn't have local file access
|
||||
if (typeof window !== 'undefined') {
|
||||
console.log('⚠️ Model verification skipped in browser environment')
|
||||
return false
|
||||
}
|
||||
private async verifyBundledModel(): Promise<boolean> {
|
||||
try {
|
||||
const engine = WASMEmbeddingEngine.getInstance()
|
||||
|
||||
// Dynamically import Node.js modules
|
||||
const fs = await import('node:fs')
|
||||
const fsPromises = await import('node:fs/promises')
|
||||
const path = await import('node:path')
|
||||
// Initialize the engine (loads bundled model)
|
||||
await engine.initialize()
|
||||
|
||||
const modelBasePath = path.join(this.modelPath, ...CRITICAL_MODEL_CONFIG.modelName.split('/'))
|
||||
console.log(`🔍 Debug: Checking model at path: ${modelBasePath}`)
|
||||
console.log(`🔍 Debug: Model path components: ${this.modelPath} + ${CRITICAL_MODEL_CONFIG.modelName.split('/')}`)
|
||||
|
||||
// Check critical files
|
||||
const criticalFiles = [
|
||||
'onnx/model.onnx',
|
||||
'tokenizer.json',
|
||||
'config.json'
|
||||
]
|
||||
|
||||
for (const file of criticalFiles) {
|
||||
const filePath = path.join(modelBasePath, file)
|
||||
console.log(`🔍 Debug: Checking file: ${filePath}`)
|
||||
// Test embedding generation
|
||||
const testEmbedding = await engine.embed('test verification')
|
||||
|
||||
if (!fs.existsSync(filePath)) {
|
||||
console.log(`❌ Missing critical file: ${file} at ${filePath}`)
|
||||
// Verify dimensions
|
||||
if (testEmbedding.length !== CRITICAL_MODEL_CONFIG.embeddingDimensions) {
|
||||
console.error(
|
||||
`❌ CRITICAL: Model dimension mismatch!\n` +
|
||||
`Expected: ${CRITICAL_MODEL_CONFIG.embeddingDimensions}\n` +
|
||||
`Got: ${testEmbedding.length}`
|
||||
)
|
||||
return false
|
||||
}
|
||||
|
||||
// Verify size for critical files
|
||||
if (CRITICAL_MODEL_CONFIG.modelSize[file]) {
|
||||
const stats = await fsPromises.stat(filePath)
|
||||
const expectedSize = CRITICAL_MODEL_CONFIG.modelSize[file]
|
||||
|
||||
if (Math.abs(stats.size - expectedSize) > 1000) { // Allow 1KB variance
|
||||
console.error(
|
||||
`❌ CRITICAL: Model file size mismatch!\n` +
|
||||
`File: ${file}\n` +
|
||||
`Expected: ${expectedSize} bytes\n` +
|
||||
`Actual: ${stats.size} bytes\n` +
|
||||
`This indicates model corruption or version mismatch!`
|
||||
)
|
||||
return false
|
||||
}
|
||||
|
||||
// Verify normalization (should be unit length)
|
||||
const norm = Math.sqrt(testEmbedding.reduce((sum, v) => sum + v * v, 0))
|
||||
if (Math.abs(norm - 1.0) > 0.01) {
|
||||
console.error(`❌ CRITICAL: Embeddings not normalized! Norm: ${norm}`)
|
||||
return false
|
||||
}
|
||||
|
||||
// SHA256 verification for ultimate security
|
||||
if (CRITICAL_MODEL_CONFIG.modelHash && CRITICAL_MODEL_CONFIG.modelHash[file]) {
|
||||
const hash = await this.computeFileHash(filePath)
|
||||
if (hash !== CRITICAL_MODEL_CONFIG.modelHash[file]) {
|
||||
console.error(
|
||||
`❌ CRITICAL: Model hash mismatch for ${file}!\n` +
|
||||
`Expected: ${CRITICAL_MODEL_CONFIG.modelHash[file]}\n` +
|
||||
`Got: ${hash}\n` +
|
||||
`This indicates model tampering or corruption!`
|
||||
)
|
||||
return false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
/**
|
||||
* Compute SHA256 hash of a file
|
||||
*/
|
||||
private async computeFileHash(filePath: string): Promise<string> {
|
||||
try {
|
||||
const { readFile } = await import('node:fs/promises')
|
||||
const { createHash } = await import('node:crypto')
|
||||
const fileBuffer = await readFile(filePath)
|
||||
const hash = createHash('sha256').update(fileBuffer).digest('hex')
|
||||
return hash
|
||||
|
||||
return true
|
||||
} catch (error) {
|
||||
console.error(`Failed to compute hash for ${filePath}:`, error)
|
||||
return ''
|
||||
console.error('❌ Model verification failed:', error)
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Download model from a fallback source
|
||||
*/
|
||||
private async downloadFromSource(source: any): Promise<void> {
|
||||
if (source.type === 'transformers') {
|
||||
// Use transformers.js native download
|
||||
const { pipeline } = await import('@huggingface/transformers')
|
||||
env.cacheDir = this.modelPath
|
||||
env.allowRemoteModels = true
|
||||
|
||||
const extractor = await pipeline(
|
||||
'feature-extraction',
|
||||
CRITICAL_MODEL_CONFIG.modelName
|
||||
)
|
||||
|
||||
// Test the model
|
||||
const test = await extractor('test', { pooling: 'mean', normalize: true })
|
||||
if (test.data.length !== CRITICAL_MODEL_CONFIG.embeddingDimensions) {
|
||||
throw new Error(
|
||||
`CRITICAL: Model dimension mismatch! ` +
|
||||
`Expected ${CRITICAL_MODEL_CONFIG.embeddingDimensions}, ` +
|
||||
`got ${test.data.length}`
|
||||
)
|
||||
}
|
||||
} else if (source.type === 'tarball') {
|
||||
// Tarball extraction would require additional dependencies
|
||||
// Skip this source and try next fallback
|
||||
console.warn(`⚠️ Tarball extraction not available for ${source.name}. Trying next source...`)
|
||||
return // Will continue to next source in the loop
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Configure transformers.js to use verified local model
|
||||
*/
|
||||
private configureTransformers(): void {
|
||||
env.localModelPath = this.modelPath
|
||||
env.allowRemoteModels = false // Force local only after verification
|
||||
console.log('🔒 Transformers configured to use verified local model')
|
||||
}
|
||||
|
||||
/**
|
||||
* Detect where models should be stored
|
||||
*/
|
||||
private detectModelPath(): string {
|
||||
// Browser always uses default path
|
||||
if (typeof window !== 'undefined') {
|
||||
return './models'
|
||||
}
|
||||
|
||||
// Use require for synchronous access in Node.js
|
||||
try {
|
||||
const fs = require('node:fs')
|
||||
const path = require('node:path')
|
||||
|
||||
const candidates = [
|
||||
process.env.BRAINY_MODELS_PATH,
|
||||
'./models',
|
||||
path.join(process.cwd(), 'models'),
|
||||
path.join(process.env.HOME || '', '.brainy', 'models'),
|
||||
'/opt/models', // Lambda/container path
|
||||
env.cacheDir
|
||||
]
|
||||
|
||||
for (const candidatePath of candidates) {
|
||||
if (candidatePath && fs.existsSync(candidatePath)) {
|
||||
const modelPath = path.join(candidatePath, ...CRITICAL_MODEL_CONFIG.modelName.split('/'))
|
||||
if (fs.existsSync(path.join(modelPath, 'onnx', 'model.onnx'))) {
|
||||
return candidatePath // Return the models directory, not its parent
|
||||
}
|
||||
}
|
||||
}
|
||||
} catch (e) {
|
||||
// If Node.js modules not available, return default
|
||||
}
|
||||
|
||||
// Default
|
||||
return './models'
|
||||
}
|
||||
|
||||
/**
|
||||
* Get model status for diagnostics
|
||||
*/
|
||||
async getStatus(): Promise<{
|
||||
verified: boolean
|
||||
path: string
|
||||
lastVerification: Date | null
|
||||
modelName: string
|
||||
dimensions: number
|
||||
bundled: boolean
|
||||
}> {
|
||||
return {
|
||||
verified: this.isVerified,
|
||||
path: this.modelPath,
|
||||
lastVerification: this.lastVerification,
|
||||
modelName: CRITICAL_MODEL_CONFIG.modelName,
|
||||
dimensions: CRITICAL_MODEL_CONFIG.embeddingDimensions
|
||||
dimensions: CRITICAL_MODEL_CONFIG.embeddingDimensions,
|
||||
bundled: CRITICAL_MODEL_CONFIG.bundled
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* Force re-verification (for testing)
|
||||
*/
|
||||
|
|
@ -336,4 +137,4 @@ export class ModelGuardian {
|
|||
}
|
||||
|
||||
// Export singleton instance
|
||||
export const modelGuardian = ModelGuardian.getInstance()
|
||||
export const modelGuardian = ModelGuardian.getInstance()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
}
|
||||
|
|
|
|||
268
src/embeddings/wasm/AssetLoader.ts
Normal file
268
src/embeddings/wasm/AssetLoader.ts
Normal file
|
|
@ -0,0 +1,268 @@
|
|||
/**
|
||||
* Asset Loader
|
||||
*
|
||||
* Resolves paths to model files (ONNX model, vocabulary) across environments.
|
||||
* Handles Node.js, Bun, and bundled scenarios.
|
||||
*
|
||||
* Asset Resolution Order:
|
||||
* 1. Environment variable: BRAINY_MODEL_PATH
|
||||
* 2. Package-relative: node_modules/@soulcraft/brainy/assets/models/
|
||||
* 3. Project-relative: ./assets/models/
|
||||
*/
|
||||
|
||||
import { MODEL_CONSTANTS } from './types.js'
|
||||
|
||||
// Cache resolved paths
|
||||
let cachedModelDir: string | null = null
|
||||
let cachedVocab: Record<string, number> | null = null
|
||||
|
||||
/**
|
||||
* Asset loader for model files
|
||||
*/
|
||||
export class AssetLoader {
|
||||
private modelDir: string | null = null
|
||||
|
||||
/**
|
||||
* Get the model directory path
|
||||
*/
|
||||
async getModelDir(): Promise<string> {
|
||||
if (this.modelDir) {
|
||||
return this.modelDir
|
||||
}
|
||||
|
||||
if (cachedModelDir) {
|
||||
this.modelDir = cachedModelDir
|
||||
return cachedModelDir
|
||||
}
|
||||
|
||||
// Try to resolve model directory
|
||||
const resolved = await this.resolveModelDir()
|
||||
this.modelDir = resolved
|
||||
cachedModelDir = resolved
|
||||
return resolved
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve the model directory across environments
|
||||
*/
|
||||
private async resolveModelDir(): Promise<string> {
|
||||
// 1. Check environment variable
|
||||
if (typeof process !== 'undefined' && process.env?.BRAINY_MODEL_PATH) {
|
||||
const envPath = process.env.BRAINY_MODEL_PATH
|
||||
if (await this.pathExists(envPath)) {
|
||||
return envPath
|
||||
}
|
||||
}
|
||||
|
||||
// 2. Try common locations
|
||||
const modelName = MODEL_CONSTANTS.MODEL_NAME + '-q8'
|
||||
const possiblePaths = [
|
||||
// Package assets (when installed as dependency)
|
||||
`./assets/models/${modelName}`,
|
||||
`./node_modules/@soulcraft/brainy/assets/models/${modelName}`,
|
||||
// Development paths
|
||||
`../assets/models/${modelName}`,
|
||||
// Absolute from package root
|
||||
this.getPackageRootPath(`assets/models/${modelName}`),
|
||||
].filter(Boolean) as string[]
|
||||
|
||||
for (const path of possiblePaths) {
|
||||
if (await this.pathExists(path)) {
|
||||
return path
|
||||
}
|
||||
}
|
||||
|
||||
// If no path found, return default (will error on use)
|
||||
return `./assets/models/${modelName}`
|
||||
}
|
||||
|
||||
/**
|
||||
* Get package root path (Node.js/Bun only)
|
||||
*/
|
||||
private getPackageRootPath(relativePath: string): string | null {
|
||||
if (typeof process === 'undefined') {
|
||||
return null
|
||||
}
|
||||
|
||||
try {
|
||||
// Use __dirname equivalent
|
||||
const url = new URL(import.meta.url)
|
||||
const currentDir = url.pathname.replace(/\/[^/]*$/, '')
|
||||
// Go up from src/embeddings/wasm to package root
|
||||
const packageRoot = currentDir.replace(/\/src\/embeddings\/wasm$/, '')
|
||||
return `${packageRoot}/${relativePath}`
|
||||
} catch {
|
||||
return null
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if path exists (works in Node.js/Bun)
|
||||
*/
|
||||
private async pathExists(path: string): Promise<boolean> {
|
||||
if (typeof process === 'undefined') {
|
||||
// Browser - check via fetch
|
||||
try {
|
||||
const response = await fetch(path, { method: 'HEAD' })
|
||||
return response.ok
|
||||
} catch {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// Node.js/Bun
|
||||
try {
|
||||
const fs = await import('node:fs/promises')
|
||||
await fs.access(path)
|
||||
return true
|
||||
} catch {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get path to ONNX model file
|
||||
*/
|
||||
async getModelPath(): Promise<string> {
|
||||
const dir = await this.getModelDir()
|
||||
return `${dir}/model.onnx`
|
||||
}
|
||||
|
||||
/**
|
||||
* Get path to vocabulary file
|
||||
*/
|
||||
async getVocabPath(): Promise<string> {
|
||||
const dir = await this.getModelDir()
|
||||
return `${dir}/vocab.json`
|
||||
}
|
||||
|
||||
/**
|
||||
* Load vocabulary from JSON file
|
||||
*/
|
||||
async loadVocab(): Promise<Record<string, number>> {
|
||||
if (cachedVocab) {
|
||||
return cachedVocab
|
||||
}
|
||||
|
||||
const vocabPath = await this.getVocabPath()
|
||||
|
||||
if (typeof process !== 'undefined') {
|
||||
// Node.js/Bun - read from filesystem
|
||||
try {
|
||||
const fs = await import('node:fs/promises')
|
||||
const content = await fs.readFile(vocabPath, 'utf-8')
|
||||
cachedVocab = JSON.parse(content)
|
||||
return cachedVocab!
|
||||
} catch (error) {
|
||||
throw new Error(
|
||||
`Failed to load vocabulary from ${vocabPath}: ${error instanceof Error ? error.message : String(error)}`
|
||||
)
|
||||
}
|
||||
} else {
|
||||
// Browser - fetch
|
||||
try {
|
||||
const response = await fetch(vocabPath)
|
||||
if (!response.ok) {
|
||||
throw new Error(`HTTP ${response.status}`)
|
||||
}
|
||||
cachedVocab = await response.json()
|
||||
return cachedVocab!
|
||||
} catch (error) {
|
||||
throw new Error(
|
||||
`Failed to fetch vocabulary from ${vocabPath}: ${error instanceof Error ? error.message : String(error)}`
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Load model as ArrayBuffer (for ONNX session)
|
||||
*/
|
||||
async loadModel(): Promise<ArrayBuffer> {
|
||||
const modelPath = await this.getModelPath()
|
||||
|
||||
if (typeof process !== 'undefined') {
|
||||
// Node.js/Bun - read from filesystem
|
||||
try {
|
||||
const fs = await import('node:fs/promises')
|
||||
const buffer = await fs.readFile(modelPath)
|
||||
// Convert Node.js Buffer to ArrayBuffer
|
||||
return new Uint8Array(buffer).buffer as ArrayBuffer
|
||||
} catch (error) {
|
||||
throw new Error(
|
||||
`Failed to load model from ${modelPath}: ${error instanceof Error ? error.message : String(error)}`
|
||||
)
|
||||
}
|
||||
} else {
|
||||
// Browser - fetch
|
||||
try {
|
||||
const response = await fetch(modelPath)
|
||||
if (!response.ok) {
|
||||
throw new Error(`HTTP ${response.status}`)
|
||||
}
|
||||
return await response.arrayBuffer()
|
||||
} catch (error) {
|
||||
throw new Error(
|
||||
`Failed to fetch model from ${modelPath}: ${error instanceof Error ? error.message : String(error)}`
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Verify all required assets exist
|
||||
*/
|
||||
async verifyAssets(): Promise<{
|
||||
valid: boolean
|
||||
modelPath: string
|
||||
vocabPath: string
|
||||
errors: string[]
|
||||
}> {
|
||||
const errors: string[] = []
|
||||
const modelPath = await this.getModelPath()
|
||||
const vocabPath = await this.getVocabPath()
|
||||
|
||||
if (!(await this.pathExists(modelPath))) {
|
||||
errors.push(`Model file not found: ${modelPath}`)
|
||||
}
|
||||
|
||||
if (!(await this.pathExists(vocabPath))) {
|
||||
errors.push(`Vocabulary file not found: ${vocabPath}`)
|
||||
}
|
||||
|
||||
return {
|
||||
valid: errors.length === 0,
|
||||
modelPath,
|
||||
vocabPath,
|
||||
errors,
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Clear cached paths (for testing)
|
||||
*/
|
||||
clearCache(): void {
|
||||
this.modelDir = null
|
||||
cachedModelDir = null
|
||||
cachedVocab = null
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create asset loader instance
|
||||
*/
|
||||
export function createAssetLoader(): AssetLoader {
|
||||
return new AssetLoader()
|
||||
}
|
||||
|
||||
/**
|
||||
* Singleton asset loader
|
||||
*/
|
||||
let singletonLoader: AssetLoader | null = null
|
||||
|
||||
export function getAssetLoader(): AssetLoader {
|
||||
if (!singletonLoader) {
|
||||
singletonLoader = new AssetLoader()
|
||||
}
|
||||
return singletonLoader
|
||||
}
|
||||
156
src/embeddings/wasm/EmbeddingPostProcessor.ts
Normal file
156
src/embeddings/wasm/EmbeddingPostProcessor.ts
Normal file
|
|
@ -0,0 +1,156 @@
|
|||
/**
|
||||
* Embedding Post-Processor
|
||||
*
|
||||
* Converts raw ONNX model output to final embedding vectors.
|
||||
* Implements mean pooling and L2 normalization as used by sentence-transformers.
|
||||
*
|
||||
* Pipeline:
|
||||
* 1. Mean Pooling: Average token embeddings (weighted by attention mask)
|
||||
* 2. L2 Normalization: Normalize to unit length for cosine similarity
|
||||
*/
|
||||
|
||||
import { MODEL_CONSTANTS } from './types.js'
|
||||
|
||||
/**
|
||||
* Post-processor for converting ONNX output to sentence embeddings
|
||||
*/
|
||||
export class EmbeddingPostProcessor {
|
||||
private hiddenSize: number
|
||||
|
||||
constructor(hiddenSize: number = MODEL_CONSTANTS.HIDDEN_SIZE) {
|
||||
this.hiddenSize = hiddenSize
|
||||
}
|
||||
|
||||
/**
|
||||
* Mean pool token embeddings weighted by attention mask
|
||||
*
|
||||
* @param hiddenStates - Raw model output [seqLen * hiddenSize] flattened
|
||||
* @param attentionMask - Attention mask [seqLen] (1 for real tokens, 0 for padding)
|
||||
* @param seqLen - Sequence length
|
||||
* @returns Mean-pooled embedding [hiddenSize]
|
||||
*/
|
||||
meanPool(
|
||||
hiddenStates: Float32Array,
|
||||
attentionMask: number[],
|
||||
seqLen: number
|
||||
): Float32Array {
|
||||
const result = new Float32Array(this.hiddenSize)
|
||||
|
||||
// Sum of attention mask (number of real tokens)
|
||||
let maskSum = 0
|
||||
for (let i = 0; i < seqLen; i++) {
|
||||
maskSum += attentionMask[i]
|
||||
}
|
||||
|
||||
// Avoid division by zero
|
||||
if (maskSum === 0) {
|
||||
maskSum = 1
|
||||
}
|
||||
|
||||
// Compute weighted sum for each dimension
|
||||
for (let dim = 0; dim < this.hiddenSize; dim++) {
|
||||
let sum = 0
|
||||
for (let pos = 0; pos < seqLen; pos++) {
|
||||
// Get hidden state at [pos, dim]
|
||||
const value = hiddenStates[pos * this.hiddenSize + dim]
|
||||
// Weight by attention mask
|
||||
sum += value * attentionMask[pos]
|
||||
}
|
||||
// Mean pool
|
||||
result[dim] = sum / maskSum
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
/**
|
||||
* L2 normalize embedding to unit length
|
||||
*
|
||||
* @param embedding - Input embedding
|
||||
* @returns Normalized embedding with ||x|| = 1
|
||||
*/
|
||||
normalize(embedding: Float32Array): Float32Array {
|
||||
// Compute L2 norm
|
||||
let sumSquares = 0
|
||||
for (let i = 0; i < embedding.length; i++) {
|
||||
sumSquares += embedding[i] * embedding[i]
|
||||
}
|
||||
|
||||
const norm = Math.sqrt(sumSquares)
|
||||
|
||||
// Avoid division by zero
|
||||
if (norm === 0) {
|
||||
return embedding
|
||||
}
|
||||
|
||||
// Normalize
|
||||
const result = new Float32Array(embedding.length)
|
||||
for (let i = 0; i < embedding.length; i++) {
|
||||
result[i] = embedding[i] / norm
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
/**
|
||||
* Full post-processing pipeline: mean pool then normalize
|
||||
*
|
||||
* @param hiddenStates - Raw model output [seqLen * hiddenSize]
|
||||
* @param attentionMask - Attention mask [seqLen]
|
||||
* @param seqLen - Sequence length
|
||||
* @returns Final normalized embedding [hiddenSize]
|
||||
*/
|
||||
process(
|
||||
hiddenStates: Float32Array,
|
||||
attentionMask: number[],
|
||||
seqLen: number
|
||||
): Float32Array {
|
||||
const pooled = this.meanPool(hiddenStates, attentionMask, seqLen)
|
||||
return this.normalize(pooled)
|
||||
}
|
||||
|
||||
/**
|
||||
* Process batch of embeddings
|
||||
*
|
||||
* @param hiddenStates - Raw model output [batchSize * seqLen * hiddenSize]
|
||||
* @param attentionMasks - Attention masks [batchSize][seqLen]
|
||||
* @param batchSize - Number of sequences in batch
|
||||
* @param seqLen - Sequence length (same for all in batch due to padding)
|
||||
* @returns Array of normalized embeddings
|
||||
*/
|
||||
processBatch(
|
||||
hiddenStates: Float32Array,
|
||||
attentionMasks: number[][],
|
||||
batchSize: number,
|
||||
seqLen: number
|
||||
): Float32Array[] {
|
||||
const results: Float32Array[] = []
|
||||
const sequenceSize = seqLen * this.hiddenSize
|
||||
|
||||
for (let b = 0; b < batchSize; b++) {
|
||||
// Extract this sequence's hidden states
|
||||
const start = b * sequenceSize
|
||||
const seqHiddenStates = hiddenStates.slice(start, start + sequenceSize)
|
||||
|
||||
// Process
|
||||
const embedding = this.process(seqHiddenStates, attentionMasks[b], seqLen)
|
||||
results.push(embedding)
|
||||
}
|
||||
|
||||
return results
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert Float32Array to number array
|
||||
*/
|
||||
toNumberArray(embedding: Float32Array): number[] {
|
||||
return Array.from(embedding)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a post-processor with default configuration
|
||||
*/
|
||||
export function createPostProcessor(): EmbeddingPostProcessor {
|
||||
return new EmbeddingPostProcessor(MODEL_CONSTANTS.HIDDEN_SIZE)
|
||||
}
|
||||
193
src/embeddings/wasm/ONNXInferenceEngine.ts
Normal file
193
src/embeddings/wasm/ONNXInferenceEngine.ts
Normal file
|
|
@ -0,0 +1,193 @@
|
|||
/**
|
||||
* ONNX Inference Engine
|
||||
*
|
||||
* Direct ONNX Runtime Web wrapper for running model inference.
|
||||
* Uses WASM backend for universal compatibility (Node.js, Bun, Browser).
|
||||
*
|
||||
* This replaces transformers.js dependency with direct ONNX control.
|
||||
*/
|
||||
|
||||
import * as ort from 'onnxruntime-web'
|
||||
import { InferenceConfig, MODEL_CONSTANTS } from './types.js'
|
||||
|
||||
// Configure ONNX Runtime for WASM-only
|
||||
ort.env.wasm.numThreads = 1 // Single-threaded for stability
|
||||
ort.env.wasm.simd = true // Enable SIMD where available
|
||||
|
||||
/**
|
||||
* ONNX Inference Engine using onnxruntime-web
|
||||
*/
|
||||
export class ONNXInferenceEngine {
|
||||
private session: ort.InferenceSession | null = null
|
||||
private initialized = false
|
||||
private modelPath: string
|
||||
private config: InferenceConfig
|
||||
|
||||
constructor(config: Partial<InferenceConfig> = {}) {
|
||||
this.modelPath = config.modelPath ?? ''
|
||||
this.config = {
|
||||
modelPath: this.modelPath,
|
||||
numThreads: config.numThreads ?? 1,
|
||||
enableSimd: config.enableSimd ?? true,
|
||||
enableCpuMemArena: config.enableCpuMemArena ?? false,
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Initialize the ONNX session
|
||||
*/
|
||||
async initialize(modelPath?: string): Promise<void> {
|
||||
if (this.initialized && this.session) {
|
||||
return
|
||||
}
|
||||
|
||||
const path = modelPath ?? this.modelPath
|
||||
if (!path) {
|
||||
throw new Error('Model path is required')
|
||||
}
|
||||
|
||||
try {
|
||||
// Configure session options
|
||||
const sessionOptions: ort.InferenceSession.SessionOptions = {
|
||||
executionProviders: ['wasm'],
|
||||
graphOptimizationLevel: 'all',
|
||||
enableCpuMemArena: this.config.enableCpuMemArena,
|
||||
// Additional WASM-specific options
|
||||
executionMode: 'sequential',
|
||||
}
|
||||
|
||||
// Load model from file path or URL
|
||||
this.session = await ort.InferenceSession.create(path, sessionOptions)
|
||||
|
||||
this.initialized = true
|
||||
} catch (error) {
|
||||
this.initialized = false
|
||||
this.session = null
|
||||
throw new Error(
|
||||
`Failed to initialize ONNX session: ${error instanceof Error ? error.message : String(error)}`
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Run inference on tokenized input
|
||||
*
|
||||
* @param inputIds - Token IDs [batchSize, seqLen]
|
||||
* @param attentionMask - Attention mask [batchSize, seqLen]
|
||||
* @param tokenTypeIds - Token type IDs [batchSize, seqLen] (optional, defaults to zeros)
|
||||
* @returns Hidden states [batchSize, seqLen, hiddenSize]
|
||||
*/
|
||||
async infer(
|
||||
inputIds: number[][],
|
||||
attentionMask: number[][],
|
||||
tokenTypeIds?: number[][]
|
||||
): Promise<Float32Array> {
|
||||
if (!this.session) {
|
||||
throw new Error('Session not initialized. Call initialize() first.')
|
||||
}
|
||||
|
||||
const batchSize = inputIds.length
|
||||
const seqLen = inputIds[0].length
|
||||
|
||||
// Convert to BigInt64Array (ONNX int64 type)
|
||||
const inputIdsFlat = new BigInt64Array(batchSize * seqLen)
|
||||
const attentionMaskFlat = new BigInt64Array(batchSize * seqLen)
|
||||
const tokenTypeIdsFlat = new BigInt64Array(batchSize * seqLen)
|
||||
|
||||
for (let b = 0; b < batchSize; b++) {
|
||||
for (let s = 0; s < seqLen; s++) {
|
||||
const idx = b * seqLen + s
|
||||
inputIdsFlat[idx] = BigInt(inputIds[b][s])
|
||||
attentionMaskFlat[idx] = BigInt(attentionMask[b][s])
|
||||
tokenTypeIdsFlat[idx] = tokenTypeIds
|
||||
? BigInt(tokenTypeIds[b][s])
|
||||
: BigInt(0)
|
||||
}
|
||||
}
|
||||
|
||||
// Create ONNX tensors
|
||||
const inputIdsTensor = new ort.Tensor('int64', inputIdsFlat, [batchSize, seqLen])
|
||||
const attentionMaskTensor = new ort.Tensor('int64', attentionMaskFlat, [batchSize, seqLen])
|
||||
const tokenTypeIdsTensor = new ort.Tensor('int64', tokenTypeIdsFlat, [batchSize, seqLen])
|
||||
|
||||
try {
|
||||
// Run inference
|
||||
const feeds = {
|
||||
input_ids: inputIdsTensor,
|
||||
attention_mask: attentionMaskTensor,
|
||||
token_type_ids: tokenTypeIdsTensor,
|
||||
}
|
||||
|
||||
const results = await this.session.run(feeds)
|
||||
|
||||
// Extract last_hidden_state (the output we need for mean pooling)
|
||||
// Model outputs: last_hidden_state [batch, seq, hidden] and pooler_output [batch, hidden]
|
||||
const output = results.last_hidden_state ?? results.token_embeddings
|
||||
|
||||
if (!output) {
|
||||
throw new Error('Model did not return expected output tensor')
|
||||
}
|
||||
|
||||
return output.data as Float32Array
|
||||
} finally {
|
||||
// Dispose tensors to free memory
|
||||
inputIdsTensor.dispose()
|
||||
attentionMaskTensor.dispose()
|
||||
tokenTypeIdsTensor.dispose()
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Infer single sequence (convenience method)
|
||||
*/
|
||||
async inferSingle(
|
||||
inputIds: number[],
|
||||
attentionMask: number[],
|
||||
tokenTypeIds?: number[]
|
||||
): Promise<Float32Array> {
|
||||
return this.infer(
|
||||
[inputIds],
|
||||
[attentionMask],
|
||||
tokenTypeIds ? [tokenTypeIds] : undefined
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if initialized
|
||||
*/
|
||||
isInitialized(): boolean {
|
||||
return this.initialized
|
||||
}
|
||||
|
||||
/**
|
||||
* Get model input/output names (for debugging)
|
||||
*/
|
||||
getModelInfo(): { inputs: readonly string[]; outputs: readonly string[] } | null {
|
||||
if (!this.session) {
|
||||
return null
|
||||
}
|
||||
|
||||
return {
|
||||
inputs: this.session.inputNames,
|
||||
outputs: this.session.outputNames,
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Dispose of the session and free resources
|
||||
*/
|
||||
async dispose(): Promise<void> {
|
||||
if (this.session) {
|
||||
// Release the session
|
||||
this.session = null
|
||||
}
|
||||
this.initialized = false
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create an inference engine with default configuration
|
||||
*/
|
||||
export function createInferenceEngine(modelPath: string): ONNXInferenceEngine {
|
||||
return new ONNXInferenceEngine({ modelPath })
|
||||
}
|
||||
291
src/embeddings/wasm/WASMEmbeddingEngine.ts
Normal file
291
src/embeddings/wasm/WASMEmbeddingEngine.ts
Normal file
|
|
@ -0,0 +1,291 @@
|
|||
/**
|
||||
* WASM Embedding Engine
|
||||
*
|
||||
* The main embedding engine that combines all components:
|
||||
* - WordPieceTokenizer: Text → Token IDs
|
||||
* - ONNXInferenceEngine: Token IDs → Hidden States
|
||||
* - EmbeddingPostProcessor: Hidden States → Normalized Embedding
|
||||
*
|
||||
* This replaces transformers.js with a clean, production-grade implementation.
|
||||
*
|
||||
* Features:
|
||||
* - Singleton pattern (one model instance)
|
||||
* - Lazy initialization
|
||||
* - Batch processing support
|
||||
* - Zero runtime dependencies
|
||||
*/
|
||||
|
||||
import { WordPieceTokenizer } from './WordPieceTokenizer.js'
|
||||
import { ONNXInferenceEngine } from './ONNXInferenceEngine.js'
|
||||
import { EmbeddingPostProcessor } from './EmbeddingPostProcessor.js'
|
||||
import { getAssetLoader } from './AssetLoader.js'
|
||||
import { EmbeddingResult, EngineStats, MODEL_CONSTANTS } from './types.js'
|
||||
|
||||
// Global singleton instance
|
||||
let globalInstance: WASMEmbeddingEngine | null = null
|
||||
let globalInitPromise: Promise<void> | null = null
|
||||
|
||||
/**
|
||||
* WASM-based embedding engine
|
||||
*/
|
||||
export class WASMEmbeddingEngine {
|
||||
private tokenizer: WordPieceTokenizer | null = null
|
||||
private inference: ONNXInferenceEngine | null = null
|
||||
private postProcessor: EmbeddingPostProcessor | null = null
|
||||
private initialized = false
|
||||
private embedCount = 0
|
||||
private totalProcessingTimeMs = 0
|
||||
|
||||
private constructor() {
|
||||
// Private constructor for singleton
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the singleton instance
|
||||
*/
|
||||
static getInstance(): WASMEmbeddingEngine {
|
||||
if (!globalInstance) {
|
||||
globalInstance = new WASMEmbeddingEngine()
|
||||
}
|
||||
return globalInstance
|
||||
}
|
||||
|
||||
/**
|
||||
* Initialize all components
|
||||
*/
|
||||
async initialize(): Promise<void> {
|
||||
// Already initialized
|
||||
if (this.initialized) {
|
||||
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 WASM Embedding Engine...')
|
||||
|
||||
try {
|
||||
const assetLoader = getAssetLoader()
|
||||
|
||||
// Verify assets exist
|
||||
const verification = await assetLoader.verifyAssets()
|
||||
if (!verification.valid) {
|
||||
throw new Error(
|
||||
`Missing model assets:\n${verification.errors.join('\n')}\n\n` +
|
||||
`Expected model at: ${verification.modelPath}\n` +
|
||||
`Expected vocab at: ${verification.vocabPath}\n\n` +
|
||||
`Run 'npm run download-model' to download the model files.`
|
||||
)
|
||||
}
|
||||
|
||||
// Load vocabulary and create tokenizer
|
||||
console.log('📖 Loading vocabulary...')
|
||||
const vocab = await assetLoader.loadVocab()
|
||||
this.tokenizer = new WordPieceTokenizer(vocab)
|
||||
console.log(`✅ Vocabulary loaded: ${this.tokenizer.vocabSize} tokens`)
|
||||
|
||||
// Initialize ONNX inference engine
|
||||
console.log('🧠 Loading ONNX model...')
|
||||
const modelPath = await assetLoader.getModelPath()
|
||||
this.inference = new ONNXInferenceEngine({ modelPath })
|
||||
await this.inference.initialize(modelPath)
|
||||
console.log('✅ ONNX model loaded')
|
||||
|
||||
// Create post-processor
|
||||
this.postProcessor = new EmbeddingPostProcessor(MODEL_CONSTANTS.HIDDEN_SIZE)
|
||||
|
||||
this.initialized = true
|
||||
const initTime = Date.now() - startTime
|
||||
console.log(`✅ WASM Embedding Engine ready in ${initTime}ms`)
|
||||
|
||||
} catch (error) {
|
||||
this.initialized = false
|
||||
this.tokenizer = null
|
||||
this.inference = null
|
||||
this.postProcessor = null
|
||||
throw new Error(
|
||||
`Failed to initialize WASM Embedding Engine: ${error instanceof Error ? error.message : String(error)}`
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate embedding for text
|
||||
*/
|
||||
async embed(text: string): Promise<number[]> {
|
||||
const result = await this.embedWithMetadata(text)
|
||||
return result.embedding
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate embedding with metadata
|
||||
*/
|
||||
async embedWithMetadata(text: string): Promise<EmbeddingResult> {
|
||||
// Ensure initialized
|
||||
if (!this.initialized) {
|
||||
await this.initialize()
|
||||
}
|
||||
|
||||
if (!this.tokenizer || !this.inference || !this.postProcessor) {
|
||||
throw new Error('Engine not properly initialized')
|
||||
}
|
||||
|
||||
const startTime = Date.now()
|
||||
|
||||
// 1. Tokenize
|
||||
const tokenized = this.tokenizer.encode(text)
|
||||
|
||||
// 2. Run inference
|
||||
const hiddenStates = await this.inference.inferSingle(
|
||||
tokenized.inputIds,
|
||||
tokenized.attentionMask,
|
||||
tokenized.tokenTypeIds
|
||||
)
|
||||
|
||||
// 3. Post-process (mean pool + normalize)
|
||||
const embedding = this.postProcessor.process(
|
||||
hiddenStates,
|
||||
tokenized.attentionMask,
|
||||
tokenized.inputIds.length
|
||||
)
|
||||
|
||||
const processingTimeMs = Date.now() - startTime
|
||||
this.embedCount++
|
||||
this.totalProcessingTimeMs += processingTimeMs
|
||||
|
||||
return {
|
||||
embedding: Array.from(embedding),
|
||||
tokenCount: tokenized.tokenCount,
|
||||
processingTimeMs,
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Batch embed multiple texts
|
||||
*/
|
||||
async embedBatch(texts: string[]): Promise<number[][]> {
|
||||
// Ensure initialized
|
||||
if (!this.initialized) {
|
||||
await this.initialize()
|
||||
}
|
||||
|
||||
if (!this.tokenizer || !this.inference || !this.postProcessor) {
|
||||
throw new Error('Engine not properly initialized')
|
||||
}
|
||||
|
||||
if (texts.length === 0) {
|
||||
return []
|
||||
}
|
||||
|
||||
// Tokenize all texts
|
||||
const batch = this.tokenizer.encodeBatch(texts)
|
||||
const seqLen = batch.inputIds[0].length
|
||||
|
||||
// Run batch inference
|
||||
const hiddenStates = await this.inference.infer(
|
||||
batch.inputIds,
|
||||
batch.attentionMask,
|
||||
batch.tokenTypeIds
|
||||
)
|
||||
|
||||
// Post-process each result
|
||||
const embeddings = this.postProcessor.processBatch(
|
||||
hiddenStates,
|
||||
batch.attentionMask,
|
||||
texts.length,
|
||||
seqLen
|
||||
)
|
||||
|
||||
this.embedCount += texts.length
|
||||
|
||||
return embeddings.map(e => Array.from(e))
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if initialized
|
||||
*/
|
||||
isInitialized(): boolean {
|
||||
return this.initialized
|
||||
}
|
||||
|
||||
/**
|
||||
* Get engine statistics
|
||||
*/
|
||||
getStats(): EngineStats {
|
||||
return {
|
||||
initialized: this.initialized,
|
||||
embedCount: this.embedCount,
|
||||
totalProcessingTimeMs: this.totalProcessingTimeMs,
|
||||
avgProcessingTimeMs: this.embedCount > 0
|
||||
? this.totalProcessingTimeMs / this.embedCount
|
||||
: 0,
|
||||
modelName: MODEL_CONSTANTS.MODEL_NAME,
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Dispose and free resources
|
||||
*/
|
||||
async dispose(): Promise<void> {
|
||||
if (this.inference) {
|
||||
await this.inference.dispose()
|
||||
this.inference = null
|
||||
}
|
||||
this.tokenizer = null
|
||||
this.postProcessor = null
|
||||
this.initialized = false
|
||||
}
|
||||
|
||||
/**
|
||||
* Reset singleton (for testing)
|
||||
*/
|
||||
static resetInstance(): void {
|
||||
if (globalInstance) {
|
||||
globalInstance.dispose()
|
||||
}
|
||||
globalInstance = null
|
||||
globalInitPromise = null
|
||||
}
|
||||
}
|
||||
|
||||
// Export singleton access
|
||||
export const wasmEmbeddingEngine = WASMEmbeddingEngine.getInstance()
|
||||
|
||||
/**
|
||||
* Convenience function to get embeddings
|
||||
*/
|
||||
export async function embed(text: string): Promise<number[]> {
|
||||
return wasmEmbeddingEngine.embed(text)
|
||||
}
|
||||
|
||||
/**
|
||||
* Convenience function for batch embeddings
|
||||
*/
|
||||
export async function embedBatch(texts: string[]): Promise<number[][]> {
|
||||
return wasmEmbeddingEngine.embedBatch(texts)
|
||||
}
|
||||
|
||||
/**
|
||||
* Get embedding stats
|
||||
*/
|
||||
export function getEmbeddingStats(): EngineStats {
|
||||
return wasmEmbeddingEngine.getStats()
|
||||
}
|
||||
316
src/embeddings/wasm/WordPieceTokenizer.ts
Normal file
316
src/embeddings/wasm/WordPieceTokenizer.ts
Normal file
|
|
@ -0,0 +1,316 @@
|
|||
/**
|
||||
* WordPiece Tokenizer for BERT-based models
|
||||
*
|
||||
* Implements the WordPiece tokenization algorithm used by all-MiniLM-L6-v2.
|
||||
* This is a clean, dependency-free implementation.
|
||||
*
|
||||
* Algorithm:
|
||||
* 1. Normalize text (lowercase for uncased models)
|
||||
* 2. Split on whitespace and punctuation
|
||||
* 3. Apply WordPiece subword tokenization
|
||||
* 4. Add special tokens ([CLS], [SEP])
|
||||
* 5. Generate attention mask
|
||||
*/
|
||||
|
||||
import {
|
||||
TokenizerConfig,
|
||||
TokenizedInput,
|
||||
SPECIAL_TOKENS,
|
||||
MODEL_CONSTANTS,
|
||||
} from './types.js'
|
||||
|
||||
/**
|
||||
* WordPiece tokenizer for BERT-based sentence transformers
|
||||
*/
|
||||
export class WordPieceTokenizer {
|
||||
private vocab: Map<string, number>
|
||||
private reverseVocab: Map<number, string>
|
||||
private config: TokenizerConfig
|
||||
|
||||
constructor(vocab: Map<string, number> | Record<string, number>, config?: Partial<TokenizerConfig>) {
|
||||
// Convert Record to Map if needed
|
||||
this.vocab = vocab instanceof Map ? vocab : new Map(Object.entries(vocab))
|
||||
|
||||
// Build reverse vocab for debugging
|
||||
this.reverseVocab = new Map()
|
||||
for (const [token, id] of this.vocab) {
|
||||
this.reverseVocab.set(id, token)
|
||||
}
|
||||
|
||||
// Default config for all-MiniLM-L6-v2
|
||||
this.config = {
|
||||
vocab: this.vocab,
|
||||
unkTokenId: config?.unkTokenId ?? SPECIAL_TOKENS.UNK,
|
||||
clsTokenId: config?.clsTokenId ?? SPECIAL_TOKENS.CLS,
|
||||
sepTokenId: config?.sepTokenId ?? SPECIAL_TOKENS.SEP,
|
||||
padTokenId: config?.padTokenId ?? SPECIAL_TOKENS.PAD,
|
||||
maxLength: config?.maxLength ?? MODEL_CONSTANTS.MAX_SEQUENCE_LENGTH,
|
||||
doLowerCase: config?.doLowerCase ?? true,
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Tokenize text into token IDs
|
||||
*/
|
||||
encode(text: string): TokenizedInput {
|
||||
// 1. Normalize
|
||||
let normalizedText = text
|
||||
if (this.config.doLowerCase) {
|
||||
normalizedText = text.toLowerCase()
|
||||
}
|
||||
|
||||
// 2. Clean and split into words
|
||||
const words = this.basicTokenize(normalizedText)
|
||||
|
||||
// 3. Apply WordPiece to each word
|
||||
const tokens: number[] = [this.config.clsTokenId]
|
||||
|
||||
for (const word of words) {
|
||||
const wordTokens = this.wordPieceTokenize(word)
|
||||
// Check if adding these tokens would exceed max length (accounting for [SEP])
|
||||
if (tokens.length + wordTokens.length + 1 > this.config.maxLength) {
|
||||
break
|
||||
}
|
||||
tokens.push(...wordTokens)
|
||||
}
|
||||
|
||||
tokens.push(this.config.sepTokenId)
|
||||
|
||||
// 4. Generate attention mask and token type IDs
|
||||
const attentionMask = new Array(tokens.length).fill(1)
|
||||
const tokenTypeIds = new Array(tokens.length).fill(0)
|
||||
|
||||
return {
|
||||
inputIds: tokens,
|
||||
attentionMask,
|
||||
tokenTypeIds,
|
||||
tokenCount: tokens.length - 2, // Exclude [CLS] and [SEP]
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Encode with padding to fixed length
|
||||
*/
|
||||
encodeWithPadding(text: string, targetLength?: number): TokenizedInput {
|
||||
const result = this.encode(text)
|
||||
const padLength = targetLength ?? this.config.maxLength
|
||||
|
||||
// Pad to target length
|
||||
while (result.inputIds.length < padLength) {
|
||||
result.inputIds.push(this.config.padTokenId)
|
||||
result.attentionMask.push(0)
|
||||
result.tokenTypeIds.push(0)
|
||||
}
|
||||
|
||||
// Truncate if longer (shouldn't happen with proper encode())
|
||||
if (result.inputIds.length > padLength) {
|
||||
result.inputIds.length = padLength
|
||||
result.attentionMask.length = padLength
|
||||
result.tokenTypeIds.length = padLength
|
||||
// Ensure [SEP] is at the end
|
||||
result.inputIds[padLength - 1] = this.config.sepTokenId
|
||||
result.attentionMask[padLength - 1] = 1
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
/**
|
||||
* Batch encode multiple texts
|
||||
*/
|
||||
encodeBatch(texts: string[]): {
|
||||
inputIds: number[][]
|
||||
attentionMask: number[][]
|
||||
tokenTypeIds: number[][]
|
||||
} {
|
||||
const results = texts.map((text) => this.encode(text))
|
||||
|
||||
// Find max length in batch
|
||||
const maxLen = Math.max(...results.map((r) => r.inputIds.length))
|
||||
|
||||
// Pad all to same length
|
||||
const inputIds: number[][] = []
|
||||
const attentionMask: number[][] = []
|
||||
const tokenTypeIds: number[][] = []
|
||||
|
||||
for (const result of results) {
|
||||
const padded = this.encodeWithPadding(
|
||||
'', // Not used since we're modifying result
|
||||
maxLen
|
||||
)
|
||||
// Copy original values
|
||||
for (let i = 0; i < result.inputIds.length; i++) {
|
||||
padded.inputIds[i] = result.inputIds[i]
|
||||
padded.attentionMask[i] = result.attentionMask[i]
|
||||
padded.tokenTypeIds[i] = result.tokenTypeIds[i]
|
||||
}
|
||||
// Pad the rest
|
||||
for (let i = result.inputIds.length; i < maxLen; i++) {
|
||||
padded.inputIds[i] = this.config.padTokenId
|
||||
padded.attentionMask[i] = 0
|
||||
padded.tokenTypeIds[i] = 0
|
||||
}
|
||||
inputIds.push(padded.inputIds.slice(0, maxLen))
|
||||
attentionMask.push(padded.attentionMask.slice(0, maxLen))
|
||||
tokenTypeIds.push(padded.tokenTypeIds.slice(0, maxLen))
|
||||
}
|
||||
|
||||
return { inputIds, attentionMask, tokenTypeIds }
|
||||
}
|
||||
|
||||
/**
|
||||
* Basic tokenization: split on whitespace and punctuation
|
||||
*/
|
||||
private basicTokenize(text: string): string[] {
|
||||
// Clean whitespace
|
||||
text = text.trim().replace(/\s+/g, ' ')
|
||||
|
||||
if (!text) {
|
||||
return []
|
||||
}
|
||||
|
||||
const words: string[] = []
|
||||
let currentWord = ''
|
||||
|
||||
for (const char of text) {
|
||||
if (this.isWhitespace(char)) {
|
||||
if (currentWord) {
|
||||
words.push(currentWord)
|
||||
currentWord = ''
|
||||
}
|
||||
} else if (this.isPunctuation(char)) {
|
||||
if (currentWord) {
|
||||
words.push(currentWord)
|
||||
currentWord = ''
|
||||
}
|
||||
words.push(char)
|
||||
} else {
|
||||
currentWord += char
|
||||
}
|
||||
}
|
||||
|
||||
if (currentWord) {
|
||||
words.push(currentWord)
|
||||
}
|
||||
|
||||
return words
|
||||
}
|
||||
|
||||
/**
|
||||
* WordPiece tokenization for a single word
|
||||
*/
|
||||
private wordPieceTokenize(word: string): number[] {
|
||||
if (!word) {
|
||||
return []
|
||||
}
|
||||
|
||||
// Check if whole word is in vocabulary
|
||||
if (this.vocab.has(word)) {
|
||||
return [this.vocab.get(word)!]
|
||||
}
|
||||
|
||||
const tokens: number[] = []
|
||||
let start = 0
|
||||
|
||||
while (start < word.length) {
|
||||
let end = word.length
|
||||
let foundToken = false
|
||||
|
||||
while (start < end) {
|
||||
let substr = word.slice(start, end)
|
||||
|
||||
// Add ## prefix for subwords (not at start of word)
|
||||
if (start > 0) {
|
||||
substr = '##' + substr
|
||||
}
|
||||
|
||||
if (this.vocab.has(substr)) {
|
||||
tokens.push(this.vocab.get(substr)!)
|
||||
foundToken = true
|
||||
break
|
||||
}
|
||||
|
||||
end--
|
||||
}
|
||||
|
||||
if (!foundToken) {
|
||||
// Unknown character - use [UNK] for single character
|
||||
tokens.push(this.config.unkTokenId)
|
||||
start++
|
||||
} else {
|
||||
start = end
|
||||
}
|
||||
}
|
||||
|
||||
return tokens
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if character is whitespace
|
||||
*/
|
||||
private isWhitespace(char: string): boolean {
|
||||
return /\s/.test(char)
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if character is punctuation
|
||||
*/
|
||||
private isPunctuation(char: string): boolean {
|
||||
const code = char.charCodeAt(0)
|
||||
// ASCII punctuation ranges
|
||||
if (
|
||||
(code >= 33 && code <= 47) || // !"#$%&'()*+,-./
|
||||
(code >= 58 && code <= 64) || // :;<=>?@
|
||||
(code >= 91 && code <= 96) || // [\]^_`
|
||||
(code >= 123 && code <= 126) // {|}~
|
||||
) {
|
||||
return true
|
||||
}
|
||||
// Unicode punctuation categories
|
||||
return /[\u2000-\u206F\u2E00-\u2E7F\\'!"#$%&()*+,\-./:;<=>?@\[\]^_`{|}~]/.test(char)
|
||||
}
|
||||
|
||||
/**
|
||||
* Decode token IDs back to text (for debugging)
|
||||
*/
|
||||
decode(tokenIds: number[]): string {
|
||||
const tokens: string[] = []
|
||||
for (const id of tokenIds) {
|
||||
const token = this.reverseVocab.get(id)
|
||||
if (token && !['[CLS]', '[SEP]', '[PAD]'].includes(token)) {
|
||||
if (token.startsWith('##')) {
|
||||
// Subword - append without space
|
||||
if (tokens.length > 0) {
|
||||
tokens[tokens.length - 1] += token.slice(2)
|
||||
} else {
|
||||
tokens.push(token.slice(2))
|
||||
}
|
||||
} else {
|
||||
tokens.push(token)
|
||||
}
|
||||
}
|
||||
}
|
||||
return tokens.join(' ')
|
||||
}
|
||||
|
||||
/**
|
||||
* Get vocabulary size
|
||||
*/
|
||||
get vocabSize(): number {
|
||||
return this.vocab.size
|
||||
}
|
||||
|
||||
/**
|
||||
* Get max sequence length
|
||||
*/
|
||||
get maxLength(): number {
|
||||
return this.config.maxLength
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create tokenizer from vocabulary JSON
|
||||
*/
|
||||
export function createTokenizer(vocabJson: Record<string, number>): WordPieceTokenizer {
|
||||
return new WordPieceTokenizer(vocabJson)
|
||||
}
|
||||
33
src/embeddings/wasm/index.ts
Normal file
33
src/embeddings/wasm/index.ts
Normal file
|
|
@ -0,0 +1,33 @@
|
|||
/**
|
||||
* WASM Embedding Engine - Public Exports
|
||||
*
|
||||
* Clean, production-grade embedding engine using direct ONNX WASM.
|
||||
* No transformers.js dependency, no runtime downloads, works everywhere.
|
||||
*/
|
||||
|
||||
// Main engine
|
||||
export {
|
||||
WASMEmbeddingEngine,
|
||||
wasmEmbeddingEngine,
|
||||
embed,
|
||||
embedBatch,
|
||||
getEmbeddingStats,
|
||||
} from './WASMEmbeddingEngine.js'
|
||||
|
||||
// Components (for advanced use)
|
||||
export { WordPieceTokenizer, createTokenizer } from './WordPieceTokenizer.js'
|
||||
export { ONNXInferenceEngine, createInferenceEngine } from './ONNXInferenceEngine.js'
|
||||
export { EmbeddingPostProcessor, createPostProcessor } from './EmbeddingPostProcessor.js'
|
||||
export { AssetLoader, getAssetLoader, createAssetLoader } from './AssetLoader.js'
|
||||
|
||||
// Types
|
||||
export type {
|
||||
TokenizerConfig,
|
||||
TokenizedInput,
|
||||
InferenceConfig,
|
||||
EmbeddingResult,
|
||||
EngineStats,
|
||||
ModelConfig,
|
||||
} from './types.js'
|
||||
|
||||
export { SPECIAL_TOKENS, MODEL_CONSTANTS } from './types.js'
|
||||
122
src/embeddings/wasm/types.ts
Normal file
122
src/embeddings/wasm/types.ts
Normal file
|
|
@ -0,0 +1,122 @@
|
|||
/**
|
||||
* Type definitions for WASM Embedding Engine
|
||||
*
|
||||
* Clean, production-grade types for direct ONNX WASM embeddings.
|
||||
*/
|
||||
|
||||
/**
|
||||
* Tokenizer configuration for WordPiece
|
||||
*/
|
||||
export interface TokenizerConfig {
|
||||
/** Vocabulary mapping word → token ID */
|
||||
vocab: Map<string, number>
|
||||
/** [UNK] token ID (100 for BERT-based models) */
|
||||
unkTokenId: number
|
||||
/** [CLS] token ID (101 for BERT-based models) */
|
||||
clsTokenId: number
|
||||
/** [SEP] token ID (102 for BERT-based models) */
|
||||
sepTokenId: number
|
||||
/** [PAD] token ID (0 for BERT-based models) */
|
||||
padTokenId: number
|
||||
/** Maximum sequence length (512 for all-MiniLM-L6-v2) */
|
||||
maxLength: number
|
||||
/** Whether to lowercase input (true for uncased models) */
|
||||
doLowerCase: boolean
|
||||
}
|
||||
|
||||
/**
|
||||
* Result of tokenization
|
||||
*/
|
||||
export interface TokenizedInput {
|
||||
/** Token IDs including [CLS] and [SEP] */
|
||||
inputIds: number[]
|
||||
/** Attention mask (1 for real tokens, 0 for padding) */
|
||||
attentionMask: number[]
|
||||
/** Token type IDs (all 0 for single sentence) */
|
||||
tokenTypeIds: number[]
|
||||
/** Number of tokens (excluding special tokens) */
|
||||
tokenCount: number
|
||||
}
|
||||
|
||||
/**
|
||||
* ONNX inference engine configuration
|
||||
*/
|
||||
export interface InferenceConfig {
|
||||
/** Path to ONNX model file */
|
||||
modelPath: string
|
||||
/** Path to WASM files directory */
|
||||
wasmPath?: string
|
||||
/** Number of threads (1 for universal compatibility) */
|
||||
numThreads: number
|
||||
/** Enable SIMD if available */
|
||||
enableSimd: boolean
|
||||
/** Enable CPU memory arena (false for memory efficiency) */
|
||||
enableCpuMemArena: boolean
|
||||
}
|
||||
|
||||
/**
|
||||
* Embedding result with metadata
|
||||
*/
|
||||
export interface EmbeddingResult {
|
||||
/** 384-dimensional embedding vector */
|
||||
embedding: number[]
|
||||
/** Number of tokens processed */
|
||||
tokenCount: number
|
||||
/** Processing time in milliseconds */
|
||||
processingTimeMs: number
|
||||
}
|
||||
|
||||
/**
|
||||
* Engine statistics
|
||||
*/
|
||||
export interface EngineStats {
|
||||
/** Whether the engine is initialized */
|
||||
initialized: boolean
|
||||
/** Total number of embeddings generated */
|
||||
embedCount: number
|
||||
/** Total processing time in milliseconds */
|
||||
totalProcessingTimeMs: number
|
||||
/** Average processing time per embedding */
|
||||
avgProcessingTimeMs: number
|
||||
/** Model name */
|
||||
modelName: string
|
||||
}
|
||||
|
||||
/**
|
||||
* Model configuration (from config.json)
|
||||
*/
|
||||
export interface ModelConfig {
|
||||
/** Model architecture type */
|
||||
architectures: string[]
|
||||
/** Hidden size (384 for all-MiniLM-L6-v2) */
|
||||
hidden_size: number
|
||||
/** Number of attention heads */
|
||||
num_attention_heads: number
|
||||
/** Number of hidden layers */
|
||||
num_hidden_layers: number
|
||||
/** Vocabulary size */
|
||||
vocab_size: number
|
||||
/** Maximum position embeddings */
|
||||
max_position_embeddings: number
|
||||
}
|
||||
|
||||
/**
|
||||
* Special token IDs for BERT-based models
|
||||
*/
|
||||
export const SPECIAL_TOKENS = {
|
||||
PAD: 0,
|
||||
UNK: 100,
|
||||
CLS: 101,
|
||||
SEP: 102,
|
||||
MASK: 103,
|
||||
} as const
|
||||
|
||||
/**
|
||||
* Model constants for all-MiniLM-L6-v2
|
||||
*/
|
||||
export const MODEL_CONSTANTS = {
|
||||
HIDDEN_SIZE: 384,
|
||||
MAX_SEQUENCE_LENGTH: 512,
|
||||
VOCAB_SIZE: 30522,
|
||||
MODEL_NAME: 'all-MiniLM-L6-v2',
|
||||
} as const
|
||||
51
src/setup.ts
51
src/setup.ts
|
|
@ -1,46 +1,31 @@
|
|||
/**
|
||||
* CRITICAL: This file is imported for its side effects to patch the environment
|
||||
* for Node.js compatibility before any other library code runs.
|
||||
* Brainy Setup - Minimal Polyfills
|
||||
*
|
||||
* It ensures that by the time Transformers.js/ONNX Runtime is imported by any other
|
||||
* module, the necessary compatibility fixes for the current Node.js
|
||||
* environment are already in place.
|
||||
* ARCHITECTURE (v7.0.0):
|
||||
* Brainy uses direct ONNX WASM for embeddings.
|
||||
* No transformers.js dependency, no hacks required.
|
||||
*
|
||||
* This file MUST be imported as the first import in unified.ts to prevent
|
||||
* race conditions with library initialization. Failure to do so may
|
||||
* result in errors like "TextEncoder is not a constructor" when the package
|
||||
* is used in Node.js environments.
|
||||
* This file provides minimal polyfills for cross-environment compatibility:
|
||||
* - TextEncoder/TextDecoder for older environments
|
||||
*
|
||||
* The package.json file marks this file as having side effects to prevent
|
||||
* tree-shaking by bundlers, ensuring the patch is always applied.
|
||||
* BENEFITS:
|
||||
* - Clean codebase with no workarounds
|
||||
* - Works everywhere: Node.js, Bun, Bun --compile, browsers, Deno
|
||||
* - No platform-specific binaries
|
||||
* - Model bundled in package (no runtime downloads)
|
||||
*/
|
||||
|
||||
// Get the appropriate global object for the current environment
|
||||
const globalObj = (() => {
|
||||
if (typeof globalThis !== 'undefined') return globalThis
|
||||
if (typeof global !== 'undefined') return global
|
||||
if (typeof self !== 'undefined') return self
|
||||
return null // No global object available
|
||||
})()
|
||||
// ============================================================================
|
||||
// TextEncoder/TextDecoder Polyfills
|
||||
// ============================================================================
|
||||
const globalObj = globalThis ?? global ?? self
|
||||
|
||||
// Define TextEncoder and TextDecoder globally to make sure they're available
|
||||
// Now works across all environments: Node.js, serverless, and other server environments
|
||||
if (globalObj) {
|
||||
if (!globalObj.TextEncoder) {
|
||||
globalObj.TextEncoder = TextEncoder
|
||||
}
|
||||
if (!globalObj.TextDecoder) {
|
||||
globalObj.TextDecoder = TextDecoder
|
||||
}
|
||||
|
||||
// Create special global constructors for library compatibility
|
||||
(globalObj as any).__TextEncoder__ = TextEncoder
|
||||
if (!globalObj.TextEncoder) globalObj.TextEncoder = TextEncoder
|
||||
if (!globalObj.TextDecoder) globalObj.TextDecoder = TextDecoder
|
||||
;(globalObj as any).__TextEncoder__ = TextEncoder
|
||||
;(globalObj as any).__TextDecoder__ = TextDecoder
|
||||
}
|
||||
|
||||
// Also import normally for ES modules environments
|
||||
import { applyTensorFlowPatch } from './utils/textEncoding.js'
|
||||
|
||||
// Apply the TextEncoder/TextDecoder compatibility patch
|
||||
applyTensorFlowPatch()
|
||||
console.log('Applied TextEncoder/TextDecoder patch via ES modules in setup.ts')
|
||||
|
|
|
|||
|
|
@ -1,271 +1,47 @@
|
|||
/**
|
||||
* Embedding functions for converting data to vectors using Transformers.js
|
||||
* Complete rewrite to eliminate TensorFlow.js and use ONNX-based models
|
||||
* Embedding functions for converting data to vectors
|
||||
*
|
||||
* Uses direct ONNX WASM for universal compatibility.
|
||||
* No transformers.js dependency - clean, production-grade implementation.
|
||||
*/
|
||||
|
||||
import { EmbeddingFunction, EmbeddingModel, Vector } from '../coreTypes.js'
|
||||
import { executeInThread } from './workerUtils.js'
|
||||
import { isBrowser } from './environment.js'
|
||||
// @ts-ignore - Transformers.js is now the primary embedding library
|
||||
import { pipeline, env } from '@huggingface/transformers'
|
||||
|
||||
// CRITICAL: Disable ONNX memory arena to prevent 4-8GB allocation
|
||||
// This is needed for BOTH production and testing - reduces memory by 50-75%
|
||||
if (typeof process !== 'undefined' && process.env) {
|
||||
process.env.ORT_DISABLE_MEMORY_ARENA = '1'
|
||||
process.env.ORT_DISABLE_MEMORY_PATTERN = '1'
|
||||
// Force single-threaded operation for maximum stability (Node.js 22 LTS)
|
||||
process.env.ORT_INTRA_OP_NUM_THREADS = '1' // Single thread for operators
|
||||
process.env.ORT_INTER_OP_NUM_THREADS = '1' // Single thread for sessions
|
||||
process.env.ORT_NUM_THREADS = '1' // Additional safety override
|
||||
}
|
||||
import { embeddingManager } from '../embeddings/EmbeddingManager.js'
|
||||
|
||||
/**
|
||||
* Detect the best available GPU device for the current environment
|
||||
*/
|
||||
export async function detectBestDevice(): Promise<'cpu' | 'webgpu' | 'cuda'> {
|
||||
// Browser environment - check for WebGPU support
|
||||
if (isBrowser()) {
|
||||
if (typeof navigator !== 'undefined' && 'gpu' in navigator) {
|
||||
try {
|
||||
const adapter = await (navigator as any).gpu?.requestAdapter()
|
||||
if (adapter) {
|
||||
return 'webgpu'
|
||||
}
|
||||
} catch (error) {
|
||||
// WebGPU not available or failed to initialize
|
||||
}
|
||||
}
|
||||
return 'cpu'
|
||||
}
|
||||
|
||||
// Node.js environment - check for CUDA support
|
||||
try {
|
||||
// Check if ONNX Runtime GPU packages are available
|
||||
// This is a simple heuristic - in production you might want more sophisticated detection
|
||||
const hasGpu = process.env.CUDA_VISIBLE_DEVICES !== undefined ||
|
||||
process.env.ONNXRUNTIME_GPU_ENABLED === 'true'
|
||||
return hasGpu ? 'cuda' : 'cpu'
|
||||
} catch (error) {
|
||||
return 'cpu'
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve device string to actual device configuration
|
||||
*/
|
||||
export async function resolveDevice(device: string = 'auto'): Promise<string> {
|
||||
if (device === 'auto') {
|
||||
return await detectBestDevice()
|
||||
}
|
||||
|
||||
// Map 'gpu' to appropriate GPU type for current environment
|
||||
if (device === 'gpu') {
|
||||
const detected = await detectBestDevice()
|
||||
return detected === 'cpu' ? 'cpu' : detected
|
||||
}
|
||||
|
||||
return device
|
||||
}
|
||||
|
||||
/**
|
||||
* Transformers.js Sentence Encoder embedding model
|
||||
* Uses ONNX Runtime for fast, offline embeddings with smaller models
|
||||
* Default model: all-MiniLM-L6-v2 (384 dimensions, ~90MB)
|
||||
* TransformerEmbedding options (kept for backward compatibility)
|
||||
*/
|
||||
export interface TransformerEmbeddingOptions {
|
||||
/** Model name/path to use - defaults to all-MiniLM-L6-v2 */
|
||||
/** Model name - only all-MiniLM-L6-v2 is supported */
|
||||
model?: string
|
||||
/** Whether to enable verbose logging */
|
||||
verbose?: boolean
|
||||
/** Custom cache directory for models */
|
||||
/** Custom cache directory - ignored (model is bundled) */
|
||||
cacheDir?: string
|
||||
/** Force local files only (no downloads) */
|
||||
/** Force local files only - ignored (model is bundled) */
|
||||
localFilesOnly?: boolean
|
||||
/** Model precision: 'q8' = 75% smaller quantized model, 'fp32' = full precision (default) */
|
||||
/** Model precision - always q8 */
|
||||
precision?: 'fp32' | 'q8'
|
||||
/** Device to run inference on - 'auto' detects best available */
|
||||
/** Device - always WASM */
|
||||
device?: 'auto' | 'cpu' | 'webgpu' | 'cuda' | 'gpu'
|
||||
}
|
||||
|
||||
/**
|
||||
* TransformerEmbedding - Sentence embeddings using WASM ONNX
|
||||
*
|
||||
* This class delegates all work to EmbeddingManager which uses
|
||||
* the direct ONNX WASM engine. Kept for backward compatibility.
|
||||
*/
|
||||
export class TransformerEmbedding implements EmbeddingModel {
|
||||
private extractor: any = null
|
||||
private initialized = false
|
||||
private verbose: boolean = true
|
||||
private options: Required<TransformerEmbeddingOptions>
|
||||
private verbose: boolean
|
||||
|
||||
/**
|
||||
* Create a new TransformerEmbedding instance
|
||||
*/
|
||||
constructor(options: TransformerEmbeddingOptions = {}) {
|
||||
this.verbose = options.verbose !== undefined ? options.verbose : true
|
||||
|
||||
// PRODUCTION-READY MODEL CONFIGURATION
|
||||
// Priority order: explicit option > environment variable > smart default
|
||||
|
||||
let localFilesOnly: boolean
|
||||
|
||||
if (options.localFilesOnly !== undefined) {
|
||||
// 1. Explicit option takes highest priority
|
||||
localFilesOnly = options.localFilesOnly
|
||||
} else if (process.env.BRAINY_ALLOW_REMOTE_MODELS === 'false') {
|
||||
// 2. Environment variable explicitly disables remote models (legacy support)
|
||||
localFilesOnly = true
|
||||
} else if (process.env.NODE_ENV === 'development') {
|
||||
// 3. Development mode allows remote models
|
||||
localFilesOnly = false
|
||||
} else if (isBrowser()) {
|
||||
// 4. Browser defaults to allowing remote models
|
||||
localFilesOnly = false
|
||||
} else {
|
||||
// 5. Node.js production: try local first, but allow remote as fallback
|
||||
// This is the NEW production-friendly default
|
||||
localFilesOnly = false
|
||||
}
|
||||
|
||||
this.options = {
|
||||
model: options.model || 'Xenova/all-MiniLM-L6-v2',
|
||||
verbose: this.verbose,
|
||||
cacheDir: options.cacheDir || './models',
|
||||
localFilesOnly: localFilesOnly,
|
||||
precision: options.precision || 'fp32', // Clean and clear!
|
||||
device: options.device || 'auto'
|
||||
}
|
||||
|
||||
// ULTRA-CAREFUL: Runtime warnings for q8 usage
|
||||
if (this.options.precision === 'q8') {
|
||||
const confirmed = process.env.BRAINY_Q8_CONFIRMED === 'true'
|
||||
if (!confirmed && this.verbose) {
|
||||
console.warn('🚨 Q8 MODEL WARNING:')
|
||||
console.warn(' • Q8 creates different embeddings than fp32')
|
||||
console.warn(' • Q8 is incompatible with existing fp32 data')
|
||||
console.warn(' • Only use q8 for new projects or when explicitly migrating')
|
||||
console.warn(' • Set BRAINY_Q8_CONFIRMED=true to silence this warning')
|
||||
console.warn(' • Q8 model is 75% smaller but may have slightly reduced accuracy')
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
if (this.verbose) {
|
||||
this.logger('log', `Embedding config: precision=${this.options.precision}, localFilesOnly=${localFilesOnly}, model=${this.options.model}`)
|
||||
console.log('[TransformerEmbedding] Using WASM ONNX backend (delegating to EmbeddingManager)')
|
||||
}
|
||||
|
||||
// Configure transformers.js environment
|
||||
if (!isBrowser()) {
|
||||
// Set cache directory for Node.js
|
||||
env.cacheDir = this.options.cacheDir
|
||||
// Prioritize local models for offline operation
|
||||
env.allowRemoteModels = !this.options.localFilesOnly
|
||||
env.allowLocalModels = true
|
||||
} else {
|
||||
// Browser configuration
|
||||
// Allow both local and remote models, but prefer local if available
|
||||
env.allowLocalModels = true
|
||||
env.allowRemoteModels = true
|
||||
// Force the configuration to ensure it's applied
|
||||
if (this.verbose) {
|
||||
this.logger('log', `Browser env config - allowLocalModels: ${env.allowLocalModels}, allowRemoteModels: ${env.allowRemoteModels}, localFilesOnly: ${this.options.localFilesOnly}`)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the default cache directory for models
|
||||
*/
|
||||
private async getDefaultCacheDir(): Promise<string> {
|
||||
if (isBrowser()) {
|
||||
return './models' // Browser default
|
||||
}
|
||||
|
||||
// Check for bundled models in the package
|
||||
const possiblePaths = [
|
||||
// In the installed package
|
||||
'./node_modules/@soulcraft/brainy/models',
|
||||
// In development/source
|
||||
'./models',
|
||||
'./dist/../models',
|
||||
// Alternative locations
|
||||
'../models',
|
||||
'../../models'
|
||||
]
|
||||
|
||||
// Check if we're in Node.js and try to find the bundled models
|
||||
if (typeof process !== 'undefined' && process.versions?.node) {
|
||||
try {
|
||||
// Use dynamic import instead of require for ES modules compatibility
|
||||
const { createRequire } = await import('module')
|
||||
const require = createRequire(import.meta.url)
|
||||
|
||||
const path = require('node:path')
|
||||
const fs = require('node:fs')
|
||||
|
||||
// Try to resolve the package location
|
||||
try {
|
||||
const brainyPackagePath = require.resolve('@soulcraft/brainy/package.json')
|
||||
const brainyPackageDir = path.dirname(brainyPackagePath)
|
||||
const bundledModelsPath = path.join(brainyPackageDir, 'models')
|
||||
|
||||
if (fs.existsSync(bundledModelsPath)) {
|
||||
this.logger('log', `Using bundled models from package: ${bundledModelsPath}`)
|
||||
return bundledModelsPath
|
||||
}
|
||||
} catch (e) {
|
||||
// Not installed as package, continue
|
||||
}
|
||||
|
||||
// Try relative paths from current location
|
||||
for (const relativePath of possiblePaths) {
|
||||
const fullPath = path.resolve(relativePath)
|
||||
if (fs.existsSync(fullPath)) {
|
||||
this.logger('log', `Using bundled models from: ${fullPath}`)
|
||||
return fullPath
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
// Silently fall back to default path if module detection fails
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to default cache directory
|
||||
return './models'
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if we're running in a test environment
|
||||
*/
|
||||
private isTestEnvironment(): boolean {
|
||||
// Always use real implementation - no more mocking
|
||||
return false
|
||||
}
|
||||
|
||||
/**
|
||||
* Log message only if verbose mode is enabled
|
||||
*/
|
||||
private logger(level: 'log' | 'warn' | 'error', message: string, ...args: any[]): void {
|
||||
if (level === 'error' || this.verbose) {
|
||||
console[level](`[TransformerEmbedding] ${message}`, ...args)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate mock embeddings for unit tests
|
||||
*/
|
||||
private getMockEmbedding(data: string | string[]): Vector {
|
||||
// Use the same mock logic as setup-unit.ts for consistency
|
||||
const input = Array.isArray(data) ? data.join(' ') : data
|
||||
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
|
||||
}
|
||||
|
||||
return vector
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -276,127 +52,16 @@ export class TransformerEmbedding implements EmbeddingModel {
|
|||
return
|
||||
}
|
||||
|
||||
// In unit test mode, skip real model initialization to prevent ONNX conflicts
|
||||
if (process.env.BRAINY_UNIT_TEST === 'true' || (globalThis as any).__BRAINY_UNIT_TEST__) {
|
||||
this.initialized = true
|
||||
this.logger('log', '🧪 Using mocked embeddings for unit tests')
|
||||
return
|
||||
}
|
||||
|
||||
try {
|
||||
|
||||
// Resolve device configuration and cache directory
|
||||
const device = await resolveDevice(this.options.device)
|
||||
const cacheDir = this.options.cacheDir === './models'
|
||||
? await this.getDefaultCacheDir()
|
||||
: this.options.cacheDir
|
||||
|
||||
this.logger('log', `Loading Transformer model: ${this.options.model} on device: ${device}`)
|
||||
|
||||
const startTime = Date.now()
|
||||
|
||||
// Use the configured precision from EmbeddingManager
|
||||
const { embeddingManager } = await import('../embeddings/EmbeddingManager.js')
|
||||
let actualType = embeddingManager.getPrecision()
|
||||
|
||||
// CRITICAL: Control which model precision transformers.js uses
|
||||
// Q8 models use quantized int8 weights for 75% size reduction
|
||||
// Always use Q8 for optimal balance
|
||||
|
||||
actualType = 'q8' // Always Q8
|
||||
this.logger('log', '🎯 Using Q8 quantized model (75% smaller, 99% accuracy)')
|
||||
|
||||
// Load the feature extraction pipeline with memory optimizations
|
||||
const pipelineOptions: any = {
|
||||
cache_dir: cacheDir,
|
||||
local_files_only: isBrowser() ? false : this.options.localFilesOnly,
|
||||
// CRITICAL: Specify dtype for model precision
|
||||
dtype: 'q8',
|
||||
// CRITICAL: For Q8, explicitly use quantized model
|
||||
quantized: true,
|
||||
// CRITICAL: ONNX memory optimizations
|
||||
session_options: {
|
||||
enableCpuMemArena: false, // Disable pre-allocated memory arena
|
||||
enableMemPattern: false, // Disable memory pattern optimization
|
||||
interOpNumThreads: 1, // Force single thread for V8 stability
|
||||
intraOpNumThreads: 1, // Force single thread for V8 stability
|
||||
graphOptimizationLevel: 'disabled' // Disable threading optimizations
|
||||
}
|
||||
}
|
||||
|
||||
// Add device configuration for GPU acceleration
|
||||
if (device !== 'cpu') {
|
||||
pipelineOptions.device = device
|
||||
this.logger('log', `🚀 GPU acceleration enabled: ${device}`)
|
||||
}
|
||||
|
||||
if (this.verbose) {
|
||||
this.logger('log', `Pipeline options: ${JSON.stringify(pipelineOptions)}`)
|
||||
}
|
||||
|
||||
try {
|
||||
// For Q8 models, we need to explicitly specify the model file
|
||||
if (actualType === 'q8' && !isBrowser()) {
|
||||
try {
|
||||
// Check if quantized model exists (Node.js only)
|
||||
const { join } = await import('node:path')
|
||||
const { existsSync } = await import('node:fs')
|
||||
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')
|
||||
actualType = 'q8' // Always Q8
|
||||
}
|
||||
} catch (error) {
|
||||
// Skip model path check in browser or if imports fail
|
||||
this.logger('log', '🌐 Skipping local model check in browser environment')
|
||||
}
|
||||
}
|
||||
|
||||
this.extractor = await pipeline('feature-extraction', this.options.model, pipelineOptions)
|
||||
} catch (gpuError: any) {
|
||||
// Fallback to CPU if GPU initialization fails
|
||||
if (device !== 'cpu') {
|
||||
this.logger('warn', `GPU initialization failed, falling back to CPU: ${gpuError?.message || gpuError}`)
|
||||
const cpuOptions = { ...pipelineOptions }
|
||||
delete cpuOptions.device
|
||||
this.extractor = await pipeline('feature-extraction', this.options.model, cpuOptions)
|
||||
} else {
|
||||
// PRODUCTION-READY ERROR HANDLING
|
||||
// If local_files_only is true and models are missing, try enabling remote downloads
|
||||
if (pipelineOptions.local_files_only && gpuError?.message?.includes('local_files_only')) {
|
||||
this.logger('warn', 'Local models not found, attempting remote download as fallback...')
|
||||
|
||||
try {
|
||||
const remoteOptions = { ...pipelineOptions, local_files_only: false }
|
||||
this.extractor = await pipeline('feature-extraction', this.options.model, remoteOptions)
|
||||
this.logger('log', '✅ Successfully downloaded and loaded model from remote')
|
||||
|
||||
// Update the configuration to reflect what actually worked
|
||||
this.options.localFilesOnly = false
|
||||
} catch (remoteError: any) {
|
||||
// Both local and remote failed - throw comprehensive error
|
||||
const errorMsg = `Failed to load embedding model "${this.options.model}". ` +
|
||||
`Local models not found and remote download failed. ` +
|
||||
`To fix: 1) Run "npm run download-models", ` +
|
||||
`2) Check your internet connection, or ` +
|
||||
`3) Use a custom embedding function.`
|
||||
throw new Error(errorMsg)
|
||||
}
|
||||
} else {
|
||||
throw gpuError
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const loadTime = Date.now() - startTime
|
||||
this.logger('log', `✅ Model loaded successfully in ${loadTime}ms`)
|
||||
|
||||
await embeddingManager.init()
|
||||
this.initialized = true
|
||||
|
||||
if (this.verbose) {
|
||||
console.log('[TransformerEmbedding] Initialized via EmbeddingManager (WASM)')
|
||||
}
|
||||
} catch (error) {
|
||||
this.logger('error', 'Failed to initialize Transformer embedding model:', error)
|
||||
throw new Error(`Transformer embedding initialization failed: ${error}`)
|
||||
console.error('[TransformerEmbedding] Failed to initialize:', error)
|
||||
throw new Error(`TransformerEmbedding initialization failed: ${error}`)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -404,177 +69,93 @@ export class TransformerEmbedding implements EmbeddingModel {
|
|||
* Generate embeddings for text data
|
||||
*/
|
||||
public async embed(data: string | string[]): Promise<Vector> {
|
||||
// In unit test mode, return mock embeddings
|
||||
if (process.env.BRAINY_UNIT_TEST === 'true' || (globalThis as any).__BRAINY_UNIT_TEST__) {
|
||||
return this.getMockEmbedding(data)
|
||||
}
|
||||
|
||||
if (!this.initialized) {
|
||||
await this.init()
|
||||
}
|
||||
|
||||
try {
|
||||
// Handle different input types
|
||||
let textToEmbed: string[]
|
||||
|
||||
if (typeof data === 'string') {
|
||||
// Handle empty string case
|
||||
if (data.trim() === '') {
|
||||
// Return a zero vector of 384 dimensions (all-MiniLM-L6-v2 standard)
|
||||
return new Array(384).fill(0)
|
||||
}
|
||||
textToEmbed = [data]
|
||||
} else if (Array.isArray(data) && data.every((item) => typeof item === 'string')) {
|
||||
// Handle empty array or array with empty strings
|
||||
if (data.length === 0 || data.every((item) => item.trim() === '')) {
|
||||
return new Array(384).fill(0)
|
||||
}
|
||||
// Filter out empty strings
|
||||
textToEmbed = data.filter((item) => item.trim() !== '')
|
||||
if (textToEmbed.length === 0) {
|
||||
return new Array(384).fill(0)
|
||||
}
|
||||
} else {
|
||||
throw new Error('TransformerEmbedding only supports string or string[] data')
|
||||
}
|
||||
// Delegate to EmbeddingManager
|
||||
return embeddingManager.embed(data)
|
||||
}
|
||||
|
||||
// Ensure the extractor is available
|
||||
if (!this.extractor) {
|
||||
throw new Error('Transformer embedding model is not available')
|
||||
}
|
||||
|
||||
// Generate embeddings with mean pooling and normalization
|
||||
const result = await this.extractor(textToEmbed, {
|
||||
pooling: 'mean',
|
||||
normalize: true
|
||||
})
|
||||
|
||||
// Extract the embedding data
|
||||
let embedding: number[]
|
||||
|
||||
if (textToEmbed.length === 1) {
|
||||
// Single text input - return first embedding
|
||||
embedding = Array.from(result.data.slice(0, 384))
|
||||
} else {
|
||||
// Multiple texts - return first embedding (maintain compatibility)
|
||||
embedding = Array.from(result.data.slice(0, 384))
|
||||
}
|
||||
|
||||
// Validate embedding dimensions
|
||||
if (embedding.length !== 384) {
|
||||
this.logger('warn', `Unexpected embedding dimension: ${embedding.length}, expected 384`)
|
||||
// Pad or truncate to 384 dimensions
|
||||
if (embedding.length < 384) {
|
||||
embedding = [...embedding, ...new Array(384 - embedding.length).fill(0)]
|
||||
} else {
|
||||
embedding = embedding.slice(0, 384)
|
||||
}
|
||||
}
|
||||
|
||||
return embedding
|
||||
} catch (error) {
|
||||
this.logger('error', 'Error generating embeddings:', error)
|
||||
throw new Error(`Failed to generate embeddings: ${error}`)
|
||||
/**
|
||||
* Get the embedding function
|
||||
*/
|
||||
getEmbeddingFunction(): EmbeddingFunction {
|
||||
return async (data: string | string[] | Record<string, unknown>): Promise<Vector> => {
|
||||
return this.embed(data as string | string[])
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Dispose of the model and free resources
|
||||
* Check if initialized
|
||||
*/
|
||||
public async dispose(): Promise<void> {
|
||||
if (this.extractor && typeof this.extractor.dispose === 'function') {
|
||||
await this.extractor.dispose()
|
||||
}
|
||||
this.extractor = null
|
||||
this.initialized = false
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the dimension of embeddings produced by this model
|
||||
*/
|
||||
public getDimension(): number {
|
||||
return 384
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if the model is initialized
|
||||
*/
|
||||
public isInitialized(): boolean {
|
||||
isInitialized(): boolean {
|
||||
return this.initialized
|
||||
}
|
||||
|
||||
/**
|
||||
* Dispose resources (no-op for WASM engine)
|
||||
*/
|
||||
async dispose(): Promise<void> {
|
||||
this.initialized = false
|
||||
}
|
||||
}
|
||||
|
||||
// Legacy alias for backward compatibility
|
||||
export const UniversalSentenceEncoder = TransformerEmbedding
|
||||
/**
|
||||
* Create a simple embedding function using the default TransformerEmbedding
|
||||
* This is the recommended way to create an embedding function for Brainy
|
||||
*/
|
||||
export function createEmbeddingFunction(options: TransformerEmbeddingOptions = {}): EmbeddingFunction {
|
||||
return embeddingManager.getEmbeddingFunction()
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a new embedding model instance
|
||||
* Create a TransformerEmbedding instance (backward compatibility)
|
||||
*/
|
||||
export function createEmbeddingModel(options?: TransformerEmbeddingOptions): EmbeddingModel {
|
||||
export function createTransformerEmbedding(options: TransformerEmbeddingOptions = {}): TransformerEmbedding {
|
||||
return new TransformerEmbedding(options)
|
||||
}
|
||||
|
||||
/**
|
||||
* Default embedding function using the unified EmbeddingManager
|
||||
* Simple, clean, reliable - no more layers of indirection
|
||||
* Convenience function to detect best device (always returns 'wasm')
|
||||
*/
|
||||
export const defaultEmbeddingFunction: EmbeddingFunction = async (data: string | string[] | Record<string, unknown>): Promise<Vector> => {
|
||||
const { embed } = await import('../embeddings/EmbeddingManager.js')
|
||||
return await embed(data)
|
||||
export async function detectBestDevice(): Promise<'cpu' | 'webgpu' | 'cuda' | 'wasm'> {
|
||||
return 'wasm' as any
|
||||
}
|
||||
|
||||
/**
|
||||
* Create an embedding function with custom options
|
||||
* NOTE: Options are validated but the singleton EmbeddingManager is always used
|
||||
* Resolve device string (always returns 'wasm')
|
||||
*/
|
||||
export function createEmbeddingFunction(options: TransformerEmbeddingOptions = {}): EmbeddingFunction {
|
||||
return async (data: string | string[] | Record<string, unknown>): Promise<Vector> => {
|
||||
const { embeddingManager } = await import('../embeddings/EmbeddingManager.js')
|
||||
export async function resolveDevice(_device: string = 'auto'): Promise<string> {
|
||||
return 'wasm'
|
||||
}
|
||||
|
||||
// Validate precision if specified
|
||||
// Precision is always Q8 now
|
||||
|
||||
return await embeddingManager.embed(data)
|
||||
/**
|
||||
* Default embedding function (backward compatibility)
|
||||
*/
|
||||
export const defaultEmbeddingFunction: EmbeddingFunction = embeddingManager.getEmbeddingFunction()
|
||||
|
||||
/**
|
||||
* UniversalSentenceEncoder alias (backward compatibility)
|
||||
*/
|
||||
export const UniversalSentenceEncoder = TransformerEmbedding
|
||||
|
||||
/**
|
||||
* Batch embed function (backward compatibility)
|
||||
*/
|
||||
export async function batchEmbed(texts: string[]): Promise<Vector[]> {
|
||||
const results: Vector[] = []
|
||||
for (const text of texts) {
|
||||
results.push(await embeddingManager.embed(text))
|
||||
}
|
||||
return results
|
||||
}
|
||||
|
||||
/**
|
||||
* Batch embedding function for processing multiple texts efficiently
|
||||
*/
|
||||
export async function batchEmbed(
|
||||
texts: string[],
|
||||
options: TransformerEmbeddingOptions = {}
|
||||
): Promise<Vector[]> {
|
||||
const embedder = new TransformerEmbedding(options)
|
||||
await embedder.init()
|
||||
|
||||
const embeddings: Vector[] = []
|
||||
|
||||
// Process in batches for memory efficiency
|
||||
const batchSize = 32
|
||||
for (let i = 0; i < texts.length; i += batchSize) {
|
||||
const batch = texts.slice(i, i + batchSize)
|
||||
|
||||
for (const text of batch) {
|
||||
const embedding = await embedder.embed(text)
|
||||
embeddings.push(embedding)
|
||||
}
|
||||
}
|
||||
|
||||
await embedder.dispose()
|
||||
return embeddings
|
||||
}
|
||||
|
||||
/**
|
||||
* Embedding functions for specific model types
|
||||
* Embedding functions registry (backward compatibility)
|
||||
*/
|
||||
export const embeddingFunctions = {
|
||||
/** Default lightweight model (all-MiniLM-L6-v2, 384 dimensions) */
|
||||
default: defaultEmbeddingFunction,
|
||||
|
||||
/** Create custom embedding function */
|
||||
create: createEmbeddingFunction,
|
||||
|
||||
/** Batch processing */
|
||||
batch: batchEmbed
|
||||
}
|
||||
transformer: createEmbeddingFunction,
|
||||
default: createEmbeddingFunction,
|
||||
}
|
||||
|
|
|
|||
165
tests/integration/bun-compile-test.ts
Normal file
165
tests/integration/bun-compile-test.ts
Normal file
|
|
@ -0,0 +1,165 @@
|
|||
/**
|
||||
* Bun Compile Test
|
||||
*
|
||||
* Tests that Brainy works when compiled with `bun build --compile`.
|
||||
* This verifies:
|
||||
* 1. No native binaries required (pure WASM)
|
||||
* 2. Model is properly bundled
|
||||
* 3. Embeddings work without network access
|
||||
*
|
||||
* To test manually:
|
||||
* bun build tests/integration/bun-compile-test.ts --compile --outfile /tmp/brainy-bun-test
|
||||
* /tmp/brainy-bun-test
|
||||
*/
|
||||
|
||||
import { Brainy } from '../../dist/index.js'
|
||||
import { embeddingManager } from '../../dist/embeddings/EmbeddingManager.js'
|
||||
import { WASMEmbeddingEngine } from '../../dist/embeddings/wasm/index.js'
|
||||
|
||||
async function testBunCompile() {
|
||||
const results: { test: string; passed: boolean; error?: string }[] = []
|
||||
|
||||
console.log('🚀 Brainy Bun Compile Test\n')
|
||||
console.log('Runtime:', typeof Bun !== 'undefined' ? `Bun ${Bun.version}` : `Node.js ${process.version}`)
|
||||
console.log('')
|
||||
|
||||
// Test 1: WASM Engine initialization
|
||||
try {
|
||||
console.log('1. Testing WASM Engine initialization...')
|
||||
const engine = WASMEmbeddingEngine.getInstance()
|
||||
await engine.initialize()
|
||||
console.log(' ✅ WASM Engine initialized')
|
||||
results.push({ test: 'WASM Engine init', passed: true })
|
||||
} catch (error) {
|
||||
console.log(' ❌ WASM Engine failed:', (error as Error).message)
|
||||
results.push({ test: 'WASM Engine init', passed: false, error: (error as Error).message })
|
||||
}
|
||||
|
||||
// Test 2: Generate embedding
|
||||
try {
|
||||
console.log('2. Testing embedding generation...')
|
||||
const engine = WASMEmbeddingEngine.getInstance()
|
||||
const embedding = await engine.embed('Hello from Bun!')
|
||||
if (embedding.length !== 384) {
|
||||
throw new Error(`Expected 384 dimensions, got ${embedding.length}`)
|
||||
}
|
||||
console.log(` ✅ Generated ${embedding.length}-dim embedding`)
|
||||
results.push({ test: 'Embedding generation', passed: true })
|
||||
} catch (error) {
|
||||
console.log(' ❌ Embedding failed:', (error as Error).message)
|
||||
results.push({ test: 'Embedding generation', passed: false, error: (error as Error).message })
|
||||
}
|
||||
|
||||
// Test 3: Semantic similarity
|
||||
try {
|
||||
console.log('3. Testing semantic similarity...')
|
||||
const engine = WASMEmbeddingEngine.getInstance()
|
||||
const catEmb = await engine.embed('The cat sat on the mat')
|
||||
const felineEmb = await engine.embed('A feline rests on a rug')
|
||||
const stockEmb = await engine.embed('Stock market crash')
|
||||
|
||||
const cosineSim = (a: number[], b: number[]) => {
|
||||
let dot = 0
|
||||
for (let i = 0; i < a.length; i++) dot += a[i] * b[i]
|
||||
return dot // Already normalized
|
||||
}
|
||||
|
||||
const similarSim = cosineSim(catEmb, felineEmb)
|
||||
const dissimilarSim = cosineSim(catEmb, stockEmb)
|
||||
|
||||
if (similarSim <= dissimilarSim) {
|
||||
throw new Error(`Semantic similarity failed: similar=${similarSim.toFixed(4)}, dissimilar=${dissimilarSim.toFixed(4)}`)
|
||||
}
|
||||
console.log(` ✅ Semantic similarity works (similar: ${similarSim.toFixed(4)} > dissimilar: ${dissimilarSim.toFixed(4)})`)
|
||||
results.push({ test: 'Semantic similarity', passed: true })
|
||||
} catch (error) {
|
||||
console.log(' ❌ Semantic similarity failed:', (error as Error).message)
|
||||
results.push({ test: 'Semantic similarity', passed: false, error: (error as Error).message })
|
||||
}
|
||||
|
||||
// Test 4: EmbeddingManager
|
||||
try {
|
||||
console.log('4. Testing EmbeddingManager...')
|
||||
await embeddingManager.init()
|
||||
const emb = await embeddingManager.embed('Test via manager')
|
||||
if (emb.length !== 384) {
|
||||
throw new Error(`Expected 384 dimensions, got ${emb.length}`)
|
||||
}
|
||||
console.log(' ✅ EmbeddingManager works')
|
||||
results.push({ test: 'EmbeddingManager', passed: true })
|
||||
} catch (error) {
|
||||
console.log(' ❌ EmbeddingManager failed:', (error as Error).message)
|
||||
results.push({ test: 'EmbeddingManager', passed: false, error: (error as Error).message })
|
||||
}
|
||||
|
||||
// Test 5: Full Brainy initialization with fresh memory storage
|
||||
try {
|
||||
console.log('5. Testing Brainy initialization (in-memory)...')
|
||||
const brain = new Brainy({
|
||||
storage: 'memory',
|
||||
storageOptions: { path: ':memory:' }
|
||||
})
|
||||
await brain.init()
|
||||
console.log(' ✅ Brainy initialized')
|
||||
results.push({ test: 'Brainy init', passed: true })
|
||||
|
||||
// Test 6: Add document
|
||||
console.log('6. Testing document add...')
|
||||
const docId = await brain.add({ data: 'Machine learning concepts', type: 'concept' })
|
||||
if (!docId || typeof docId !== 'string') {
|
||||
throw new Error('Document add returned no ID')
|
||||
}
|
||||
console.log(` ✅ Document added: ${docId}`)
|
||||
results.push({ test: 'Document add', passed: true })
|
||||
|
||||
// Test 7: Search
|
||||
console.log('7. Testing semantic search...')
|
||||
const searchResults = await brain.find('AI')
|
||||
console.log(` ✅ Search returned ${searchResults.length} results`)
|
||||
results.push({ test: 'Semantic search', passed: true })
|
||||
|
||||
// Test 8: Get document
|
||||
console.log('8. Testing document retrieval...')
|
||||
const retrieved = await brain.get(docId)
|
||||
if (!retrieved) {
|
||||
throw new Error('Document not found')
|
||||
}
|
||||
console.log(' ✅ Document retrieved')
|
||||
results.push({ test: 'Document retrieval', passed: true })
|
||||
|
||||
await brain.close()
|
||||
} catch (error) {
|
||||
console.log(' ❌ Brainy test failed:', (error as Error).message)
|
||||
results.push({ test: 'Brainy operations', passed: false, error: (error as Error).message })
|
||||
}
|
||||
|
||||
// Summary
|
||||
console.log('\n' + '='.repeat(50))
|
||||
console.log('SUMMARY')
|
||||
console.log('='.repeat(50))
|
||||
|
||||
const passed = results.filter(r => r.passed).length
|
||||
const failed = results.filter(r => !r.passed).length
|
||||
|
||||
for (const r of results) {
|
||||
console.log(`${r.passed ? '✅' : '❌'} ${r.test}${r.error ? `: ${r.error}` : ''}`)
|
||||
}
|
||||
|
||||
console.log('')
|
||||
console.log(`Passed: ${passed}/${results.length}`)
|
||||
console.log(`Failed: ${failed}/${results.length}`)
|
||||
|
||||
if (failed > 0) {
|
||||
console.log('\n❌ Some tests failed!')
|
||||
process.exit(1)
|
||||
} else {
|
||||
console.log('\n✅ All tests passed! Brainy works with Bun compile.')
|
||||
process.exit(0)
|
||||
}
|
||||
}
|
||||
|
||||
// Run tests
|
||||
testBunCompile().catch(error => {
|
||||
console.error('Fatal error:', error)
|
||||
process.exit(1)
|
||||
})
|
||||
142
tests/integration/wasm-embeddings.test.ts
Normal file
142
tests/integration/wasm-embeddings.test.ts
Normal file
|
|
@ -0,0 +1,142 @@
|
|||
/**
|
||||
* WASM Embedding Integration Test
|
||||
*
|
||||
* Tests the actual WASM embedding engine with real model inference.
|
||||
* NO mocks - this loads the real ONNX model and generates real embeddings.
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeAll } from 'vitest'
|
||||
import { WASMEmbeddingEngine } from '../../src/embeddings/wasm/WASMEmbeddingEngine.js'
|
||||
import { embeddingManager } from '../../src/embeddings/EmbeddingManager.js'
|
||||
|
||||
// Ensure we're NOT in mock mode for these tests
|
||||
beforeAll(() => {
|
||||
delete process.env.BRAINY_UNIT_TEST
|
||||
;(globalThis as any).__BRAINY_UNIT_TEST__ = false
|
||||
})
|
||||
|
||||
describe('WASM Embedding Engine - Real Embeddings', () => {
|
||||
it('should initialize the WASM engine', async () => {
|
||||
const engine = WASMEmbeddingEngine.getInstance()
|
||||
await engine.initialize()
|
||||
expect(engine.isInitialized()).toBe(true)
|
||||
})
|
||||
|
||||
it('should generate 384-dimensional embeddings', async () => {
|
||||
const engine = WASMEmbeddingEngine.getInstance()
|
||||
const embedding = await engine.embed('Hello world')
|
||||
|
||||
expect(embedding).toBeInstanceOf(Array)
|
||||
expect(embedding.length).toBe(384)
|
||||
expect(typeof embedding[0]).toBe('number')
|
||||
})
|
||||
|
||||
it('should produce consistent embeddings for same input', async () => {
|
||||
const engine = WASMEmbeddingEngine.getInstance()
|
||||
|
||||
const emb1 = await engine.embed('test consistency')
|
||||
const emb2 = await engine.embed('test consistency')
|
||||
|
||||
expect(emb1).toEqual(emb2)
|
||||
})
|
||||
|
||||
it('should produce normalized embeddings (unit length)', async () => {
|
||||
const engine = WASMEmbeddingEngine.getInstance()
|
||||
const embedding = await engine.embed('normalize test')
|
||||
|
||||
// Calculate L2 norm
|
||||
const norm = Math.sqrt(embedding.reduce((sum, v) => sum + v * v, 0))
|
||||
|
||||
// Should be approximately 1.0 (normalized)
|
||||
expect(norm).toBeCloseTo(1.0, 4)
|
||||
})
|
||||
|
||||
it('should produce semantically meaningful embeddings', async () => {
|
||||
const engine = WASMEmbeddingEngine.getInstance()
|
||||
|
||||
// Similar sentences
|
||||
const catEmb = await engine.embed('The cat sat on the mat')
|
||||
const felineEmb = await engine.embed('A feline rests on a rug')
|
||||
|
||||
// Dissimilar sentence
|
||||
const stockEmb = await engine.embed('The stock market crashed today')
|
||||
|
||||
// Cosine similarity function
|
||||
const cosineSim = (a: number[], b: number[]) => {
|
||||
let dot = 0, normA = 0, normB = 0
|
||||
for (let i = 0; i < a.length; i++) {
|
||||
dot += a[i] * b[i]
|
||||
normA += a[i] * a[i]
|
||||
normB += b[i] * b[i]
|
||||
}
|
||||
return dot / (Math.sqrt(normA) * Math.sqrt(normB))
|
||||
}
|
||||
|
||||
const similarSim = cosineSim(catEmb, felineEmb)
|
||||
const dissimilarSim = cosineSim(catEmb, stockEmb)
|
||||
|
||||
// Similar sentences should have higher similarity than dissimilar
|
||||
expect(similarSim).toBeGreaterThan(dissimilarSim)
|
||||
expect(similarSim).toBeGreaterThan(0.4) // Reasonably similar
|
||||
expect(dissimilarSim).toBeLessThan(0.3) // Not very similar
|
||||
})
|
||||
|
||||
it('should handle empty string', async () => {
|
||||
const engine = WASMEmbeddingEngine.getInstance()
|
||||
const embedding = await engine.embed('')
|
||||
|
||||
expect(embedding.length).toBe(384)
|
||||
})
|
||||
|
||||
it('should handle long text', async () => {
|
||||
const engine = WASMEmbeddingEngine.getInstance()
|
||||
const longText = 'word '.repeat(1000)
|
||||
const embedding = await engine.embed(longText)
|
||||
|
||||
expect(embedding.length).toBe(384)
|
||||
})
|
||||
|
||||
it('should work through EmbeddingManager', async () => {
|
||||
await embeddingManager.init()
|
||||
|
||||
const embedding = await embeddingManager.embed('test via manager')
|
||||
|
||||
expect(embedding.length).toBe(384)
|
||||
expect(embeddingManager.isInitialized()).toBe(true)
|
||||
})
|
||||
|
||||
it('should report correct stats', async () => {
|
||||
const engine = WASMEmbeddingEngine.getInstance()
|
||||
const stats = engine.getStats()
|
||||
|
||||
expect(stats.initialized).toBe(true)
|
||||
expect(stats.modelName).toBe('all-MiniLM-L6-v2')
|
||||
expect(stats.embedCount).toBeGreaterThan(0)
|
||||
})
|
||||
})
|
||||
|
||||
describe('WASM Embedding - Batch Operations', () => {
|
||||
it('should embed multiple texts', async () => {
|
||||
const engine = WASMEmbeddingEngine.getInstance()
|
||||
|
||||
const texts = [
|
||||
'First document about cats',
|
||||
'Second document about dogs',
|
||||
'Third document about birds'
|
||||
]
|
||||
|
||||
const embeddings = await engine.embedBatch(texts)
|
||||
|
||||
expect(embeddings.length).toBe(3)
|
||||
expect(embeddings[0].length).toBe(384)
|
||||
expect(embeddings[1].length).toBe(384)
|
||||
expect(embeddings[2].length).toBe(384)
|
||||
})
|
||||
|
||||
it('should handle empty batch', async () => {
|
||||
const engine = WASMEmbeddingEngine.getInstance()
|
||||
const embeddings = await engine.embedBatch([])
|
||||
|
||||
expect(embeddings).toEqual([])
|
||||
})
|
||||
})
|
||||
Loading…
Add table
Add a link
Reference in a new issue