feat: add optional Q8 quantized model support
- Add Q8 quantized models (75% smaller than FP32) - Enhance download scripts with model variant selection - Add smart model loading with availability detection - Implement runtime warnings for Q8 compatibility - Update documentation with Q8 usage examples - Maintain 100% backward compatibility (FP32 default) BREAKING CHANGE: None - FP32 remains default 🧠 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
parent
223311f4e7
commit
32df3ee6ae
8 changed files with 259 additions and 40 deletions
|
|
@ -9,6 +9,11 @@ const path = require('path')
|
|||
const MODEL_NAME = 'Xenova/all-MiniLM-L6-v2'
|
||||
const OUTPUT_DIR = './models'
|
||||
|
||||
// Parse command line arguments for model type selection
|
||||
const args = process.argv.slice(2)
|
||||
const downloadType = args.includes('fp32') ? 'fp32' :
|
||||
args.includes('q8') ? 'q8' : 'both'
|
||||
|
||||
async function downloadModels() {
|
||||
// Use dynamic import for ES modules in CommonJS
|
||||
const { pipeline, env } = await import('@huggingface/transformers')
|
||||
|
|
@ -16,29 +21,31 @@ async function downloadModels() {
|
|||
// Configure transformers.js to use local cache
|
||||
env.cacheDir = './models-cache'
|
||||
env.allowRemoteModels = true
|
||||
|
||||
try {
|
||||
console.log('🔄 Downloading all-MiniLM-L6-v2 model for offline bundling...')
|
||||
console.log('🧠 Brainy Model Downloader v2.8.0')
|
||||
console.log('===================================')
|
||||
console.log(` Model: ${MODEL_NAME}`)
|
||||
console.log(` Type: ${downloadType} (fp32, q8, or both)`)
|
||||
console.log(` Cache: ${env.cacheDir}`)
|
||||
console.log('')
|
||||
|
||||
// Create output directory
|
||||
await fs.mkdir(OUTPUT_DIR, { recursive: true })
|
||||
|
||||
// Load the model to force download
|
||||
console.log('📥 Loading model pipeline...')
|
||||
const extractor = await pipeline('feature-extraction', MODEL_NAME)
|
||||
// Download models based on type
|
||||
if (downloadType === 'both' || downloadType === 'fp32') {
|
||||
console.log('📥 Downloading FP32 model (full precision, 90MB)...')
|
||||
await downloadModelVariant('fp32')
|
||||
}
|
||||
|
||||
// Test the model to make sure it works
|
||||
console.log('🧪 Testing model...')
|
||||
const testResult = await extractor(['Hello world!'], {
|
||||
pooling: 'mean',
|
||||
normalize: true
|
||||
})
|
||||
|
||||
console.log(`✅ Model test successful! Embedding dimensions: ${testResult.data.length}`)
|
||||
if (downloadType === 'both' || downloadType === 'q8') {
|
||||
console.log('📥 Downloading Q8 model (quantized, 23MB)...')
|
||||
await downloadModelVariant('q8')
|
||||
}
|
||||
|
||||
// Copy ALL model files from cache to our models directory
|
||||
console.log('📋 Copying ALL model files to bundle directory...')
|
||||
console.log('📋 Copying model files to bundle directory...')
|
||||
|
||||
const cacheDir = path.resolve(env.cacheDir)
|
||||
const outputDir = path.resolve(OUTPUT_DIR)
|
||||
|
|
@ -62,22 +69,89 @@ async function downloadModels() {
|
|||
console.log(` Total size: ${await calculateDirectorySize(outputDir)} MB`)
|
||||
console.log(` Location: ${outputDir}`)
|
||||
|
||||
// Create a marker file
|
||||
// Create a marker file with downloaded model info
|
||||
const markerData = {
|
||||
model: MODEL_NAME,
|
||||
bundledAt: new Date().toISOString(),
|
||||
version: '2.8.0',
|
||||
downloadType: downloadType,
|
||||
models: {}
|
||||
}
|
||||
|
||||
// Check which models were downloaded
|
||||
const fp32Path = path.join(outputDir, 'Xenova/all-MiniLM-L6-v2/onnx/model.onnx')
|
||||
const q8Path = path.join(outputDir, 'Xenova/all-MiniLM-L6-v2/onnx/model_quantized.onnx')
|
||||
|
||||
if (await fileExists(fp32Path)) {
|
||||
const stats = await fs.stat(fp32Path)
|
||||
markerData.models.fp32 = {
|
||||
file: 'onnx/model.onnx',
|
||||
size: stats.size,
|
||||
sizeFormatted: `${Math.round(stats.size / (1024 * 1024))}MB`
|
||||
}
|
||||
}
|
||||
|
||||
if (await fileExists(q8Path)) {
|
||||
const stats = await fs.stat(q8Path)
|
||||
markerData.models.q8 = {
|
||||
file: 'onnx/model_quantized.onnx',
|
||||
size: stats.size,
|
||||
sizeFormatted: `${Math.round(stats.size / (1024 * 1024))}MB`
|
||||
}
|
||||
}
|
||||
|
||||
await fs.writeFile(
|
||||
path.join(outputDir, '.brainy-models-bundled'),
|
||||
JSON.stringify({
|
||||
model: MODEL_NAME,
|
||||
bundledAt: new Date().toISOString(),
|
||||
version: '1.0.0'
|
||||
}, null, 2)
|
||||
JSON.stringify(markerData, null, 2)
|
||||
)
|
||||
|
||||
console.log('')
|
||||
console.log('✅ Download complete! Available models:')
|
||||
if (markerData.models.fp32) {
|
||||
console.log(` • FP32: ${markerData.models.fp32.sizeFormatted} (full precision)`)
|
||||
}
|
||||
if (markerData.models.q8) {
|
||||
console.log(` • Q8: ${markerData.models.q8.sizeFormatted} (quantized, 75% smaller)`)
|
||||
}
|
||||
console.log('')
|
||||
console.log('Air-gap deployment ready! 🚀')
|
||||
|
||||
} catch (error) {
|
||||
console.error('❌ Error downloading models:', error)
|
||||
process.exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
// Download a specific model variant
|
||||
async function downloadModelVariant(dtype) {
|
||||
const { pipeline } = await import('@huggingface/transformers')
|
||||
|
||||
try {
|
||||
// Load the model to force download
|
||||
const extractor = await pipeline('feature-extraction', MODEL_NAME, {
|
||||
dtype: dtype,
|
||||
cache_dir: './models-cache'
|
||||
})
|
||||
|
||||
// Test the model
|
||||
const testResult = await extractor(['Hello world!'], {
|
||||
pooling: 'mean',
|
||||
normalize: true
|
||||
})
|
||||
|
||||
console.log(` ✅ ${dtype.toUpperCase()} model downloaded and tested (${testResult.data.length} dimensions)`)
|
||||
|
||||
// Dispose to free memory
|
||||
if (extractor.dispose) {
|
||||
await extractor.dispose()
|
||||
}
|
||||
|
||||
} catch (error) {
|
||||
console.error(` ❌ Failed to download ${dtype} model:`, error)
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
async function findModelDirectories(baseDir, modelName) {
|
||||
const dirs = []
|
||||
|
||||
|
|
@ -141,6 +215,15 @@ async function dirExists(dir) {
|
|||
}
|
||||
}
|
||||
|
||||
async function fileExists(file) {
|
||||
try {
|
||||
const stats = await fs.stat(file)
|
||||
return stats.isFile()
|
||||
} catch (error) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
async function copyDirectory(src, dest) {
|
||||
await fs.mkdir(dest, { recursive: true })
|
||||
const entries = await fs.readdir(src, { withFileTypes: true })
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue