Struct LlamaModelWrapper 

pub struct LlamaModelWrapper {
    pub model_dir: Option<PathBuf>,
    /* private fields */
}

Fields

model_dir: Option<PathBuf>

Implementations

impl LlamaModelWrapper

pub fn from_varbuilder( vb: VarBuilder<'_>, config: &ModelDefinition, device: CandleDevice, dtype: DType, ) -> Result<Self>

pub fn forward_prefill( &self, input_ids: &Tensor, cache_key: &str, ) -> Result<Tensor>

pub fn forward_decode( &self, token_id: &Tensor, pos: usize, cache_key: &str, ) -> Result<Tensor>

pub fn export_kv_cache( &self, cache_key: &str, ) -> Option<Vec<(Tensor, Tensor, usize, usize)>>

pub fn release_cache(&self, cache_key: &str)

pub fn config(&self) -> &Config

pub fn device(&self) -> &CandleDevice

pub fn candle_device(&self) -> &CandleDevice

pub fn dtype(&self) -> DType

pub fn set_model_dir(&mut self, dir: PathBuf)