139 lines
4.0 KiB
Rust
139 lines
4.0 KiB
Rust
//! Configuration types for speculative decoding
|
|
|
|
use super::types::DraftModelType;
|
|
use crate::{InferenceError, InferenceResult};
|
|
|
|
/// Adaptive configuration for draft length optimization
|
|
#[derive(Debug, Clone)]
|
|
pub struct AdaptiveConfig {
|
|
/// Minimum draft length to try
|
|
pub min_length: usize,
|
|
/// Maximum draft length to try
|
|
pub max_length: usize,
|
|
/// Target acceptance rate for adaptation
|
|
pub target_acceptance_rate: f32,
|
|
/// How aggressively to adjust draft length
|
|
pub adjustment_factor: f32,
|
|
}
|
|
|
|
impl Default for AdaptiveConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
min_length: 1,
|
|
max_length: 8,
|
|
target_acceptance_rate: 0.8,
|
|
adjustment_factor: 0.1,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Configuration for speculative decoding
|
|
#[derive(Debug, Clone)]
|
|
pub struct SpeculativeConfig {
|
|
/// Type of draft model to use
|
|
pub draft_model_type: DraftModelType,
|
|
/// Maximum number of tokens to generate in draft
|
|
pub max_draft_tokens: usize,
|
|
/// Minimum confidence threshold for acceptance
|
|
pub acceptance_threshold: f32,
|
|
/// Number of additional lookahead tokens
|
|
pub lookahead_tokens: usize,
|
|
/// Whether to track performance metrics
|
|
pub performance_tracking: bool,
|
|
/// Adaptive draft length configuration
|
|
pub adaptive_config: Option<AdaptiveConfig>,
|
|
}
|
|
|
|
impl SpeculativeConfig {
|
|
/// Create a new configuration with sensible defaults
|
|
#[must_use]
|
|
pub fn new() -> Self {
|
|
Self {
|
|
draft_model_type: DraftModelType::SmallModel("default".to_string()),
|
|
max_draft_tokens: 4,
|
|
acceptance_threshold: 0.7,
|
|
lookahead_tokens: 1,
|
|
performance_tracking: false,
|
|
adaptive_config: None,
|
|
}
|
|
}
|
|
|
|
/// Set the draft model type
|
|
#[must_use]
|
|
pub fn with_draft_model(mut self, model_type: DraftModelType) -> Self {
|
|
self.draft_model_type = model_type;
|
|
self
|
|
}
|
|
|
|
/// Set maximum draft tokens
|
|
#[must_use]
|
|
pub fn with_max_draft_tokens(mut self, max_tokens: usize) -> Self {
|
|
self.max_draft_tokens = max_tokens;
|
|
self
|
|
}
|
|
|
|
/// Set acceptance threshold
|
|
#[must_use]
|
|
pub fn with_acceptance_threshold(mut self, threshold: f32) -> Self {
|
|
self.acceptance_threshold = threshold;
|
|
self
|
|
}
|
|
|
|
/// Set lookahead tokens
|
|
#[must_use]
|
|
pub fn with_lookahead_tokens(mut self, lookahead: usize) -> Self {
|
|
self.lookahead_tokens = lookahead;
|
|
self
|
|
}
|
|
|
|
/// Enable performance tracking
|
|
#[must_use]
|
|
pub fn with_performance_tracking(mut self, enabled: bool) -> Self {
|
|
self.performance_tracking = enabled;
|
|
self
|
|
}
|
|
|
|
/// Enable adaptive draft length
|
|
#[must_use]
|
|
pub fn with_adaptive_draft_length(mut self, config: AdaptiveConfig) -> Self {
|
|
self.adaptive_config = Some(config);
|
|
self
|
|
}
|
|
|
|
/// Validate configuration parameters
|
|
pub fn validate(&self) -> InferenceResult<()> {
|
|
if self.max_draft_tokens == 0 {
|
|
return Err(InferenceError::invalid_request(
|
|
"max_draft_tokens must be greater than 0",
|
|
));
|
|
}
|
|
|
|
if !(0.0..=1.0).contains(&self.acceptance_threshold) {
|
|
return Err(InferenceError::invalid_request(
|
|
"acceptance_threshold must be between 0.0 and 1.0",
|
|
));
|
|
}
|
|
|
|
if let Some(ref adaptive) = self.adaptive_config {
|
|
if adaptive.min_length == 0 || adaptive.min_length > adaptive.max_length {
|
|
return Err(InferenceError::invalid_request(
|
|
"adaptive config: min_length must be > 0 and <= max_length",
|
|
));
|
|
}
|
|
if !(0.0..=1.0).contains(&adaptive.target_acceptance_rate) {
|
|
return Err(InferenceError::invalid_request(
|
|
"adaptive config: target_acceptance_rate must be between 0.0 and 1.0",
|
|
));
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl Default for SpeculativeConfig {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|