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