mirror of
https://github.com/signerless/llm-checker.git
synced 2026-10-19 09:54:45 +00:00
e40e215303
This commit adds a comprehensive AI-powered model selection system using machine learning: 🧠 **AI Model Selector Framework** - TabTransformer model architecture for hardware-aware model selection - Lightweight quantized ONNX model (<150KB) for fast inference - Comprehensive Python training pipeline with data collection and preprocessing - JavaScript runtime with ONNX integration for real-time model selection 🔬 **Machine Learning Pipeline** - benchmark_collector.py: Automated performance data collection across different hardware - dataset_aggregator.py: Data preprocessing and training label generation - train_model.py: TabTransformer training with quantization and ONNX export - Hardware feature engineering with categorical and numerical features 🎯 **Smart Selection Logic** - AI-powered predictions based on hardware specifications - Robust fallback heuristics when AI model unavailable - Confidence scoring and multi-model ranking - Hardware tier classification (CPU, RAM, VRAM) ⚡ **New CLI Command: ai-check** - `ai-check --status`: Check AI model training status - `ai-check --benchmark`: Collect performance benchmarks - `ai-check --train`: Train the AI model - `ai-check --models <list>`: Select from specific models - `ai-check --prompt <text>`: Run with selected model 🔧 **Technical Features** - Optional ONNX Runtime dependency with graceful fallback - System information detection with multiple fallback methods - Cross-platform compatibility (macOS, Linux, Windows) - Integration with existing Ollama client and hardware detection 📊 **Model Architecture** - 2-layer TabTransformer with 32-dim embeddings - Binary classification for optimal model selection - Features: model size, CPU cores/freq, RAM/VRAM, GPU model, platform - Target: >90% AUC, <150KB model size, <100ms inference This enables intelligent, hardware-optimized model selection that learns from real performance data across diverse system configurations. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
469 lines
No EOL
18 KiB
Python
469 lines
No EOL
18 KiB
Python
#!/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() |