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
|
|
@ -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
|
||||
Loading…
Add table
Add a link
Reference in a new issue