mullama 0.3.0

Comprehensive Rust bindings for llama.cpp with memory-safe API and advanced features
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

use crate::{context::Context, error::MullamaError, sys};

/// Represents the complete state of a context that can be saved and restored
#[derive(Debug)]
pub struct Session {
    pub data: Vec<u8>,
}

impl Session {
    /// Create a session from a context's current state
    pub fn from_context(context: &Context) -> Result<Self, MullamaError> {
        // Get the required buffer size
        let size = unsafe { sys::llama_state_get_size(context.as_ptr()) };

        // Allocate buffer
        let mut data = vec![0u8; size];

        // Get the state data
        let written =
            unsafe { sys::llama_state_get_data(context.as_ptr(), data.as_mut_ptr(), size) };

        if written != size {
            return Err(MullamaError::SessionError(
                "Failed to get session data".to_string(),
            ));
        }

        // Truncate to actual size written
        data.truncate(written);

        Ok(Session { data })
    }

    /// Restore a session to a context
    pub fn restore_to_context(&self, context: &mut Context) -> Result<(), MullamaError> {
        let read = unsafe {
            sys::llama_state_set_data(context.as_ptr(), self.data.as_ptr(), self.data.len())
        };

        if read != self.data.len() {
            return Err(MullamaError::SessionError(
                "Failed to restore session data".to_string(),
            ));
        }

        Ok(())
    }

    /// Save session to a file
    pub fn save_to_file(&self, path: impl AsRef<std::path::Path>) -> Result<(), MullamaError> {
        std::fs::write(path, &self.data).map_err(MullamaError::IoError)
    }

    /// Load session from a file
    pub fn load_from_file(path: impl AsRef<std::path::Path>) -> Result<Self, MullamaError> {
        let data = std::fs::read(path).map_err(MullamaError::IoError)?;
        Ok(Session { data })
    }
}