callm 0.2.0

Run Generative AI models directly on your hardware
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

use super::ModelImpl;
use crate::{device::DeviceConfig, error::CallmError};
use candle_core::Tensor;
use candle_nn::VarBuilder;
use candle_transformers::models::llama::{Cache, Config, Llama as Model};
use std::path::Path;
use std::sync::Arc;

const USE_KV_CACHE: bool = true;

pub struct ModelLlama {
    model: Model,
    cache: Cache,
    config: Config,
    device: Arc<DeviceConfig>,
}

impl ModelLlama {
    pub fn from_paths<P: AsRef<Path>>(
        paths: &[P],
        config: &Config,
        device: Arc<DeviceConfig>,
    ) -> Result<Self, CallmError> {
        // NOTE: unsafe inherited from memmap2::MmapOptions
        let vb = unsafe {
            VarBuilder::from_mmaped_safetensors(
                paths,
                device.candle_dtype(),
                device.candle_device(),
            )?
        };

        Ok(Self {
            model: Model::load(vb, config)?,
            cache: Self::spawn_kv_cache(config, &device)?,
            config: config.clone(),
            device,
        })
    }

    fn spawn_kv_cache(config: &Config, device: &DeviceConfig) -> Result<Cache, CallmError> {
        let cache = Cache::new(
            USE_KV_CACHE,
            device.candle_dtype(),
            config,
            device.candle_device(),
        )?;

        Ok(cache)
    }
}

impl ModelImpl for ModelLlama {
    fn forward(&mut self, input: &Tensor, index_pos: usize) -> Result<Tensor, CallmError> {
        Ok(self.model.forward(input, index_pos, &mut self.cache)?)
    }

    fn clear_kv_cache(&mut self) -> Result<(), CallmError> {
        self.cache = Self::spawn_kv_cache(&self.config, &self.device)?;
        Ok(())
    }
}