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