diff --git a/bin/enhanced_cli.js b/bin/enhanced_cli.js index 87f5354..a327374 100755 --- a/bin/enhanced_cli.js +++ b/bin/enhanced_cli.js @@ -881,4 +881,223 @@ program } }); +program + .command('ai-check') + .description('AI-powered model selection and execution') + .option('-m, --models ', 'Specific models to choose from') + .option('--prompt ', 'Prompt to run with selected model') + .option('--benchmark', 'Benchmark available models for training data') + .option('--train', 'Train the AI selector model') + .option('--status', 'Show AI model training status') + .action(async (options) => { + const AIModelSelector = require('../src/ai/model-selector'); + const { spawn } = require('child_process'); + + try { + const aiSelector = new AIModelSelector(); + + if (options.status) { + const status = aiSelector.getTrainingStatus(); + console.log('\n' + chalk.bgMagenta.white.bold(' ๐Ÿง  AI MODEL STATUS ')); + console.log(chalk.magenta('โ•ญ' + 'โ”€'.repeat(50))); + console.log(chalk.magenta('โ”‚') + ` Status: ${getStatusColor(status.status)(status.status.replace('_', ' ').toUpperCase())}`); + + if (status.status === 'trained') { + console.log(chalk.magenta('โ”‚') + ` Model size: ${chalk.green.bold(status.modelSize + ' KB')}`); + console.log(chalk.magenta('โ”‚') + ` Version: ${chalk.cyan(status.version)}`); + console.log(chalk.magenta('โ”‚') + ` Features: ${chalk.yellow(status.features)}`); + console.log(chalk.magenta('โ”‚') + ` Last updated: ${chalk.gray(new Date(status.lastUpdated).toLocaleDateString())}`); + console.log(chalk.magenta('โ”‚') + ` Ready for use: ${chalk.green.bold('YES')}`); + } else if (status.status === 'not_trained') { + console.log(chalk.magenta('โ”‚') + ` To get started:`); + console.log(chalk.magenta('โ”‚') + ` 1. ${chalk.cyan.bold('npm run benchmark')} - Collect performance data`); + console.log(chalk.magenta('โ”‚') + ` 2. ${chalk.cyan.bold('npm run train-ai')} - Train the AI model`); + console.log(chalk.magenta('โ”‚') + ` 3. ${chalk.cyan.bold('npm run ai-check')} - Use AI selection`); + } else { + console.log(chalk.magenta('โ”‚') + ` Issue: ${chalk.red('Model files corrupted or incomplete')}`); + console.log(chalk.magenta('โ”‚') + ` Solution: ${chalk.cyan.bold('npm run train-ai')}`); + } + + console.log(chalk.magenta('โ•ฐ')); + return; + } + + if (options.benchmark) { + const spinner = ora('๐Ÿ”ฌ Collecting benchmark data...').start(); + + try { + const benchmarkProcess = spawn('python', [ + 'ml-model/python/benchmark_collector.py' + ], { stdio: 'inherit' }); + + benchmarkProcess.on('close', (code) => { + if (code === 0) { + spinner.succeed('โœ… Benchmark data collected!'); + console.log('\nNext steps:'); + console.log(` 1. ${chalk.cyan.bold('npm run train-ai')} - Train the AI model`); + console.log(` 2. ${chalk.cyan.bold('npm run ai-check')} - Use AI selection`); + } else { + spinner.fail('โŒ Benchmark collection failed'); + } + }); + + return; + } catch (error) { + spinner.fail('โŒ Failed to start benchmark collection'); + console.error(chalk.red('Error:'), error.message); + return; + } + } + + if (options.train) { + const spinner = ora('๐Ÿง  Training AI model...').start(); + + try { + const trainProcess = spawn('python', [ + 'ml-model/python/train_model.py' + ], { stdio: 'inherit' }); + + trainProcess.on('close', (code) => { + if (code === 0) { + spinner.succeed('โœ… AI model trained successfully!'); + console.log('\nYou can now use AI-powered selection:'); + console.log(` ${chalk.cyan.bold('npm run ai-check')}`); + } else { + spinner.fail('โŒ AI model training failed'); + } + }); + + return; + } catch (error) { + spinner.fail('โŒ Failed to start training'); + console.error(chalk.red('Error:'), error.message); + return; + } + } + + // Main AI selection logic + const spinner = ora('๐Ÿง  AI-powered model selection...').start(); + + // Get system info + const checker = new LLMChecker(); + const systemInfo = await checker.getSystemInfo(); + + // Get available models or use provided ones + let candidateModels = options.models; + + if (!candidateModels) { + spinner.text = '๐Ÿ“‹ Getting available Ollama models...'; + const OllamaClient = require('../src/ollama/client'); + const client = new OllamaClient(); + + try { + const models = await client.listModels(); + candidateModels = models.map(m => m.name || m.model); + + if (candidateModels.length === 0) { + spinner.fail('โŒ No Ollama models found'); + console.log('\nInstall some models first:'); + console.log(' ollama pull llama2:7b'); + console.log(' ollama pull mistral:7b'); + console.log(' ollama pull phi3:mini'); + return; + } + } catch (error) { + spinner.fail('โŒ Failed to get Ollama models'); + console.error(chalk.red('Error:'), error.message); + return; + } + } + + spinner.text = '๐Ÿค– Analyzing optimal model selection...'; + + // Use AI selector + const result = await aiSelector.selectBestModel(candidateModels, { + cpu_cores: systemInfo.cpu?.cores || 4, + cpu_freq_max: (systemInfo.cpu?.speed || 3000) / 1000, + total_ram_gb: (systemInfo.memory?.total || 8 * 1024 * 1024 * 1024) / (1024 ** 3), + gpu_vram_gb: (systemInfo.graphics?.vram || 0) / 1024, + gpu_model_normalized: systemInfo.graphics?.model || 'cpu_only' + }); + + spinner.succeed('โœ… AI selection completed!'); + + // Display results + console.log('\n' + chalk.bgMagenta.white.bold(' ๐Ÿง  AI MODEL SELECTION ')); + console.log(chalk.magenta('โ•ญ' + 'โ”€'.repeat(60))); + console.log(chalk.magenta('โ”‚') + ` ๐Ÿ† Selected Model: ${chalk.green.bold(result.bestModel)}`); + console.log(chalk.magenta('โ”‚') + ` ๐ŸŽฏ Selection Method: ${chalk.cyan(result.method.toUpperCase())}`); + console.log(chalk.magenta('โ”‚') + ` ๐Ÿ“Š Confidence: ${getConfidenceColor(result.confidence)(Math.round(result.confidence * 100) + '%')}`); + + if (result.reason) { + console.log(chalk.magenta('โ”‚') + ` ๐Ÿ’ก Reason: ${chalk.yellow(result.reason)}`); + } + + if (result.allPredictions && result.allPredictions.length > 1) { + console.log(chalk.magenta('โ”‚')); + console.log(chalk.magenta('โ”‚') + ` ${chalk.bold.white('Top Candidates:')}`); + result.allPredictions.slice(0, 5).forEach((pred, i) => { + const score = Math.round(pred.score * 100); + const icon = i === 0 ? '๐Ÿฅ‡' : i === 1 ? '๐Ÿฅˆ' : i === 2 ? '๐Ÿฅ‰' : ' '; + console.log(chalk.magenta('โ”‚') + ` ${icon} ${pred.model}: ${getConfidenceColor(pred.score)(score + '%')}`); + }); + } + + console.log(chalk.magenta('โ•ฐ')); + + // Display system info + if (result.systemSpecs) { + const specs = result.systemSpecs; + console.log('\n' + chalk.bgBlue.white.bold(' ๐Ÿ’ป SYSTEM ANALYSIS ')); + console.log(chalk.blue('โ•ญ' + 'โ”€'.repeat(45))); + console.log(chalk.blue('โ”‚') + ` CPU: ${chalk.green(specs.cpu_cores + ' cores')} @ ${chalk.cyan(specs.cpu_freq_max?.toFixed(1) + ' GHz')}`); + console.log(chalk.blue('โ”‚') + ` RAM: ${chalk.green(specs.total_ram_gb?.toFixed(1) + ' GB')}`); + console.log(chalk.blue('โ”‚') + ` GPU: ${chalk.yellow(specs.gpu_model_normalized || 'CPU Only')}`); + console.log(chalk.blue('โ”‚') + ` VRAM: ${chalk.green((specs.gpu_vram_gb || 0).toFixed(1) + ' GB')}`); + console.log(chalk.blue('โ•ฐ')); + } + + // Execute the model if requested + if (options.prompt || process.stdout.isTTY) { + console.log(chalk.magenta.bold('\n๐Ÿš€ Launching selected model...')); + + const args = ['run', result.bestModel]; + if (options.prompt) { + args.push(options.prompt); + } + + const ollamaProcess = spawn('ollama', args, { + stdio: 'inherit' + }); + + ollamaProcess.on('error', (error) => { + console.error(chalk.red('Failed to launch Ollama:'), error.message); + }); + } else { + console.log(chalk.gray('\nTo run the selected model:')); + console.log(chalk.cyan.bold(`ollama run ${result.bestModel}`)); + } + + } catch (error) { + console.error(chalk.red('โŒ AI selection failed:'), error.message); + process.exit(1); + } + }); + +function getStatusColor(status) { + const colors = { + 'TRAINED': chalk.green, + 'NOT TRAINED': chalk.yellow, + 'CORRUPTED': chalk.red + }; + return colors[status] || chalk.gray; +} + +function getConfidenceColor(confidence) { + if (confidence >= 0.8) return chalk.green.bold; + if (confidence >= 0.6) return chalk.yellow.bold; + if (confidence >= 0.4) return chalk.orange.bold; + return chalk.red.bold; +} + program.parse(); \ No newline at end of file diff --git a/ml-model/README.md b/ml-model/README.md new file mode 100644 index 0000000..99f49e3 --- /dev/null +++ b/ml-model/README.md @@ -0,0 +1,199 @@ +# AI Model Selector + +This module implements an AI-powered model selector using a lightweight TabTransformer model that learns from performance benchmarks to recommend the optimal Ollama model for any given hardware configuration. + +## ๐Ÿง  How It Works + +The AI selector uses a small (<150KB) quantized ONNX model trained on hardware specifications and model performance data to predict which model will perform best on your system. + +### Features + +- **Hardware-Aware Selection**: Considers CPU cores, RAM, GPU VRAM, and model architecture +- **Lightweight Model**: Quantized ONNX model under 150KB +- **Fallback Heuristics**: Works even without trained model using smart heuristics +- **Continuous Learning**: Can be retrained with new benchmark data +- **Fast Inference**: Sub-second model selection + +## ๐Ÿš€ Quick Start + +### 1. Check AI Model Status +```bash +npm run ai-check -- --status +``` + +### 2. Collect Benchmark Data (Optional) +```bash +npm run benchmark +``` + +### 3. Train AI Model (Optional) +```bash +npm run train-ai +``` + +### 4. Use AI Selection +```bash +npm run ai-check +npm run ai-check -- --models llama2:7b mistral:7b phi3:mini +npm run ai-check -- --prompt "Explain machine learning" +``` + +## ๐Ÿ“Š Architecture + +### TabTransformer Model +- **Input Features**: Hardware specs (categorical + numerical) +- **Architecture**: 2-layer transformer with 32-dim embeddings +- **Output**: Binary classification (best model probability) +- **Size**: <150KB quantized ONNX + +### Feature Engineering +**Categorical Features:** +- `model_id`: Model identifier +- `gpu_model_normalized`: GPU category +- `hw_platform`: Operating system +- `ram_tier`, `cpu_tier`, `vram_tier`: Hardware capability tiers + +**Numerical Features:** +- `model_size_numeric`: Model parameters in billions +- `hw_cpu_cores`: CPU core count +- `hw_cpu_freq_max`: Maximum CPU frequency +- `hw_total_ram_gb`: System RAM in GB +- `hw_gpu_vram_gb`: GPU VRAM in GB + +## ๐Ÿ”ง Development + +### Python Requirements +```bash +cd ml-model +pip install -r requirements.txt +``` + +### Training Pipeline +1. **Data Collection**: `benchmark_collector.py` + - Runs performance tests on available models + - Collects hardware specifications + - Saves data as Parquet files + +2. **Data Aggregation**: `dataset_aggregator.py` + - Combines benchmark data from multiple machines + - Creates training labels (best model per hardware config) + - Preprocesses features + +3. **Model Training**: `train_model.py` + - Trains TabTransformer on processed data + - Exports to ONNX and quantizes to INT8 + - Validates performance (target: >90% AUC) + +### JavaScript Runtime +- **index.js**: ONNX runtime for model inference +- **cli.js**: Standalone CLI tool +- **test.js**: Testing utilities + +## ๐Ÿ“ˆ Performance Metrics + +The model is trained to achieve: +- **>90% AUC** on validation set +- **<150KB** model size after quantization +- **<100ms** inference time +- **>80% accuracy** on hardware compatibility + +## ๐Ÿ”„ Continuous Improvement + +The model can be continuously improved by: +1. Running benchmarks on new hardware configurations +2. Adding new models to the training set +3. Retraining periodically with updated data +4. Fine-tuning hyperparameters based on performance + +## ๐Ÿ› ๏ธ API Reference + +### AIModelSelector Class + +```javascript +const selector = new AIModelSelector(); + +// Initialize (loads ONNX model) +await selector.initialize(); + +// Select best model +const result = await selector.selectBestModel( + ['llama2:7b', 'mistral:7b'], + systemSpecs +); + +// Fallback selection +const fallback = selector.selectModelHeuristic(models, specs); +``` + +### CLI Commands + +```bash +# AI-powered selection +llm-checker ai-check + +# With specific models +llm-checker ai-check -m llama2:7b mistral:7b + +# With prompt +llm-checker ai-check --prompt "Hello world" + +# Check training status +llm-checker ai-check --status + +# Collect benchmarks +llm-checker ai-check --benchmark + +# Train model +llm-checker ai-check --train +``` + +## ๐Ÿ” Troubleshooting + +### Common Issues + +1. **"ONNX model not found"** + - Run `npm run train-ai` to train the model first + - Or collect benchmarks with `npm run benchmark` + +2. **"Python not found"** + - Install Python โ‰ฅ3.10 + - Install required packages: `pip install -r requirements.txt` + +3. **"No models found"** + - Install Ollama models: `ollama pull llama2:7b` + +4. **Training fails with low AUC** + - Collect more diverse benchmark data + - Run benchmarks on different hardware configurations + +### Debug Mode +```bash +llm-checker ai-check --debug +``` + +## ๐Ÿ“ File Structure + +``` +ml-model/ +โ”œโ”€โ”€ README.md # This file +โ”œโ”€โ”€ requirements.txt # Python dependencies +โ”œโ”€โ”€ python/ # Training pipeline +โ”‚ โ”œโ”€โ”€ benchmark_collector.py +โ”‚ โ”œโ”€โ”€ dataset_aggregator.py +โ”‚ โ””โ”€โ”€ train_model.py +โ”œโ”€โ”€ js/ # JavaScript runtime +โ”‚ โ”œโ”€โ”€ package.json +โ”‚ โ”œโ”€โ”€ index.js +โ”‚ โ”œโ”€โ”€ cli.js +โ”‚ โ””โ”€โ”€ test.js +โ”œโ”€โ”€ data/ # Training data +โ”‚ โ”œโ”€โ”€ raw/ # Benchmark parquet files +โ”‚ โ””โ”€โ”€ processed/ # Processed training data +โ””โ”€โ”€ trained/ # Trained model artifacts + โ”œโ”€โ”€ model_quantized.onnx + โ”œโ”€โ”€ metadata.json + โ”œโ”€โ”€ scaler.joblib + โ””โ”€โ”€ label_encoders.joblib +``` + +This AI-powered approach ensures optimal model selection tailored to your specific hardware, maximizing performance while minimizing resource usage. \ No newline at end of file diff --git a/ml-model/js/cli.js b/ml-model/js/cli.js new file mode 100755 index 0000000..ae0fb47 --- /dev/null +++ b/ml-model/js/cli.js @@ -0,0 +1,194 @@ +#!/usr/bin/env node + +const AIModelSelector = require('./index'); +const { spawn } = require('child_process'); + +class OllamaAISelector { + constructor() { + this.selector = new AIModelSelector(); + } + + async getAvailableModels() { + return new Promise((resolve, reject) => { + const ollama = spawn('ollama', ['list']); + let output = ''; + let error = ''; + + ollama.stdout.on('data', (data) => { + output += data.toString(); + }); + + ollama.stderr.on('data', (data) => { + error += data.toString(); + }); + + ollama.on('close', (code) => { + if (code !== 0) { + reject(new Error(`Ollama list failed: ${error}`)); + return; + } + + // Parse ollama list output + const lines = output.trim().split('\n').slice(1); // Skip header + const models = lines + .filter(line => line.trim()) + .map(line => line.split(/\s+/)[0]) + .filter(model => model && !model.includes('NAME')); + + resolve(models); + }); + }); + } + + async runOllamaModel(modelName, prompt) { + console.log(`๐Ÿš€ Running Ollama with ${modelName}...`); + + return new Promise((resolve, reject) => { + const args = ['run', modelName]; + if (prompt) { + args.push(prompt); + } + + const ollama = spawn('ollama', args, { + stdio: 'inherit' // Stream output directly to terminal + }); + + ollama.on('close', (code) => { + if (code === 0) { + resolve(); + } else { + reject(new Error(`Ollama exited with code ${code}`)); + } + }); + + ollama.on('error', (error) => { + reject(new Error(`Failed to start Ollama: ${error.message}`)); + }); + }); + } + + async selectAndRun(candidateModels = null, prompt = null) { + try { + console.log('๐Ÿง  AI-Powered Ollama Model Selector'); + console.log('=' * 40); + + // Get available models if not provided + if (!candidateModels) { + console.log('๐Ÿ“‹ Getting available models...'); + candidateModels = await this.getAvailableModels(); + + if (candidateModels.length === 0) { + console.log('โŒ No Ollama models found. Install some models first:'); + console.log(' ollama pull llama2:7b'); + console.log(' ollama pull mistral:7b'); + console.log(' ollama pull phi3:mini'); + return; + } + + console.log(`โœ… Found ${candidateModels.length} models: ${candidateModels.join(', ')}`); + } + + // Use AI to select best model + let result; + try { + console.log('๐Ÿ” Analyzing your hardware...'); + result = await this.selector.predictBestModel(candidateModels); + console.log('โœ… AI selection completed'); + } catch (error) { + console.log(`โš ๏ธ AI selection failed: ${error.message}`); + console.log('๐Ÿ”„ Falling back to heuristic selection...'); + + const specs = await this.selector.getSystemSpecs(); + result = this.selector.selectModelHeuristic(candidateModels, specs); + } + + // Display results + console.log('\n๐Ÿ“Š Selection Results:'); + console.log(`๐Ÿ† Best Model: ${result.bestModel}`); + + if (result.systemSpecs) { + const specs = result.systemSpecs; + console.log(`๐Ÿ’ป System: ${specs.cpu_cores} cores, ${specs.total_ram_gb.toFixed(1)}GB RAM`); + console.log(`๐ŸŽฎ GPU: ${specs.gpu_model_normalized} (${specs.gpu_vram_gb.toFixed(1)}GB VRAM)`); + } + + if (result.allPredictions) { + console.log('\n๐Ÿ“ˆ All Model Scores:'); + result.allPredictions.slice(0, 5).forEach((pred, i) => { + const score = (pred.score * 100).toFixed(1); + const icon = i === 0 ? '๐Ÿฅ‡' : i === 1 ? '๐Ÿฅˆ' : i === 2 ? '๐Ÿฅ‰' : ' '; + console.log(`${icon} ${pred.model}: ${score}% (${pred.size}B params)`); + }); + } + + // Run the selected model + console.log(`\n๐Ÿš€ Launching ${result.bestModel}...`); + await this.runOllamaModel(result.bestModel, prompt); + + } catch (error) { + console.error('โŒ Error:', error.message); + process.exit(1); + } + } +} + +// CLI interface +async function main() { + const args = process.argv.slice(2); + + if (args.includes('--help') || args.includes('-h')) { + console.log('AI-Powered Ollama Model Selector'); + console.log(''); + console.log('Automatically selects the best Ollama model for your hardware using AI.'); + console.log(''); + console.log('Usage:'); + console.log(' ollama-select # Select from all available models'); + console.log(' ollama-select model1 model2 # Select from specific models'); + console.log(' ollama-select --test # Test the selector'); + console.log(''); + console.log('Examples:'); + console.log(' ollama-select llama2:7b mistral:7b phi3:mini'); + console.log(' ollama-select'); + console.log(''); + return; + } + + if (args.includes('--test')) { + console.log('๐Ÿงช Testing AI Model Selector...'); + + const selector = new AIModelSelector(); + try { + await selector.initialize(); + + const testModels = ['llama2:7b', 'mistral:7b', 'phi3:mini']; + const result = await selector.predictBestModel(testModels); + + console.log('โœ… Test successful!'); + console.log(`Best model: ${result.bestModel}`); + console.log('Scores:', result.allPredictions.map(p => `${p.model}: ${(p.score * 100).toFixed(1)}%`)); + + } catch (error) { + console.log('โŒ Test failed:', error.message); + + // Test fallback + console.log('๐Ÿ”„ Testing fallback heuristic...'); + const result = selector.selectModelHeuristic(testModels); + console.log(`Fallback result: ${result.bestModel} (${result.reason})`); + } + + return; + } + + // Extract model candidates and prompt + const candidateModels = args.filter(arg => !arg.startsWith('--')); + const prompt = null; // Could be extended to accept prompts + + const cli = new OllamaAISelector(); + await cli.selectAndRun(candidateModels.length > 0 ? candidateModels : null, prompt); +} + +if (require.main === module) { + main().catch(console.error); +} + +module.exports = OllamaAISelector; \ No newline at end of file diff --git a/ml-model/js/index.js b/ml-model/js/index.js new file mode 100644 index 0000000..2160b07 --- /dev/null +++ b/ml-model/js/index.js @@ -0,0 +1,370 @@ +const fs = require('fs'); +const path = require('path'); + +// Try to load optional dependencies +let ort = null; +let si = null; + +try { + ort = require('onnxruntime-node'); +} catch (e) { + console.warn('ONNX Runtime not available, using fallback mode'); +} + +try { + si = require('systeminformation'); +} catch (e) { + console.warn('systeminformation not available, using basic system detection'); +} + +class AIModelSelector { + constructor() { + this.session = null; + this.metadata = null; + this.scaler = null; + this.encoders = null; + this.initialized = false; + } + + async initialize() { + if (this.initialized) return; + + if (!ort) { + throw new Error('ONNX Runtime not available. Install with: npm install onnxruntime-node'); + } + + try { + // Load metadata + const metadataPath = path.join(__dirname, '../trained/metadata.json'); + if (!fs.existsSync(metadataPath)) { + throw new Error('Model metadata not found. Please train the model first.'); + } + + this.metadata = JSON.parse(fs.readFileSync(metadataPath, 'utf8')); + + // Load ONNX model + const modelPath = path.join(__dirname, '../trained/model_quantized.onnx'); + if (!fs.existsSync(modelPath)) { + throw new Error('ONNX model not found. Please train the model first.'); + } + + this.session = await ort.InferenceSession.create(modelPath); + + // Load preprocessing objects (simplified - we'll implement basic versions) + this.scaler = this.loadScaler(); + this.encoders = this.loadEncoders(); + + this.initialized = true; + console.log('โœ… AI Model Selector initialized successfully'); + + } catch (error) { + console.error('โŒ Failed to initialize AI Model Selector:', error.message); + throw error; + } + } + + loadScaler() { + // Simplified scaler implementation + // In production, you'd load the actual joblib scaler + return { + mean: [7.2, 3.2, 16.5, 8.2, 0], // Default values + scale: [10.1, 1.1, 15.2, 12.1, 1], + transform: function(data) { + return data.map((value, index) => { + return (value - this.mean[index]) / this.scale[index]; + }); + } + }; + } + + loadEncoders() { + // Simplified label encoders + // In production, you'd load the actual joblib encoders + return { + model_id: { + 'llama2:7b': 0, 'llama2:13b': 1, 'llama2:70b': 2, + 'mistral:7b': 3, 'phi3:mini': 4, 'gemma:7b': 5, + 'codellama:7b': 6, 'deepseek-coder:6.7b': 7 + }, + gpu_model_normalized: { + 'cpu_only': 0, 'rtx_4090': 1, 'rtx_4080': 2, 'rtx_4070': 3, + 'rtx_3090': 4, 'rtx_3080': 5, 'apple_silicon': 6, 'nvidia_gpu_other': 7 + }, + hw_platform: { + 'Darwin': 0, 'Linux': 1, 'Windows': 2 + }, + ram_tier: { + 'low': 0, 'medium': 1, 'high': 2, 'very_high': 3, 'extreme': 4 + }, + cpu_tier: { + 'low': 0, 'medium': 1, 'high': 2, 'very_high': 3, 'extreme': 4 + }, + vram_tier: { + 'none': 0, 'low': 1, 'medium': 2, 'high': 3, 'very_high': 4, 'extreme': 5 + } + }; + } + + async getSystemSpecs() { + if (si) { + try { + const [cpu, mem, graphics, osInfo] = await Promise.all([ + si.cpu(), + si.mem(), + si.graphics(), + si.osInfo() + ]); + + const specs = { + cpu_cores: cpu.cores || 4, + cpu_freq_max: (cpu.speed || 3000) / 1000, // Convert to GHz + total_ram_gb: (mem.total || 8 * 1024 * 1024 * 1024) / (1024 ** 3), + platform: osInfo.platform || 'unknown' + }; + + // GPU detection + let gpu_model = 'cpu_only'; + let gpu_vram_gb = 0; + + if (graphics.controllers && graphics.controllers.length > 0) { + const gpu = graphics.controllers[0]; + const model = (gpu.model || '').toLowerCase(); + + if (model.includes('rtx 4090')) gpu_model = 'rtx_4090'; + else if (model.includes('rtx 4080')) gpu_model = 'rtx_4080'; + else if (model.includes('rtx 4070')) gpu_model = 'rtx_4070'; + else if (model.includes('rtx 3090')) gpu_model = 'rtx_3090'; + else if (model.includes('rtx 3080')) gpu_model = 'rtx_3080'; + else if (model.includes('apple')) gpu_model = 'apple_silicon'; + else if (model.includes('nvidia') || model.includes('geforce')) gpu_model = 'nvidia_gpu_other'; + + gpu_vram_gb = (gpu.vram || 0) / 1024; // Convert MB to GB + + // Special case for Apple Silicon - use unified memory + if (gpu_model === 'apple_silicon') { + gpu_vram_gb = specs.total_ram_gb * 0.75; // Assume 75% available + } + } + + specs.gpu_model_normalized = gpu_model; + specs.gpu_vram_gb = gpu_vram_gb; + + // Create tiers + specs.ram_tier = this.getRamTier(specs.total_ram_gb); + specs.cpu_tier = this.getCpuTier(specs.cpu_cores); + specs.vram_tier = this.getVramTier(specs.gpu_vram_gb); + + return specs; + + } catch (error) { + console.error('Error getting system specs:', error); + } + } + + // Fallback to basic detection + const os = require('os'); + const platform = os.platform(); + const totalMem = os.totalmem() / (1024 ** 3); // GB + const cpuCount = os.cpus().length; + + const specs = { + cpu_cores: cpuCount, + cpu_freq_max: 3.0, // Default + total_ram_gb: totalMem, + gpu_model_normalized: platform === 'darwin' ? 'apple_silicon' : 'cpu_only', + gpu_vram_gb: platform === 'darwin' ? totalMem * 0.75 : 0, + platform: platform + }; + + specs.ram_tier = this.getRamTier(specs.total_ram_gb); + specs.cpu_tier = this.getCpuTier(specs.cpu_cores); + specs.vram_tier = this.getVramTier(specs.gpu_vram_gb); + + return specs; + } + + getRamTier(ramGB) { + if (ramGB <= 8) return 'low'; + if (ramGB <= 16) return 'medium'; + if (ramGB <= 32) return 'high'; + if (ramGB <= 64) return 'very_high'; + return 'extreme'; + } + + getCpuTier(cores) { + if (cores <= 4) return 'low'; + if (cores <= 8) return 'medium'; + if (cores <= 16) return 'high'; + if (cores <= 32) return 'very_high'; + return 'extreme'; + } + + getVramTier(vramGB) { + if (vramGB <= 0) return 'none'; + if (vramGB <= 4) return 'low'; + if (vramGB <= 8) return 'medium'; + if (vramGB <= 16) return 'high'; + if (vramGB <= 24) return 'very_high'; + return 'extreme'; + } + + getModelSize(modelId) { + // Extract model size from ID + const sizeMatch = modelId.match(/(\d+\.?\d*)[kmb]/i); + if (sizeMatch) { + const num = parseFloat(sizeMatch[1]); + const unit = sizeMatch[0].slice(-1).toLowerCase(); + + if (unit === 'k') return num / 1000; + if (unit === 'm') return num / 1000; + if (unit === 'b') return num; + } + + // Default sizes for common models + if (modelId.includes('mini')) return 3.8; + if (modelId.includes('7b')) return 7; + if (modelId.includes('13b')) return 13; + if (modelId.includes('70b')) return 70; + + return 7; // Default + } + + encodeFeatures(specs, modelId) { + const features = {}; + + // Categorical features + const categoricalFeatures = this.metadata.categorical_features; + + for (const feature of categoricalFeatures) { + let value; + + if (feature === 'model_id') { + value = this.encoders[feature][modelId] || 0; + } else if (feature === 'gpu_model_normalized') { + value = this.encoders[feature][specs.gpu_model_normalized] || 0; + } else if (feature === 'hw_platform') { + value = this.encoders[feature][specs.platform] || 0; + } else if (feature === 'ram_tier') { + value = this.encoders[feature][specs.ram_tier] || 1; + } else if (feature === 'cpu_tier') { + value = this.encoders[feature][specs.cpu_tier] || 1; + } else if (feature === 'vram_tier') { + value = this.encoders[feature][specs.vram_tier] || 0; + } else { + value = 0; + } + + features[feature] = new ort.Tensor('int32', [value], [1]); + } + + // Numerical features + const numericalData = [ + this.getModelSize(modelId), + specs.cpu_cores, + specs.cpu_freq_max, + specs.total_ram_gb, + specs.gpu_vram_gb + ]; + + // Apply scaling + const scaledData = this.scaler.transform(numericalData); + + const numericalFeatures = this.metadata.numerical_features; + for (let i = 0; i < numericalFeatures.length; i++) { + const feature = numericalFeatures[i]; + features[feature] = new ort.Tensor('float32', [scaledData[i]], [1, 1]); + } + + return features; + } + + async predictBestModel(modelCandidates, specs = null) { + if (!this.initialized) { + await this.initialize(); + } + + if (!specs) { + specs = await this.getSystemSpecs(); + } + + const predictions = []; + + for (const modelId of modelCandidates) { + try { + const features = this.encodeFeatures(specs, modelId); + const result = await this.session.run(features); + + const outputTensor = result[Object.keys(result)[0]]; + const score = outputTensor.data[0]; + + predictions.push({ + model: modelId, + score: score, + size: this.getModelSize(modelId) + }); + + } catch (error) { + console.warn(`Warning: Failed to predict for ${modelId}:`, error.message); + // Add with low score as fallback + predictions.push({ + model: modelId, + score: 0.1, + size: this.getModelSize(modelId) + }); + } + } + + // Sort by score (descending) + predictions.sort((a, b) => b.score - a.score); + + return { + bestModel: predictions[0]?.model || modelCandidates[0], + allPredictions: predictions, + systemSpecs: specs + }; + } + + // Fallback heuristic method + selectModelHeuristic(modelCandidates, specs = null) { + console.log('๐Ÿ”„ Using fallback heuristic selection...'); + + if (!specs) { + // Use simplified specs if not provided + specs = { + total_ram_gb: 8, + gpu_vram_gb: 0, + cpu_cores: 4 + }; + } + + // Simple heuristic based on available memory + const availableMemory = specs.gpu_vram_gb > 0 ? + specs.gpu_vram_gb : specs.total_ram_gb * 0.7; // 70% of RAM for CPU-only + + const suitableModels = modelCandidates + .map(model => ({ + model, + size: this.getModelSize(model), + memoryReq: this.getModelSize(model) * 1.2 // Rough estimate + })) + .filter(m => m.memoryReq <= availableMemory) + .sort((a, b) => b.size - a.size); // Prefer larger models that fit + + const bestModel = suitableModels.length > 0 ? + suitableModels[0].model : + modelCandidates.reduce((a, b) => + this.getModelSize(a) < this.getModelSize(b) ? a : b + ); + + return { + bestModel, + method: 'heuristic', + systemSpecs: specs, + reason: suitableModels.length > 0 ? + 'Best fitting model for available memory' : + 'Smallest available model (fallback)' + }; + } +} + +module.exports = AIModelSelector; \ No newline at end of file diff --git a/ml-model/js/test.js b/ml-model/js/test.js new file mode 100644 index 0000000..9be48cd --- /dev/null +++ b/ml-model/js/test.js @@ -0,0 +1,57 @@ +const AIModelSelector = require('./index'); + +async function testSelector() { + console.log('๐Ÿงช Testing AI Model Selector...'); + + const selector = new AIModelSelector(); + + try { + // Test system specs detection + console.log('\n1. Testing system specs detection...'); + const specs = await selector.getSystemSpecs(); + console.log('System specs:', { + cpu_cores: specs.cpu_cores, + ram_gb: specs.total_ram_gb.toFixed(1), + gpu: specs.gpu_model_normalized, + vram_gb: specs.gpu_vram_gb.toFixed(1) + }); + + // Test heuristic selection (always works) + console.log('\n2. Testing heuristic selection...'); + const testModels = ['llama2:7b', 'mistral:7b', 'phi3:mini', 'llama2:13b']; + const heuristicResult = selector.selectModelHeuristic(testModels, specs); + console.log('Heuristic result:', { + bestModel: heuristicResult.bestModel, + method: heuristicResult.method, + reason: heuristicResult.reason + }); + + // Test AI selection (might fail if model not trained) + console.log('\n3. Testing AI selection...'); + try { + const aiResult = await selector.predictBestModel(testModels, specs); + console.log('AI result:', { + bestModel: aiResult.bestModel, + topScores: aiResult.allPredictions.slice(0, 3).map(p => ({ + model: p.model, + score: (p.score * 100).toFixed(1) + '%' + })) + }); + console.log('โœ… AI selection working!'); + } catch (error) { + console.log('โš ๏ธ AI selection not available:', error.message); + console.log(' This is expected if the model hasn\'t been trained yet.'); + } + + console.log('\n๐ŸŽ‰ Test completed successfully!'); + + } catch (error) { + console.error('โŒ Test failed:', error); + } +} + +if (require.main === module) { + testSelector(); +} + +module.exports = testSelector; \ No newline at end of file diff --git a/ml-model/python/benchmark_collector.py b/ml-model/python/benchmark_collector.py new file mode 100644 index 0000000..b6deb04 --- /dev/null +++ b/ml-model/python/benchmark_collector.py @@ -0,0 +1,352 @@ +#!/usr/bin/env python3 +""" +Performance benchmark collector for Ollama models. +Collects hardware specs and model performance data for ML training. +""" + +import json +import time +import subprocess +import platform +import hashlib +import os +from pathlib import Path +from dataclasses import dataclass, asdict +from typing import List, Dict, Optional, Tuple +import pandas as pd +import psutil +import sys + +@dataclass +class HardwareSpec: + """Hardware specification for a given system.""" + host_id: str + cpu_cores: int + cpu_freq_max: float # GHz + total_ram_gb: float + gpu_model: str + gpu_vram_gb: float + platform: str + architecture: str + +@dataclass +class ModelBenchmark: + """Benchmark result for a specific model on specific hardware.""" + host_id: str + model_id: str + model_size: str + tokens_per_second: float + peak_ram_mb: float + peak_vram_mb: float + response_time: float + success: bool + error_message: Optional[str] = None + timestamp: str = "" + +class HardwareDetector: + """Detects hardware specifications of the current system.""" + + def __init__(self): + self.host_id = self._generate_host_id() + + def _generate_host_id(self) -> str: + """Generate a unique but anonymous host identifier.""" + system_info = f"{platform.node()}-{platform.system()}-{platform.machine()}" + return hashlib.md5(system_info.encode()).hexdigest()[:12] + + def get_hardware_spec(self) -> HardwareSpec: + """Collect comprehensive hardware specifications.""" + # CPU information + cpu_cores = psutil.cpu_count(logical=False) + cpu_freq = psutil.cpu_freq() + cpu_freq_max = cpu_freq.max / 1000 if cpu_freq else 3.0 # Default to 3GHz + + # Memory information + memory = psutil.virtual_memory() + total_ram_gb = memory.total / (1024**3) + + # GPU information + gpu_model, gpu_vram_gb = self._detect_gpu() + + return HardwareSpec( + host_id=self.host_id, + cpu_cores=cpu_cores, + cpu_freq_max=cpu_freq_max, + total_ram_gb=total_ram_gb, + gpu_model=gpu_model, + gpu_vram_gb=gpu_vram_gb, + platform=platform.system(), + architecture=platform.machine() + ) + + def _detect_gpu(self) -> Tuple[str, float]: + """Detect GPU model and VRAM.""" + try: + # Try nvidia-smi first + result = subprocess.run(['nvidia-smi', '--query-gpu=name,memory.total', + '--format=csv,noheader,nounits'], + capture_output=True, text=True, timeout=10) + if result.returncode == 0: + lines = result.stdout.strip().split('\n') + if lines and lines[0]: + parts = lines[0].split(', ') + if len(parts) >= 2: + gpu_name = parts[0].strip() + vram_mb = float(parts[1].strip()) + return gpu_name, vram_mb / 1024 # Convert to GB + except (subprocess.TimeoutExpired, subprocess.CalledProcessError, FileNotFoundError): + pass + + # Try other methods or return CPU-only + try: + # Check for Apple Silicon + if platform.system() == "Darwin" and platform.machine() == "arm64": + # Apple Silicon - use unified memory + memory = psutil.virtual_memory() + return "Apple Silicon GPU", memory.total / (1024**3) * 0.75 # Assume 75% available for GPU + except: + pass + + # Default to CPU-only + return "CPU Only", 0.0 + +class OllamaBenchmarker: + """Benchmarks Ollama model performance.""" + + def __init__(self): + self.test_prompt = "Explain in one sentence what machine learning is." + + def get_available_models(self) -> List[str]: + """Get list of locally available Ollama models.""" + try: + result = subprocess.run(['ollama', 'list'], capture_output=True, text=True, timeout=30) + if result.returncode != 0: + print(f"Error getting models: {result.stderr}") + return [] + + lines = result.stdout.strip().split('\n')[1:] # Skip header + models = [] + + for line in lines: + if line.strip(): + parts = line.split() + if parts: + model_name = parts[0] + models.append(model_name) + + return models + + except (subprocess.TimeoutExpired, subprocess.CalledProcessError, FileNotFoundError) as e: + print(f"Error: {e}") + return [] + + def benchmark_model(self, model_id: str, hardware_spec: HardwareSpec) -> ModelBenchmark: + """Benchmark a specific model.""" + print(f"Benchmarking {model_id}...") + + # Initialize benchmark result + benchmark = ModelBenchmark( + host_id=hardware_spec.host_id, + model_id=model_id, + model_size=self._extract_size_from_model_id(model_id), + tokens_per_second=0.0, + peak_ram_mb=0.0, + peak_vram_mb=0.0, + response_time=0.0, + success=False, + timestamp=time.strftime("%Y-%m-%d %H:%M:%S") + ) + + try: + # Record initial memory + process = psutil.Process() + initial_memory = process.memory_info().rss / (1024**2) # MB + + # Run the benchmark + start_time = time.time() + + result = subprocess.run([ + 'ollama', 'run', model_id, self.test_prompt + ], capture_output=True, text=True, timeout=120) + + end_time = time.time() + response_time = end_time - start_time + + if result.returncode == 0: + # Parse output for tokens/second if available + output = result.stderr if result.stderr else result.stdout + tokens_per_second = self._extract_tokens_per_second(output) + + # Estimate peak memory usage (simplified) + final_memory = process.memory_info().rss / (1024**2) # MB + peak_ram_mb = max(final_memory - initial_memory, 0) + + benchmark.tokens_per_second = tokens_per_second + benchmark.peak_ram_mb = peak_ram_mb + benchmark.peak_vram_mb = 0.0 # TODO: Implement GPU memory tracking + benchmark.response_time = response_time + benchmark.success = True + + print(f" โœ… {model_id}: {tokens_per_second:.2f} tokens/s, {response_time:.2f}s") + + else: + benchmark.error_message = result.stderr + print(f" โŒ {model_id}: {result.stderr}") + + except subprocess.TimeoutExpired: + benchmark.error_message = "Timeout" + print(f" โฐ {model_id}: Timeout") + except Exception as e: + benchmark.error_message = str(e) + print(f" โŒ {model_id}: {e}") + + return benchmark + + def _extract_size_from_model_id(self, model_id: str) -> str: + """Extract model size from model identifier.""" + # Extract patterns like 7b, 13b, 70b, etc. + import re + size_match = re.search(r'(\d+\.?\d*[kmb])', model_id.lower()) + return size_match.group(1) if size_match else "unknown" + + def _extract_tokens_per_second(self, output: str) -> float: + """Extract tokens per second from Ollama output.""" + import re + + # Look for patterns like "45.67 tokens/s" or similar + patterns = [ + r'(\d+\.?\d*)\s*tokens?/s', + r'(\d+\.?\d*)\s*tok/s', + r'speed:\s*(\d+\.?\d*)', + ] + + for pattern in patterns: + match = re.search(pattern, output, re.IGNORECASE) + if match: + try: + return float(match.group(1)) + except ValueError: + continue + + # If no explicit rate found, estimate from response time and output length + # This is a rough approximation + word_count = len(output.split()) + if word_count > 10: # Reasonable response + return word_count * 1.3 # Rough tokens per word estimate + + return 0.0 + + def benchmark_all_models(self, hardware_spec: HardwareSpec) -> List[ModelBenchmark]: + """Benchmark all available models.""" + models = self.get_available_models() + + if not models: + print("No models found. Please install some models first:") + print(" ollama pull llama2:7b") + print(" ollama pull mistral:7b") + print(" ollama pull phi3:mini") + return [] + + print(f"Found {len(models)} models to benchmark") + benchmarks = [] + + for model_id in models: + benchmark = self.benchmark_model(model_id, hardware_spec) + benchmarks.append(benchmark) + + # Small delay between benchmarks + time.sleep(2) + + return benchmarks + +class BenchmarkCollector: + """Main collector that coordinates hardware detection and model benchmarking.""" + + def __init__(self, output_dir: str = "data/raw"): + self.output_dir = Path(output_dir) + self.output_dir.mkdir(parents=True, exist_ok=True) + + self.hardware_detector = HardwareDetector() + self.benchmarker = OllamaBenchmarker() + + def collect_benchmark_data(self) -> str: + """Collect complete benchmark data for this system.""" + print("๐Ÿ” Detecting hardware specifications...") + hardware_spec = self.hardware_detector.get_hardware_spec() + + print(f"Hardware detected:") + print(f" Host ID: {hardware_spec.host_id}") + print(f" CPU: {hardware_spec.cpu_cores} cores @ {hardware_spec.cpu_freq_max:.1f} GHz") + print(f" RAM: {hardware_spec.total_ram_gb:.1f} GB") + print(f" GPU: {hardware_spec.gpu_model} ({hardware_spec.gpu_vram_gb:.1f} GB VRAM)") + + print("\n๐Ÿš€ Starting model benchmarks...") + benchmarks = self.benchmarker.benchmark_all_models(hardware_spec) + + if not benchmarks: + print("No benchmarks collected.") + return "" + + # Convert to DataFrame + benchmark_data = [asdict(b) for b in benchmarks] + df = pd.DataFrame(benchmark_data) + + # Add hardware info to each row + hardware_dict = asdict(hardware_spec) + for key, value in hardware_dict.items(): + if key != 'host_id': # host_id already in benchmark data + df[f'hw_{key}'] = value + + # Save to parquet + output_file = self.output_dir / f"benchmark_{hardware_spec.host_id}_{int(time.time())}.parquet" + df.to_parquet(output_file, index=False) + + print(f"\nโœ… Benchmark data saved to: {output_file}") + print(f"Collected {len(benchmarks)} model benchmarks") + + return str(output_file) + +def main(): + """Main entry point.""" + if len(sys.argv) > 1 and sys.argv[1] == "--help": + print("Ollama Model Performance Benchmark Collector") + print("") + print("This script benchmarks all locally installed Ollama models") + print("and collects hardware specifications for ML training.") + print("") + print("Prerequisites:") + print(" - Ollama installed and running") + print(" - At least one model installed (ollama pull )") + print(" - Python packages: pandas, psutil, pyarrow") + print("") + print("Usage:") + print(" python benchmark_collector.py") + print("") + return + + print("๐Ÿง  Ollama Model Performance Benchmark Collector") + print("=" * 50) + + try: + collector = BenchmarkCollector() + output_file = collector.collect_benchmark_data() + + if output_file: + print(f"\n๐ŸŽ‰ Collection complete! Data saved to: {output_file}") + print("\nNext steps:") + print(" 1. Run this on multiple machines with different hardware") + print(" 2. Aggregate all parquet files for training") + print(" 3. Train the ML model selector") + else: + print("\nโŒ No data collected. Please install Ollama models first.") + + except KeyboardInterrupt: + print("\nโน๏ธ Collection cancelled by user") + except Exception as e: + print(f"\nโŒ Error during collection: {e}") + if "--debug" in sys.argv: + import traceback + traceback.print_exc() + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/ml-model/python/dataset_aggregator.py b/ml-model/python/dataset_aggregator.py new file mode 100644 index 0000000..1f0a65e --- /dev/null +++ b/ml-model/python/dataset_aggregator.py @@ -0,0 +1,337 @@ +#!/usr/bin/env python3 +""" +Dataset aggregator for ML model training. +Combines benchmark data from multiple hosts and creates training labels. +""" + +import pandas as pd +import numpy as np +from pathlib import Path +import sys +from typing import List, Dict, Tuple +import warnings +warnings.filterwarnings('ignore') + +class DatasetAggregator: + """Aggregates benchmark data and creates training labels.""" + + def __init__(self, raw_data_dir: str = "data/raw", output_dir: str = "data/processed"): + self.raw_data_dir = Path(raw_data_dir) + self.output_dir = Path(output_dir) + self.output_dir.mkdir(parents=True, exist_ok=True) + + def load_all_benchmarks(self) -> pd.DataFrame: + """Load and concatenate all benchmark parquet files.""" + parquet_files = list(self.raw_data_dir.glob("benchmark_*.parquet")) + + if not parquet_files: + raise FileNotFoundError(f"No benchmark files found in {self.raw_data_dir}") + + print(f"๐Ÿ“ Found {len(parquet_files)} benchmark files") + + dfs = [] + for file in parquet_files: + try: + df = pd.read_parquet(file) + print(f" โœ… Loaded {file.name}: {len(df)} records") + dfs.append(df) + except Exception as e: + print(f" โŒ Error loading {file.name}: {e}") + + if not dfs: + raise ValueError("No valid benchmark files could be loaded") + + combined_df = pd.concat(dfs, ignore_index=True) + print(f"๐Ÿ“Š Combined dataset: {len(combined_df)} total records") + + return combined_df + + def clean_and_preprocess(self, df: pd.DataFrame) -> pd.DataFrame: + """Clean and preprocess the dataset.""" + print("๐Ÿงน Cleaning and preprocessing data...") + + # Remove failed benchmarks + initial_count = len(df) + df = df[df['success'] == True].copy() + print(f" Removed {initial_count - len(df)} failed benchmarks") + + # Remove benchmarks with zero or negative performance + df = df[df['tokens_per_second'] > 0].copy() + + # Extract numeric model size + df['model_size_numeric'] = df['model_size'].apply(self._parse_model_size) + + # Normalize GPU model names + df['gpu_model_normalized'] = df['hw_gpu_model'].apply(self._normalize_gpu_name) + + # Create categorical bins for hardware specs + df['ram_tier'] = pd.cut(df['hw_total_ram_gb'], + bins=[0, 8, 16, 32, 64, float('inf')], + labels=['low', 'medium', 'high', 'very_high', 'extreme']) + + df['cpu_tier'] = pd.cut(df['hw_cpu_cores'], + bins=[0, 4, 8, 16, 32, float('inf')], + labels=['low', 'medium', 'high', 'very_high', 'extreme']) + + df['vram_tier'] = pd.cut(df['hw_gpu_vram_gb'], + bins=[-1, 0, 4, 8, 16, 24, float('inf')], + labels=['none', 'low', 'medium', 'high', 'very_high', 'extreme']) + + print(f" Final dataset: {len(df)} records") + return df + + def _parse_model_size(self, size_str: str) -> float: + """Parse model size string to numeric value in billions of parameters.""" + if pd.isna(size_str) or size_str == 'unknown': + return 1.0 # Default small size + + size_str = str(size_str).lower().strip() + + # Extract number and unit + import re + match = re.match(r'(\d+\.?\d*)([kmb]?)', size_str) + + if not match: + return 1.0 + + num = float(match.group(1)) + unit = match.group(2) + + if unit == 'k': + return num / 1000 # Convert to billions + elif unit == 'm': + return num / 1000 # Convert to billions + elif unit == 'b' or unit == '': + return num + else: + return num + + def _normalize_gpu_name(self, gpu_name: str) -> str: + """Normalize GPU names to common categories.""" + if pd.isna(gpu_name): + return 'cpu_only' + + gpu_name = str(gpu_name).lower() + + # NVIDIA GPUs + if 'rtx 4090' in gpu_name: + return 'rtx_4090' + elif 'rtx 4080' in gpu_name: + return 'rtx_4080' + elif 'rtx 4070' in gpu_name: + return 'rtx_4070' + elif 'rtx 4060' in gpu_name: + return 'rtx_4060' + elif 'rtx 3090' in gpu_name: + return 'rtx_3090' + elif 'rtx 3080' in gpu_name: + return 'rtx_3080' + elif 'rtx 3070' in gpu_name: + return 'rtx_3070' + elif 'rtx 3060' in gpu_name: + return 'rtx_3060' + elif 'gtx 1080' in gpu_name: + return 'gtx_1080' + elif 'gtx 1070' in gpu_name: + return 'gtx_1070' + elif 'gtx 1060' in gpu_name: + return 'gtx_1060' + elif 'apple silicon' in gpu_name: + return 'apple_silicon' + elif 'cpu only' in gpu_name: + return 'cpu_only' + elif 'amd' in gpu_name or 'radeon' in gpu_name: + return 'amd_gpu' + elif 'intel' in gpu_name: + return 'intel_gpu' + elif 'nvidia' in gpu_name or 'geforce' in gpu_name: + return 'nvidia_gpu_other' + else: + return 'unknown_gpu' + + def create_training_labels(self, df: pd.DataFrame) -> pd.DataFrame: + """Create binary labels for which model is best for each hardware configuration.""" + print("๐Ÿท๏ธ Creating training labels...") + + # Initialize label column + df['label_best'] = 0 + + labeled_count = 0 + + # Group by unique hardware configurations + hardware_cols = ['host_id', 'hw_cpu_cores', 'hw_cpu_freq_max', + 'hw_total_ram_gb', 'gpu_model_normalized', 'hw_gpu_vram_gb'] + + for hardware_config, group in df.groupby(hardware_cols): + if len(group) < 2: # Need at least 2 models to compare + continue + + # Find models that can actually run (reasonable performance thresholds) + viable_models = group[ + (group['tokens_per_second'] > 1.0) & # At least 1 token/s + (group['response_time'] < 60.0) & # Response within 60 seconds + (group['peak_ram_mb'] < group['hw_total_ram_gb'] * 1024 * 0.8) # Uses <80% RAM + ].copy() + + if len(viable_models) == 0: + continue + + # Calculate performance score (tokens/s weighted by efficiency) + viable_models['efficiency_score'] = ( + viable_models['tokens_per_second'] / + (viable_models['model_size_numeric'] + 1) # Avoid division by zero + ) + + # Find the best model (highest tokens/s among viable ones) + best_idx = viable_models['tokens_per_second'].idxmax() + df.loc[best_idx, 'label_best'] = 1 + labeled_count += 1 + + positive_labels = (df['label_best'] == 1).sum() + print(f" Created {positive_labels} positive labels from {labeled_count} hardware configs") + print(f" Label distribution: {positive_labels} positive, {len(df) - positive_labels} negative") + + return df + + def prepare_features(self, df: pd.DataFrame) -> pd.DataFrame: + """Prepare features for ML training.""" + print("๐Ÿ”ง Preparing features...") + + # Select and create final feature set + feature_df = df[[ + 'host_id', + 'model_id', + 'model_size_numeric', + 'hw_cpu_cores', + 'hw_cpu_freq_max', + 'hw_total_ram_gb', + 'gpu_model_normalized', + 'hw_gpu_vram_gb', + 'hw_platform', + 'ram_tier', + 'cpu_tier', + 'vram_tier', + 'tokens_per_second', + 'peak_ram_mb', + 'response_time', + 'label_best' + ]].copy() + + # Handle missing values + feature_df['hw_cpu_freq_max'].fillna(3.0, inplace=True) # Default CPU frequency + feature_df['hw_gpu_vram_gb'].fillna(0.0, inplace=True) # CPU-only systems + + print(f" Final feature set: {len(feature_df.columns)} columns, {len(feature_df)} rows") + + return feature_df + + def analyze_dataset(self, df: pd.DataFrame) -> Dict: + """Analyze the dataset and provide statistics.""" + print("\n๐Ÿ“Š Dataset Analysis:") + print("=" * 40) + + analysis = {} + + # Basic stats + analysis['total_records'] = len(df) + analysis['unique_hosts'] = df['host_id'].nunique() + analysis['unique_models'] = df['model_id'].nunique() + analysis['positive_labels'] = (df['label_best'] == 1).sum() + analysis['label_ratio'] = analysis['positive_labels'] / len(df) + + print(f"Total records: {analysis['total_records']}") + print(f"Unique hosts: {analysis['unique_hosts']}") + print(f"Unique models: {analysis['unique_models']}") + print(f"Positive labels: {analysis['positive_labels']} ({analysis['label_ratio']:.2%})") + + # Hardware distribution + print(f"\nHardware Distribution:") + print(f"CPU cores: {df['hw_cpu_cores'].min()}-{df['hw_cpu_cores'].max()} (avg: {df['hw_cpu_cores'].mean():.1f})") + print(f"RAM: {df['hw_total_ram_gb'].min():.1f}-{df['hw_total_ram_gb'].max():.1f} GB (avg: {df['hw_total_ram_gb'].mean():.1f})") + print(f"GPU VRAM: {df['hw_gpu_vram_gb'].min():.1f}-{df['hw_gpu_vram_gb'].max():.1f} GB (avg: {df['hw_gpu_vram_gb'].mean():.1f})") + + # Model distribution + print(f"\nModel Distribution:") + model_counts = df['model_id'].value_counts().head(10) + for model, count in model_counts.items(): + print(f" {model}: {count} benchmarks") + + # Performance stats for labeled models + best_models = df[df['label_best'] == 1] + if len(best_models) > 0: + print(f"\nBest Model Performance:") + print(f"Tokens/sec: {best_models['tokens_per_second'].min():.1f}-{best_models['tokens_per_second'].max():.1f} (avg: {best_models['tokens_per_second'].mean():.1f})") + print(f"Response time: {best_models['response_time'].min():.1f}-{best_models['response_time'].max():.1f}s (avg: {best_models['response_time'].mean():.1f})") + + return analysis + + def save_processed_dataset(self, df: pd.DataFrame, filename: str = "benchmarks.parquet") -> str: + """Save the processed dataset.""" + output_path = self.output_dir / filename + df.to_parquet(output_path, index=False) + print(f"\n๐Ÿ’พ Processed dataset saved to: {output_path}") + return str(output_path) + + def process_all(self) -> Tuple[str, Dict]: + """Run the complete aggregation pipeline.""" + print("๐Ÿ—๏ธ Starting dataset aggregation pipeline...") + + # Load raw data + df = self.load_all_benchmarks() + + # Clean and preprocess + df = self.clean_and_preprocess(df) + + # Create labels + df = self.create_training_labels(df) + + # Prepare features + df = self.prepare_features(df) + + # Analyze + analysis = self.analyze_dataset(df) + + # Save + output_path = self.save_processed_dataset(df) + + return output_path, analysis + +def main(): + """Main entry point.""" + if len(sys.argv) > 1 and sys.argv[1] == "--help": + print("Dataset Aggregator for ML Model Training") + print("") + print("Combines benchmark data from multiple hosts and creates training labels.") + print("") + print("Usage:") + print(" python dataset_aggregator.py") + print("") + print("Input: data/raw/benchmark_*.parquet files") + print("Output: data/processed/benchmarks.parquet") + print("") + return + + try: + aggregator = DatasetAggregator() + output_path, analysis = aggregator.process_all() + + print("\n๐ŸŽ‰ Dataset aggregation complete!") + print(f"Output: {output_path}") + + # Check if dataset is suitable for training + if analysis['positive_labels'] < 10: + print("\nโš ๏ธ Warning: Very few positive labels. Collect more benchmark data.") + elif analysis['unique_hosts'] < 3: + print("\nโš ๏ธ Warning: Data from very few hosts. Collect from more diverse hardware.") + else: + print("\nโœ… Dataset looks good for training!") + print("Next step: python train_model.py") + + except Exception as e: + print(f"\nโŒ Error during aggregation: {e}") + if "--debug" in sys.argv: + import traceback + traceback.print_exc() + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/ml-model/python/train_model.py b/ml-model/python/train_model.py new file mode 100644 index 0000000..3e9d1f6 --- /dev/null +++ b/ml-model/python/train_model.py @@ -0,0 +1,469 @@ +#!/usr/bin/env python3 +""" +TabTransformer model trainer for Ollama model selection. +Trains a lightweight transformer model to predict the best model for given hardware. +""" + +import pandas as pd +import numpy as np +import tensorflow as tf +from tensorflow import keras +from sklearn.model_selection import train_test_split +from sklearn.preprocessing import LabelEncoder, StandardScaler +from sklearn.metrics import classification_report, roc_auc_score +import joblib +import json +import sys +from pathlib import Path +from typing import Dict, Tuple, List +import warnings +warnings.filterwarnings('ignore') + +# Set random seeds for reproducibility +tf.random.set_seed(42) +np.random.seed(42) + +class TabTransformer(keras.Model): + """Lightweight TabTransformer implementation for tabular data.""" + + def __init__(self, + categorical_features: List[str], + numerical_features: List[str], + categorical_vocab_sizes: Dict[str, int], + embed_dim: int = 32, + num_heads: int = 2, + num_layers: int = 2, + dropout_rate: float = 0.1, + **kwargs): + super().__init__(**kwargs) + + self.categorical_features = categorical_features + self.numerical_features = numerical_features + self.embed_dim = embed_dim + + # Embedding layers for categorical features + self.embeddings = {} + for feature in categorical_features: + vocab_size = categorical_vocab_sizes[feature] + self.embeddings[feature] = keras.layers.Embedding( + input_dim=vocab_size, + output_dim=embed_dim, + name=f"embed_{feature}" + ) + + # Transformer layers + self.transformer_blocks = [] + for i in range(num_layers): + transformer_block = self._create_transformer_block( + embed_dim, num_heads, dropout_rate, name=f"transformer_{i}" + ) + self.transformer_blocks.append(transformer_block) + + # Dense layers for numerical features + self.numerical_dense = keras.layers.Dense( + embed_dim, activation='relu', name='numerical_dense' + ) + self.numerical_dropout = keras.layers.Dropout(dropout_rate) + + # Final classification layers + self.global_pool = keras.layers.GlobalAveragePooling1D() + self.concat = keras.layers.Concatenate() + self.final_dense1 = keras.layers.Dense(64, activation='relu', name='final_dense1') + self.final_dropout = keras.layers.Dropout(dropout_rate) + self.final_dense2 = keras.layers.Dense(32, activation='relu', name='final_dense2') + self.output_layer = keras.layers.Dense(1, activation='sigmoid', name='output') + + def _create_transformer_block(self, embed_dim, num_heads, dropout_rate, name): + """Create a transformer block.""" + inputs = keras.layers.Input(shape=(None, embed_dim)) + + # Multi-head attention + attention = keras.layers.MultiHeadAttention( + num_heads=num_heads, key_dim=embed_dim // num_heads, dropout=dropout_rate + )(inputs, inputs) + + # Add & Norm + attention = keras.layers.Dropout(dropout_rate)(attention) + attention_output = keras.layers.Add()([inputs, attention]) + attention_output = keras.layers.LayerNormalization()(attention_output) + + # Feed Forward + ff = keras.layers.Dense(embed_dim * 2, activation='relu')(attention_output) + ff = keras.layers.Dropout(dropout_rate)(ff) + ff = keras.layers.Dense(embed_dim)(ff) + + # Add & Norm + ff = keras.layers.Dropout(dropout_rate)(ff) + output = keras.layers.Add()([attention_output, ff]) + output = keras.layers.LayerNormalization()(output) + + return keras.Model(inputs=inputs, outputs=output, name=name) + + def call(self, inputs, training=None): + # Process categorical features + categorical_embeddings = [] + for feature in self.categorical_features: + embedded = self.embeddings[feature](inputs[feature]) + categorical_embeddings.append(embedded) + + # Stack categorical embeddings + if categorical_embeddings: + categorical_tensor = tf.stack(categorical_embeddings, axis=1) + + # Apply transformer blocks + transformer_output = categorical_tensor + for transformer_block in self.transformer_blocks: + transformer_output = transformer_block(transformer_output, training=training) + + # Global pooling + categorical_features = self.global_pool(transformer_output) + else: + categorical_features = tf.zeros((tf.shape(inputs[self.numerical_features[0]])[0], self.embed_dim)) + + # Process numerical features + numerical_tensor = tf.concat([inputs[feature] for feature in self.numerical_features], axis=-1) + numerical_features = self.numerical_dense(numerical_tensor) + numerical_features = self.numerical_dropout(numerical_features, training=training) + + # Combine features + combined = self.concat([categorical_features, numerical_features]) + + # Final layers + x = self.final_dense1(combined) + x = self.final_dropout(x, training=training) + x = self.final_dense2(x) + output = self.output_layer(x) + + return output + +class ModelTrainer: + """Trains the TabTransformer model.""" + + def __init__(self, data_path: str = "data/processed/benchmarks.parquet"): + self.data_path = Path(data_path) + self.model_dir = Path("ml-model/trained") + self.model_dir.mkdir(parents=True, exist_ok=True) + + # Feature definitions + self.categorical_features = [ + 'model_id', 'gpu_model_normalized', 'hw_platform', + 'ram_tier', 'cpu_tier', 'vram_tier' + ] + + self.numerical_features = [ + 'model_size_numeric', 'hw_cpu_cores', 'hw_cpu_freq_max', + 'hw_total_ram_gb', 'hw_gpu_vram_gb' + ] + + # Preprocessing objects + self.label_encoders = {} + self.scaler = StandardScaler() + self.categorical_vocab_sizes = {} + + def load_and_prepare_data(self) -> Tuple[Dict, np.ndarray, Dict, np.ndarray]: + """Load data and prepare for training.""" + print("๐Ÿ“Š Loading training data...") + + if not self.data_path.exists(): + raise FileNotFoundError(f"Training data not found: {self.data_path}") + + df = pd.read_parquet(self.data_path) + print(f"Loaded {len(df)} training samples") + + # Check for minimum data requirements + positive_samples = (df['label_best'] == 1).sum() + if positive_samples < 10: + raise ValueError(f"Not enough positive samples: {positive_samples} (need at least 10)") + + if df['host_id'].nunique() < 2: + raise ValueError("Need data from at least 2 different hardware configurations") + + print(f"Positive samples: {positive_samples} ({positive_samples/len(df):.2%})") + print(f"Unique hardware configs: {df['host_id'].nunique()}") + print(f"Unique models: {df['model_id'].nunique()}") + + # Prepare features and target + X = {} + y = df['label_best'].values.astype(np.float32) + + # Process categorical features + for feature in self.categorical_features: + if feature not in df.columns: + print(f"Warning: {feature} not in data, using default") + X[feature] = np.zeros(len(df), dtype=np.int32) + self.categorical_vocab_sizes[feature] = 1 + continue + + # Encode categorical feature + encoder = LabelEncoder() + encoded = encoder.fit_transform(df[feature].astype(str)) + + self.label_encoders[feature] = encoder + self.categorical_vocab_sizes[feature] = len(encoder.classes_) + X[feature] = encoded.astype(np.int32) + + print(f" {feature}: {len(encoder.classes_)} categories") + + # Process numerical features + numerical_data = [] + for feature in self.numerical_features: + if feature not in df.columns: + print(f"Warning: {feature} not in data, using zeros") + numerical_data.append(np.zeros(len(df))) + else: + numerical_data.append(df[feature].values) + + # Scale numerical features + numerical_array = np.column_stack(numerical_data) + numerical_scaled = self.scaler.fit_transform(numerical_array) + + # Add numerical features to X + for i, feature in enumerate(self.numerical_features): + X[feature] = numerical_scaled[:, i:i+1].astype(np.float32) + + return X, y, self.categorical_vocab_sizes, numerical_scaled + + def create_model(self, categorical_vocab_sizes: Dict[str, int]) -> TabTransformer: + """Create the TabTransformer model.""" + print("๐Ÿ—๏ธ Creating TabTransformer model...") + + model = TabTransformer( + categorical_features=self.categorical_features, + numerical_features=self.numerical_features, + categorical_vocab_sizes=categorical_vocab_sizes, + embed_dim=32, + num_heads=2, + num_layers=2, + dropout_rate=0.1 + ) + + # Compile model + model.compile( + optimizer=keras.optimizers.Adam(learning_rate=0.001), + loss='binary_crossentropy', + metrics=['accuracy', 'precision', 'recall'] + ) + + return model + + def train_model(self, model: TabTransformer, X: Dict, y: np.ndarray) -> keras.callbacks.History: + """Train the model.""" + print("๐Ÿš€ Training model...") + + # Split data + X_train, X_val, y_train, y_val = {}, {}, None, None + + # Get indices for split + indices = np.arange(len(y)) + train_indices, val_indices = train_test_split( + indices, test_size=0.2, random_state=42, stratify=y + ) + + # Split each feature + for feature in X.keys(): + X_train[feature] = X[feature][train_indices] + X_val[feature] = X[feature][val_indices] + + y_train = y[train_indices] + y_val = y[val_indices] + + print(f"Training samples: {len(y_train)}, Validation samples: {len(y_val)}") + + # Callbacks + callbacks = [ + keras.callbacks.EarlyStopping( + monitor='val_loss', patience=10, restore_best_weights=True + ), + keras.callbacks.ReduceLROnPlateau( + monitor='val_loss', factor=0.5, patience=5, min_lr=1e-6 + ) + ] + + # Train model + history = model.fit( + X_train, y_train, + validation_data=(X_val, y_val), + epochs=100, + batch_size=32, + callbacks=callbacks, + verbose=1 + ) + + # Evaluate + train_pred = model.predict(X_train) + val_pred = model.predict(X_val) + + train_auc = roc_auc_score(y_train, train_pred) + val_auc = roc_auc_score(y_val, val_pred) + + print(f"\n๐Ÿ“Š Training Results:") + print(f"Training AUC: {train_auc:.4f}") + print(f"Validation AUC: {val_auc:.4f}") + + if val_auc < 0.80: + print("โš ๏ธ Warning: Low validation AUC. Consider collecting more data.") + elif val_auc >= 0.90: + print("๐ŸŽ‰ Excellent model performance!") + else: + print("โœ… Good model performance.") + + return history + + def save_model_and_preprocessors(self, model: TabTransformer): + """Save the model and preprocessing objects.""" + print("๐Ÿ’พ Saving model and preprocessors...") + + # Save Keras model + model_path = self.model_dir / "model.keras" + model.save(model_path) + print(f" Model saved: {model_path}") + + # Save preprocessing objects + scaler_path = self.model_dir / "scaler.joblib" + joblib.dump(self.scaler, scaler_path) + print(f" Scaler saved: {scaler_path}") + + encoders_path = self.model_dir / "label_encoders.joblib" + joblib.dump(self.label_encoders, encoders_path) + print(f" Encoders saved: {encoders_path}") + + # Save model metadata + metadata = { + 'categorical_features': self.categorical_features, + 'numerical_features': self.numerical_features, + 'categorical_vocab_sizes': self.categorical_vocab_sizes, + 'feature_count': len(self.categorical_features) + len(self.numerical_features), + 'model_version': '1.0' + } + + metadata_path = self.model_dir / "metadata.json" + with open(metadata_path, 'w') as f: + json.dump(metadata, f, indent=2) + print(f" Metadata saved: {metadata_path}") + + def convert_to_onnx(self): + """Convert the model to ONNX format.""" + print("๐Ÿ”„ Converting to ONNX...") + + try: + import tf2onnx + import onnx + from onnxconverter_common import float16 + + # Load the saved model + model_path = self.model_dir / "model.keras" + model = keras.models.load_model(model_path, custom_objects={'TabTransformer': TabTransformer}) + + # Create dummy input for conversion + dummy_inputs = {} + for feature in self.categorical_features: + dummy_inputs[feature] = tf.TensorSpec(shape=[None], dtype=tf.int32, name=feature) + for feature in self.numerical_features: + dummy_inputs[feature] = tf.TensorSpec(shape=[None, 1], dtype=tf.float32, name=feature) + + # Convert to ONNX + onnx_path = self.model_dir / "model.onnx" + + model_proto, _ = tf2onnx.convert.from_keras( + model, + input_signature=dummy_inputs, + output_path=str(onnx_path) + ) + + print(f" ONNX model saved: {onnx_path}") + + # Quantize to INT8 + self._quantize_onnx_model(onnx_path) + + except ImportError as e: + print(f" โš ๏ธ ONNX conversion skipped (missing dependencies): {e}") + except Exception as e: + print(f" โŒ ONNX conversion failed: {e}") + + def _quantize_onnx_model(self, onnx_path: Path): + """Quantize ONNX model to INT8.""" + try: + from onnxruntime.quantization import quantize_dynamic, QuantType + + quantized_path = self.model_dir / "model_quantized.onnx" + + quantize_dynamic( + str(onnx_path), + str(quantized_path), + weight_type=QuantType.QInt8 + ) + + # Check file sizes + original_size = onnx_path.stat().st_size / 1024 # KB + quantized_size = quantized_path.stat().st_size / 1024 # KB + + print(f" Original model: {original_size:.1f} KB") + print(f" Quantized model: {quantized_size:.1f} KB ({quantized_size/original_size:.1%} of original)") + + if quantized_size < 150: + print(" โœ… Model size target achieved (<150 KB)") + else: + print(" โš ๏ธ Model larger than target (150 KB)") + + except ImportError as e: + print(f" โš ๏ธ Quantization skipped (missing dependencies): {e}") + except Exception as e: + print(f" โŒ Quantization failed: {e}") + + def train_complete_pipeline(self): + """Run the complete training pipeline.""" + print("๐Ÿง  Starting ML Model Training Pipeline") + print("=" * 50) + + try: + # Load and prepare data + X, y, vocab_sizes, _ = self.load_and_prepare_data() + + # Create model + model = self.create_model(vocab_sizes) + + # Train model + history = self.train_model(model, X, y) + + # Save everything + self.save_model_and_preprocessors(model) + + # Convert to ONNX + self.convert_to_onnx() + + print("\n๐ŸŽ‰ Training pipeline completed successfully!") + print(f"Model saved in: {self.model_dir}") + + return model, history + + except Exception as e: + print(f"\nโŒ Training failed: {e}") + if "--debug" in sys.argv: + import traceback + traceback.print_exc() + raise + +def main(): + """Main entry point.""" + if len(sys.argv) > 1 and sys.argv[1] == "--help": + print("TabTransformer Model Trainer") + print("") + print("Trains a lightweight transformer model to predict the best") + print("Ollama model for given hardware specifications.") + print("") + print("Requirements:") + print(" - data/processed/benchmarks.parquet (from dataset_aggregator.py)") + print(" - TensorFlow 2.16+") + print(" - scikit-learn, pandas, numpy") + print("") + print("Usage:") + print(" python train_model.py") + print(" python train_model.py --debug") + print("") + return + + trainer = ModelTrainer() + model, history = trainer.train_complete_pipeline() + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/ml-model/requirements.txt b/ml-model/requirements.txt new file mode 100644 index 0000000..f6d8fd4 --- /dev/null +++ b/ml-model/requirements.txt @@ -0,0 +1,10 @@ +pandas>=2.0.0 +scikit-learn>=1.3.0 +tensorflow>=2.16.0 +tf2onnx>=1.16.0 +onnxruntime-tools>=1.16.0 +joblib>=1.3.0 +tqdm>=4.65.0 +psutil>=5.9.0 +pyarrow>=12.0.0 +numpy>=1.24.0 \ No newline at end of file diff --git a/package.json b/package.json index e09eaab..364aa75 100644 --- a/package.json +++ b/package.json @@ -16,6 +16,9 @@ "update-db": "node bin/enhanced_cli.js update-db", "ollama": "node bin/enhanced_cli.js ollama", "list-models": "node bin/enhanced_cli.js list-models", + "ai-check": "node bin/enhanced_cli.js ai-check", + "benchmark": "cd ml-model && python python/benchmark_collector.py", + "train-ai": "cd ml-model && python python/train_model.py", "postinstall": "echo '๐ŸŽ‰ LLM Checker installed! Try: npm run check'" }, "dependencies": { @@ -73,6 +76,7 @@ "bin/", "src/", "analyzer/", + "ml-model/", "README.md", "LICENSE" ] diff --git a/src/ai/model-selector.js b/src/ai/model-selector.js new file mode 100644 index 0000000..502eda4 --- /dev/null +++ b/src/ai/model-selector.js @@ -0,0 +1,173 @@ +const path = require('path'); +const fs = require('fs'); + +class AIModelSelector { + constructor() { + this.aiSelectorPath = path.join(__dirname, '../../ml-model/js'); + this.isAvailable = this.checkAvailability(); + } + + checkAvailability() { + try { + const indexPath = path.join(this.aiSelectorPath, 'index.js'); + return fs.existsSync(indexPath); + } catch { + return false; + } + } + + async initialize() { + if (!this.isAvailable) { + throw new Error('AI Model Selector not available. Please train the model first.'); + } + + try { + const AISelector = require(path.join(this.aiSelectorPath, 'index.js')); + this.selector = new AISelector(); + await this.selector.initialize(); + return true; + } catch (error) { + throw new Error(`Failed to initialize AI selector: ${error.message}`); + } + } + + async selectBestModel(candidateModels, systemSpecs = null) { + if (!this.isAvailable) { + return this.fallbackSelection(candidateModels, systemSpecs); + } + + try { + if (!this.selector) { + await this.initialize(); + } + + const result = await this.selector.predictBestModel(candidateModels, systemSpecs); + + return { + bestModel: result.bestModel, + confidence: result.allPredictions[0]?.score || 0, + allPredictions: result.allPredictions, + method: 'ai', + systemSpecs: result.systemSpecs + }; + + } catch (error) { + console.warn(`AI selection failed: ${error.message}`); + return this.fallbackSelection(candidateModels, systemSpecs); + } + } + + fallbackSelection(candidateModels, systemSpecs = null) { + if (!systemSpecs) { + systemSpecs = { + total_ram_gb: 8, + gpu_vram_gb: 0, + cpu_cores: 4 + }; + } + + // Simple heuristic: pick largest model that fits in memory + const availableMemory = systemSpecs.gpu_vram_gb > 0 ? + systemSpecs.gpu_vram_gb : + systemSpecs.total_ram_gb * 0.7; + + const modelSizes = candidateModels.map(model => ({ + model, + size: this.estimateModelSize(model), + memoryReq: this.estimateModelSize(model) * 1.2 + })); + + const suitableModels = modelSizes + .filter(m => m.memoryReq <= availableMemory) + .sort((a, b) => b.size - a.size); + + const bestModel = suitableModels.length > 0 ? + suitableModels[0].model : + modelSizes.reduce((a, b) => a.size < b.size ? a : b).model; + + return { + bestModel, + confidence: suitableModels.length > 0 ? 0.8 : 0.5, + method: 'heuristic', + reason: suitableModels.length > 0 ? + 'Best fitting model for available memory' : + 'Smallest available model (safety fallback)', + systemSpecs + }; + } + + estimateModelSize(modelId) { + const sizeMatch = modelId.match(/(\d+\.?\d*)[kmb]/i); + if (sizeMatch) { + const num = parseFloat(sizeMatch[1]); + const unit = sizeMatch[0].slice(-1).toLowerCase(); + + if (unit === 'k') return num / 1000; + if (unit === 'm') return num / 1000; + if (unit === 'b') return num; + } + + // Default sizes for common model names + if (modelId.includes('mini')) return 3.8; + if (modelId.includes('7b')) return 7; + if (modelId.includes('13b')) return 13; + if (modelId.includes('70b')) return 70; + + return 7; // Safe default + } + + async benchmarkModel(modelId) { + // This would interface with the Python benchmarking script + const { spawn } = require('child_process'); + + return new Promise((resolve, reject) => { + const benchmark = spawn('python', [ + path.join(__dirname, '../../ml-model/python/benchmark_collector.py'), + '--single-model', modelId + ]); + + let output = ''; + let error = ''; + + benchmark.stdout.on('data', (data) => output += data); + benchmark.stderr.on('data', (data) => error += data); + + benchmark.on('close', (code) => { + if (code === 0) { + resolve({ success: true, data: output }); + } else { + reject(new Error(`Benchmark failed: ${error}`)); + } + }); + }); + } + + getTrainingStatus() { + const modelPath = path.join(__dirname, '../../ml-model/trained/model_quantized.onnx'); + const metadataPath = path.join(__dirname, '../../ml-model/trained/metadata.json'); + + const hasModel = fs.existsSync(modelPath); + const hasMetadata = fs.existsSync(metadataPath); + + if (hasModel && hasMetadata) { + try { + const metadata = JSON.parse(fs.readFileSync(metadataPath, 'utf8')); + const modelStats = fs.statSync(modelPath); + + return { + status: 'trained', + modelSize: Math.round(modelStats.size / 1024), // KB + version: metadata.model_version || '1.0', + features: metadata.feature_count || 0, + lastUpdated: modelStats.mtime.toISOString() + }; + } catch { + return { status: 'corrupted' }; + } + } + + return { status: 'not_trained' }; + } +} + +module.exports = AIModelSelector; \ No newline at end of file