Expose functions llama_load_session_file and llama_save_session_file

Author: zh217

llama-cpp-2/src/context.rs

@@ -15,6 +15,7 @@ use std::slice;
 pub mod kv_cache;
 pub mod params;
 pub mod sample;
+pub mod session;
 
 /// Safe wrapper around `llama_context`.
 #[allow(clippy::module_name_repetitions)]

llama-cpp-2/src/context/session.rs

@@ -0,0 +1,77 @@
+//! utilities for working with session files
+
+use std::ffi::{CString, NulError};
+use std::path::{Path, PathBuf};
+use crate::context::LlamaContext;
+use crate::token::LlamaToken;
+
+#[derive(Debug, Eq, PartialEq, thiserror::Error)]
+pub enum SaveSessionError {
+    #[error("Failed to save session file")]
+    FailedToSave,
+
+    #[error("null byte in string {0}")]
+    NullError(#[from] NulError),
+
+    #[error("failed to convert path {0} to str")]
+    PathToStrError(PathBuf),
+}
+
+#[derive(Debug, Eq, PartialEq, thiserror::Error)]
+pub enum LoadSessionError {
+    #[error("Failed to load session file")]
+    FailedToLoad,
+
+    #[error("null byte in string {0}")]
+    NullError(#[from] NulError),
+
+    #[error("failed to convert path {0} to str")]
+    PathToStrError(PathBuf),
+}
+
+impl LlamaContext<'_> {
+    pub fn save_session_file(&self, path_session: impl AsRef<Path>, tokens: &[LlamaToken]) -> Result<(), SaveSessionError> {
+        let path = path_session.as_ref();
+        let path = path
+            .to_str()
+            .ok_or(SaveSessionError::PathToStrError(path.to_path_buf()))?;
+
+        let cstr = CString::new(path)?;
+
+        if unsafe {
+            llama_cpp_sys_2::llama_save_session_file(
+                self.context.as_ptr(),
+                cstr.as_ptr(),
+                tokens.as_ptr() as *const i32,
+                tokens.len())
+        } {
+            Ok(())
+        } else {
+            Err(SaveSessionError::FailedToSave)
+        }
+    }
+    pub fn load_session_file(&mut self, path_session: impl AsRef<Path>, max_tokens: usize) -> Result<Vec<LlamaToken>, LoadSessionError> {
+        let path = path_session.as_ref();
+        let path = path
+            .to_str()
+            .ok_or(LoadSessionError::PathToStrError(path.to_path_buf()))?;
+
+        let cstr = CString::new(path)?;
+        let mut tokens = Vec::with_capacity(max_tokens);
+        let mut n_out = 0;
+
+        unsafe {
+            if llama_cpp_sys_2::llama_load_session_file(
+                self.context.as_ptr(),
+                cstr.as_ptr(),
+                tokens.as_mut_ptr() as *mut i32,
+                max_tokens,
+                &mut n_out) {
+                tokens.set_len(n_out);
+                Ok(tokens)
+            } else {
+                Err(LoadSessionError::FailedToLoad)
+            }
+        }
+    }
+}