**feat(models): add pre-bundled Universal Sentence Encoder for offline use**

- Introduced `@soulcraft/brainy-models` package with pre-bundled TensorFlow models for enhanced offline reliability.
- Added `index.d.ts` and `index.js` allowing offline embedding workflows with the Universal Sentence Encoder model.
- Included utility scripts for model compression, size retrieval, and availability checks.
- Added `metadata.json` and `model.json` defining the Universal Sentence Encoder configuration with offline bundling.
- Ensured comprehensive model documentation, error handling, and robust logging for seamless integration.
- Supported optional model quantization placeholders for future TensorFlow.js enhancements.

**Purpose**: Enable fully offline-ready embedding workflows via pre-bundled Universal Sentence Encoder models, ensuring maximum reliability and air-gapped environment compatibility.
This commit is contained in:
David Snelling 2025-08-01 16:22:58 -07:00
parent 93483572d8
commit e476d45fac
16 changed files with 7678 additions and 52 deletions

View file

@ -11,6 +11,19 @@ import { readFileSync, existsSync } from 'fs'
import { join, dirname } from 'path'
import { fileURLToPath } from 'url'
/**
* Helper function to safely extract error message from unknown error type
*/
function getErrorMessage(error: unknown): string {
if (error instanceof Error) {
return error.message
}
if (typeof error === 'string') {
return error
}
return String(error)
}
// Get the package directory
const __filename = fileURLToPath(import.meta.url)
const __dirname = dirname(__filename)
@ -90,7 +103,7 @@ export class BundledUniversalSentenceEncoder {
}
} catch (error) {
throw new Error(`Failed to load bundled model: ${error.message}`)
throw new Error(`Failed to load bundled model: ${getErrorMessage(error)}`)
}
}
@ -114,7 +127,7 @@ export class BundledUniversalSentenceEncoder {
return embeddings
} catch (error) {
throw new Error(`Failed to generate embeddings: ${error.message}`)
throw new Error(`Failed to generate embeddings: ${getErrorMessage(error)}`)
}
}
@ -173,6 +186,7 @@ export class BundledUniversalSentenceEncoder {
export class ModelCompressor {
/**
* Compress model weights using quantization
* Note: TensorFlow.js doesn't currently support model quantization
*/
static async quantizeModel(
modelPath: string,
@ -187,23 +201,23 @@ export class ModelCompressor {
console.log(`🗜️ Quantizing model to ${dtype}...`)
// Note: TensorFlow.js doesn't have built-in quantization yet,
// but we can implement basic weight compression
const modelArtifacts = await model.serialize()
// Save the compressed model
await tf.io.fileSystem(outputPath).save(modelArtifacts)
console.log(`✅ Compressed model saved to: ${outputPath}`)
// TensorFlow.js doesn't have built-in quantization or model serialization APIs yet
// This is a placeholder implementation that acknowledges the limitation
console.warn('⚠️ Model quantization is not yet supported in TensorFlow.js')
console.log(`📋 Model loaded successfully from: ${modelPath}`)
console.log(`📋 Target output path: ${outputPath}`)
console.log(`📋 Target dtype: ${dtype}`)
model.dispose()
throw new Error('Model quantization is not yet supported in TensorFlow.js. This feature requires server-side processing with TensorFlow Python.')
} catch (error) {
throw new Error(`Failed to compress model: ${error.message}`)
throw new Error(`Failed to compress model: ${getErrorMessage(error)}`)
}
}
/**
* Get model size information
* Get model size information by reading files from disk
*/
static async getModelSize(modelPath: string): Promise<{
totalSize: number
@ -211,22 +225,41 @@ export class ModelCompressor {
modelJsonSize: number
}> {
try {
// Load model to verify it's valid
const model = await tf.loadGraphModel(`file://${modelPath}`)
const artifacts = await model.serialize()
const weightsSize = artifacts.weightData?.byteLength || 0
const modelJsonSize = JSON.stringify(artifacts.modelTopology).length
const totalSize = weightsSize + modelJsonSize
model.dispose()
// Get model.json size
const modelJsonSize = existsSync(modelPath) ? readFileSync(modelPath).length : 0
// Calculate weights size by reading weight files
let weightsSize = 0
const modelDir = dirname(modelPath)
// Read model.json to get weight file names
if (existsSync(modelPath)) {
const modelJson = JSON.parse(readFileSync(modelPath, 'utf8'))
if (modelJson.weightsManifest) {
for (const manifest of modelJson.weightsManifest) {
for (const path of manifest.paths) {
const weightFilePath = join(modelDir, path)
if (existsSync(weightFilePath)) {
weightsSize += readFileSync(weightFilePath).length
}
}
}
}
}
const totalSize = weightsSize + modelJsonSize
return {
totalSize,
weightsSize,
modelJsonSize
}
} catch (error) {
throw new Error(`Failed to get model size: ${error.message}`)
throw new Error(`Failed to get model size: ${getErrorMessage(error)}`)
}
}
}