perf: 580x faster embedding init - separate model from WASM
Cloud Run cold starts taking 139 seconds due to 90MB WASM file with embedded 87MB model weights. WASM compilation scales with file size. Solution: Split into 2.4MB WASM (code only) + external model files. - WASM compile: 139,000ms → 6-8ms - Model load: N/A → 30-115ms - Total init: 139,000ms → 136-240ms New modelLoader.ts handles all environments: - Node.js: fs.readFile() - Bun: Bun.file() - Bun --compile: auto-embedded assets - Browser: fetch() Zero config - same API, npm package includes model files. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
parent
4fc1fb7340
commit
677e2d6624
7 changed files with 360 additions and 38 deletions
188
src/embeddings/wasm/modelLoader.ts
Normal file
188
src/embeddings/wasm/modelLoader.ts
Normal file
|
|
@ -0,0 +1,188 @@
|
|||
/**
|
||||
* Universal Model Loader for Candle Embeddings
|
||||
*
|
||||
* v7.2.0: Model weights are now loaded separately from WASM for faster initialization.
|
||||
* This reduces WASM compilation time from 139 seconds to ~3-5 seconds.
|
||||
*
|
||||
* Loads model files from appropriate source based on environment:
|
||||
*
|
||||
* | Environment | Method |
|
||||
* |----------------|-------------------------------------|
|
||||
* | Node.js | fs.readFile() |
|
||||
* | Bun | Bun.file() |
|
||||
* | Bun --compile | Bun.file() with auto-embedded asset |
|
||||
* | Browser | fetch() |
|
||||
*
|
||||
* For Bun --compile, Bun automatically embeds files accessed via Bun.file()
|
||||
* when the paths are resolvable at build time.
|
||||
*/
|
||||
|
||||
// =============================================================================
|
||||
// Type Declarations
|
||||
// =============================================================================
|
||||
|
||||
declare const Bun: {
|
||||
file(path: string): { arrayBuffer(): Promise<ArrayBuffer> }
|
||||
} | undefined
|
||||
|
||||
// =============================================================================
|
||||
// Environment Detection (evaluated once at module load)
|
||||
// =============================================================================
|
||||
|
||||
const isBun = typeof Bun !== 'undefined'
|
||||
const isNode = !isBun && typeof process !== 'undefined' && !!process.versions?.node
|
||||
|
||||
// =============================================================================
|
||||
// Types
|
||||
// =============================================================================
|
||||
|
||||
export interface ModelAssets {
|
||||
model: Uint8Array
|
||||
tokenizer: Uint8Array
|
||||
config: Uint8Array
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Public API
|
||||
// =============================================================================
|
||||
|
||||
/**
|
||||
* Load model assets from the appropriate source for the current environment.
|
||||
*
|
||||
* This function handles all environment differences internally.
|
||||
* Model files are loaded in parallel for best performance.
|
||||
*
|
||||
* @returns ModelAssets containing model, tokenizer, and config bytes
|
||||
* @throws Error if model files cannot be loaded
|
||||
*/
|
||||
export async function loadModelAssets(): Promise<ModelAssets> {
|
||||
// Bun (runtime or compiled binary)
|
||||
if (isBun) {
|
||||
return loadBunAssets()
|
||||
}
|
||||
|
||||
// Node.js
|
||||
if (isNode) {
|
||||
return loadNodeAssets()
|
||||
}
|
||||
|
||||
// Browser
|
||||
return loadBrowserAssets()
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if model files are available.
|
||||
* Useful for debugging.
|
||||
*/
|
||||
export function isModelEmbedded(): boolean {
|
||||
// In Bun --compile, files accessed via Bun.file() are auto-embedded
|
||||
return isBun
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Internal Helpers - Bun
|
||||
// =============================================================================
|
||||
|
||||
async function loadBunAssets(): Promise<ModelAssets> {
|
||||
const modelPath = await resolveAssetPath('model.safetensors')
|
||||
const tokenizerPath = await resolveAssetPath('tokenizer.json')
|
||||
const configPath = await resolveAssetPath('config.json')
|
||||
|
||||
const [model, tokenizer, config] = await Promise.all([
|
||||
Bun!.file(modelPath).arrayBuffer(),
|
||||
Bun!.file(tokenizerPath).arrayBuffer(),
|
||||
Bun!.file(configPath).arrayBuffer(),
|
||||
])
|
||||
|
||||
return {
|
||||
model: new Uint8Array(model),
|
||||
tokenizer: new Uint8Array(tokenizer),
|
||||
config: new Uint8Array(config),
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Internal Helpers - Node.js
|
||||
// =============================================================================
|
||||
|
||||
async function loadNodeAssets(): Promise<ModelAssets> {
|
||||
const fs = await import('node:fs')
|
||||
const nodePath = await import('node:path')
|
||||
const { fileURLToPath } = await import('node:url')
|
||||
|
||||
const thisDir = nodePath.dirname(fileURLToPath(import.meta.url))
|
||||
const assetsDir = nodePath.resolve(thisDir, '../../../assets/models/all-MiniLM-L6-v2')
|
||||
|
||||
// Verify assets directory exists
|
||||
if (!fs.existsSync(assetsDir)) {
|
||||
throw new Error(
|
||||
`Model assets not found: ${assetsDir}\n` +
|
||||
`Ensure @soulcraft/brainy is installed correctly.`
|
||||
)
|
||||
}
|
||||
|
||||
const [model, tokenizer, config] = await Promise.all([
|
||||
fs.promises.readFile(nodePath.join(assetsDir, 'model.safetensors')),
|
||||
fs.promises.readFile(nodePath.join(assetsDir, 'tokenizer.json')),
|
||||
fs.promises.readFile(nodePath.join(assetsDir, 'config.json')),
|
||||
])
|
||||
|
||||
return {
|
||||
model: new Uint8Array(model),
|
||||
tokenizer: new Uint8Array(tokenizer),
|
||||
config: new Uint8Array(config),
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Internal Helpers - Browser
|
||||
// =============================================================================
|
||||
|
||||
async function loadBrowserAssets(): Promise<ModelAssets> {
|
||||
// In browser, assets are served relative to the WASM location
|
||||
const baseUrl = new URL('../../../assets/models/all-MiniLM-L6-v2/', import.meta.url)
|
||||
|
||||
const [modelRes, tokenizerRes, configRes] = await Promise.all([
|
||||
fetch(new URL('model.safetensors', baseUrl)),
|
||||
fetch(new URL('tokenizer.json', baseUrl)),
|
||||
fetch(new URL('config.json', baseUrl)),
|
||||
])
|
||||
|
||||
// Check for errors
|
||||
if (!modelRes.ok) {
|
||||
throw new Error(`Failed to fetch model: ${modelRes.status} ${modelRes.statusText}`)
|
||||
}
|
||||
if (!tokenizerRes.ok) {
|
||||
throw new Error(`Failed to fetch tokenizer: ${tokenizerRes.status} ${tokenizerRes.statusText}`)
|
||||
}
|
||||
if (!configRes.ok) {
|
||||
throw new Error(`Failed to fetch config: ${configRes.status} ${configRes.statusText}`)
|
||||
}
|
||||
|
||||
const [model, tokenizer, config] = await Promise.all([
|
||||
modelRes.arrayBuffer(),
|
||||
tokenizerRes.arrayBuffer(),
|
||||
configRes.arrayBuffer(),
|
||||
])
|
||||
|
||||
return {
|
||||
model: new Uint8Array(model),
|
||||
tokenizer: new Uint8Array(tokenizer),
|
||||
config: new Uint8Array(config),
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Internal Helpers - Path Resolution
|
||||
// =============================================================================
|
||||
|
||||
/**
|
||||
* Resolve the filesystem path to a model asset file.
|
||||
*/
|
||||
async function resolveAssetPath(filename: string): Promise<string> {
|
||||
const nodePath = await import('node:path')
|
||||
const { fileURLToPath } = await import('node:url')
|
||||
|
||||
const thisDir = nodePath.dirname(fileURLToPath(import.meta.url))
|
||||
return nodePath.join(thisDir, '../../../assets/models/all-MiniLM-L6-v2', filename)
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue