brainy/brainy-models-package/scripts/download-full-models.js

#!/usr/bin/env node

/* eslint-env node */
/* eslint-disable no-console */

/**
 * Download Full Models Script for @soulcraft/brainy-models
 * 
 * This script downloads the complete Universal Sentence Encoder model
 * and saves it locally for offline use, providing maximum reliability.
 */

import fs from 'fs'
import path from 'path'
import { fileURLToPath } from 'url'
import * as tf from '@tensorflow/tfjs-node'
import * as use from '@tensorflow-models/universal-sentence-encoder'
import https from 'https'
import { promisify } from 'util'

const __filename = fileURLToPath(import.meta.url)
const __dirname = path.dirname(__filename)

const MODELS_DIR = path.join(__dirname, '..', 'models')
const USE_MODEL_DIR = path.join(MODELS_DIR, 'universal-sentence-encoder')

// Ensure directories exist
if (!fs.existsSync(MODELS_DIR)) {
  fs.mkdirSync(MODELS_DIR, { recursive: true })
}

if (!fs.existsSync(USE_MODEL_DIR)) {
  fs.mkdirSync(USE_MODEL_DIR, { recursive: true })
}

console.log('🚀 Starting full model download for @soulcraft/brainy-models...')
console.log('This will download the complete Universal Sentence Encoder model (~25MB)')
console.log('for offline use and maximum reliability.\n')

/**
 * Download a file from URL to local path
 */
async function downloadFile(url, filePath) {
  return new Promise((resolve, reject) => {
    const file = fs.createWriteStream(filePath)
    
    https.get(url, (response) => {
      if (response.statusCode !== 200) {
        reject(new Error(`Failed to download ${url}: ${response.statusCode}`))
        return
      }
      
      const totalSize = parseInt(response.headers['content-length'] || '0')
      let downloadedSize = 0
      
      response.on('data', (chunk) => {
        downloadedSize += chunk.length
        if (totalSize > 0) {
          const progress = ((downloadedSize / totalSize) * 100).toFixed(1)
          process.stdout.write(`\r📥 Downloading: ${progress}% (${downloadedSize}/${totalSize} bytes)`)
        }
      })
      
      response.pipe(file)
      
      file.on('finish', () => {
        file.close()
        console.log(`\n✅ Downloaded: ${path.basename(filePath)}`)
        resolve()
      })
      
      file.on('error', (err) => {
        fs.unlink(filePath, () => {}) // Delete partial file
        reject(err)
      })
    }).on('error', reject)
  })
}

/**
 * Download the complete Universal Sentence Encoder model
 */
async function downloadFullModel() {
  try {
    console.log('🔍 Loading model to get download URLs...')
    
    // Load the model to get access to its internal structure
    const model = await use.load()
    console.log('✅ Model loaded successfully')
    
    // Test the model to ensure it works
    console.log('🧪 Testing model functionality...')
    const testEmbedding = await model.embed(['Hello world'])
    const testArray = await testEmbedding.array()
    console.log(`✅ Model test passed - embedding dimensions: ${testArray[0].length}`)
    testEmbedding.dispose()
    
    // The Universal Sentence Encoder model URL
    const modelBaseUrl = 'https://tfhub.dev/tensorflow/tfjs-model/universal-sentence-encoder/1/default/1'
    
    console.log('📦 Downloading model files...')
    
    // Download model.json
    const modelJsonUrl = `${modelBaseUrl}/model.json`
    const modelJsonPath = path.join(USE_MODEL_DIR, 'model.json')
    await downloadFile(modelJsonUrl, modelJsonPath)
    
    // Read the model.json to get the weights manifest
    const modelJson = JSON.parse(fs.readFileSync(modelJsonPath, 'utf8'))
    
    // Download all weight files
    if (modelJson.weightsManifest) {
      for (const manifest of modelJson.weightsManifest) {
        for (const weightFile of manifest.paths) {
          const weightUrl = `${modelBaseUrl}/${weightFile}`
          const weightPath = path.join(USE_MODEL_DIR, weightFile)
          await downloadFile(weightUrl, weightPath)
        }
      }
    }
    
    // Create metadata for the bundled model
    const metadata = {
      name: 'universal-sentence-encoder',
      version: '1.0.0',
      description: 'Complete Universal Sentence Encoder model bundled for offline use',
      dimensions: 512,
      downloadDate: new Date().toISOString(),
      source: 'tensorflow-models/universal-sentence-encoder',
      approach: 'full-bundle',
      modelUrl: modelBaseUrl,
      bundledLocally: true,
      reliability: 'maximum'
    }
    
    fs.writeFileSync(
      path.join(USE_MODEL_DIR, 'metadata.json'),
      JSON.stringify(metadata, null, 2)
    )
    
    // Verify all files exist and calculate total size
    const modelFiles = fs.readdirSync(USE_MODEL_DIR)
    let totalSize = 0
    
    console.log('\n📋 Downloaded files:')
    for (const file of modelFiles) {
      const filePath = path.join(USE_MODEL_DIR, file)
      const stats = fs.statSync(filePath)
      totalSize += stats.size
      console.log(`  ✅ ${file} (${(stats.size / 1024 / 1024).toFixed(2)} MB)`)
    }
    
    console.log(`\n🎉 Model download complete!`)
    console.log(`📊 Total size: ${(totalSize / 1024 / 1024).toFixed(2)} MB`)
    console.log(`📁 Location: ${USE_MODEL_DIR}`)
    console.log(`🔒 Reliability: Maximum (fully offline)`)
    
    // Test loading the downloaded model
    console.log('\n🧪 Testing downloaded model...')
    const offlineModel = await tf.loadGraphModel(`file://${path.join(USE_MODEL_DIR, 'model.json')}`)
    console.log('✅ Offline model loads successfully')
    
    // Clean up
    model.dispose()
    offlineModel.dispose()
    
    console.log('\n✨ Full model bundling completed successfully!')
    console.log('The model is now available for offline use with maximum reliability.')
    
  } catch (error) {
    console.error('❌ Error downloading full model:', error)
    process.exit(1)
  }
}

// Run the download
downloadFullModel().catch(console.error)