use scirs2_core::numeric::Float;
use std::fmt::Debug;
use crate::error::Result;
use super::predictor::{
PredictionCache, PredictionFeatures, PredictorFitReport, PredictorNetwork, PredictorSample,
PREDICTION_FEATURE_COUNT,
};
use super::types::{
ArchitectureAdaptation, ArchitectureChange, LandscapeAnalysis, OptimizationStrategy,
PerformancePrediction, UncertaintyMethod,
};
use super::AdaptiveConfig;
#[derive(Debug)]
pub struct TransformerPerformancePredictor<
T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static,
> {
predictor_network: PredictorNetwork<T>,
prediction_cache: PredictionCache,
uncertainty_method: UncertaintyMethod,
prediction_horizon: usize,
}
impl<T: Float + Debug + scirs2_core::ndarray::ScalarOperand + Send + Sync + 'static>
TransformerPerformancePredictor<T>
{
pub fn new(config: &AdaptiveConfig<T>) -> Result<Self> {
let horizon = config.prediction_horizon.max(1);
let hidden = (2 * PREDICTION_FEATURE_COUNT).max(horizon).min(128);
let feature_width = PREDICTION_FEATURE_COUNT.max(horizon.min(64));
let ridge = {
let lr = config.adaptation_lr.to_f64().unwrap_or(0.0).abs();
let lambda = if lr > 0.0 { lr * 1e-2 } else { 1e-4 };
scirs2_core::numeric::NumCast::from(lambda.clamp(1e-8, 1.0)).unwrap_or_else(|| T::one())
};
Ok(Self {
predictor_network: PredictorNetwork::new(
vec![PREDICTION_FEATURE_COUNT, hidden, feature_width],
ridge,
)?,
prediction_cache: PredictionCache::new(8 * horizon),
uncertainty_method: UncertaintyMethod::Ensemble,
prediction_horizon: horizon,
})
}
pub fn train(&mut self, samples: &[PredictorSample]) -> Result<PredictorFitReport> {
self.prediction_cache.clear();
self.predictor_network.fit(samples)
}
pub fn is_trained(&self) -> bool {
self.predictor_network.is_trained()
}
pub fn cache_hit_rate(&self) -> f64 {
self.prediction_cache.hit_rate()
}
pub fn uncertainty_method(&self) -> UncertaintyMethod {
self.uncertainty_method
}
pub fn prediction_horizon(&self) -> usize {
self.prediction_horizon
}
pub fn extract_features(
landscape: &LandscapeAnalysis<T>,
adaptation: &ArchitectureAdaptation<T>,
) -> PredictionFeatures {
let complexity = landscape.complexity.to_f64().unwrap_or(0.0);
let difficulty = landscape.difficulty.to_f64().unwrap_or(0.0);
let mut conservative = 0.0;
let mut aggressive = 0.0;
let mut exploratory = 0.0;
for strategy in &landscape.recommended_strategies {
match strategy {
OptimizationStrategy::Conservative | OptimizationStrategy::Exploitative => {
conservative += 1.0
}
OptimizationStrategy::Aggressive => aggressive += 1.0,
OptimizationStrategy::Exploratory => exploratory += 1.0,
OptimizationStrategy::Adaptive => {
conservative += 0.5;
aggressive += 0.5;
}
}
}
let strategy_total = (conservative + aggressive + exploratory).max(1.0);
let mut layers = adaptation.adapted_config.num_transformer_layers as f64;
let mut hidden = adaptation.adapted_config.model_dimension as f64;
let mut heads = adaptation.adapted_config.num_attention_heads as f64;
let mut dropout = adaptation.adapted_config.dropout_rate;
for change in &adaptation.changes {
match change {
ArchitectureChange::LayerCountChange(n) => layers = *n as f64,
ArchitectureChange::HiddenSizeChange(n) => hidden = *n as f64,
ArchitectureChange::AttentionHeadChange(n) => heads = *n as f64,
ArchitectureChange::DropoutChange(d) => dropout = *d,
ArchitectureChange::ActivationChange(_) => {}
}
}
PredictionFeatures {
complexity,
difficulty,
landscape_confidence: landscape.confidence.to_f64().unwrap_or(0.0),
expected_improvement: adaptation.expected_improvement.to_f64().unwrap_or(0.0),
adaptation_confidence: adaptation.confidence.to_f64().unwrap_or(0.0),
change_count: (adaptation.changes.len() as f64 / 8.0).min(1.0),
layer_count: (layers / 24.0).min(1.0),
hidden_size: (hidden / 2048.0).min(1.0),
head_count: (heads / 32.0).min(1.0),
dropout: dropout.clamp(0.0, 1.0),
interaction: complexity * difficulty,
complexity_sq: complexity * complexity,
difficulty_sq: difficulty * difficulty,
strategy_conservative: conservative / strategy_total,
strategy_aggressive: aggressive / strategy_total,
strategy_exploratory: exploratory / strategy_total,
}
}
pub fn predict_improvement(
&mut self,
landscape: &LandscapeAnalysis<T>,
adaptation: &ArchitectureAdaptation<T>,
) -> Result<PerformancePrediction<T>> {
let features = Self::extract_features(landscape, adaptation);
let key = features.cache_key();
let output = match self.prediction_cache.get(&key) {
Some(cached) => cached,
None => {
let fresh = self.predictor_network.predict(&features)?;
self.prediction_cache.insert(key, fresh.clone());
fresh
}
};
let cast =
|v: f64| -> T { scirs2_core::numeric::NumCast::from(v).unwrap_or_else(|| T::zero()) };
Ok(PerformancePrediction {
convergence_improvement: cast(output.convergence_improvement),
final_performance: cast(output.final_performance),
confidence: cast(output.confidence),
uncertainty: cast(output.uncertainty),
})
}
}