Expose safe wrappers around the new sampling API types

Author: volesen

llama-cpp-2/src/lib.rs

@@ -26,6 +26,7 @@ pub mod context;
 pub mod llama_backend;
 pub mod llama_batch;
 pub mod model;
+pub mod sampling;
 pub mod timing;
 pub mod token;
 pub mod token_type;
@@ -61,6 +62,7 @@ pub enum LLamaCppError {
     /// see [`EmbeddingsError`]
     #[error(transparent)]
     EmbeddingError(#[from] EmbeddingsError),
+    // See [`LlamaSamplerError`]
 }
 
 /// There was an error while getting the chat template from a model.
@@ -193,6 +195,14 @@ pub enum LlamaLoraAdapterRemoveError {
     ErrorResult(i32),
 }
 
+/// An error that can occur when initializing a sampler.
+#[derive(Debug, Eq, PartialEq, thiserror::Error)]
+pub enum LlamaSamplerError {
+    /// llama.cpp returned null
+    #[error("null reference from llama.cpp")]
+    NullReturn,
+}
+
 /// get the time (in microseconds) according to llama.cpp
 /// ```
 /// # use llama_cpp_2::llama_time_us;

llama-cpp-2/src/sampling.rs

@@ -0,0 +1,39 @@
+//! Safe wrapper around `llama_sampler`.
+pub mod params;
+
+use std::fmt::{Debug, Formatter};
+use std::ptr::NonNull;
+
+use crate::LlamaSamplerError;
+
+/// A safe wrapper around `llama_sampler`.
+pub struct LlamaSampler {
+    pub(crate) sampler: NonNull<llama_cpp_sys_2::llama_sampler>,
+}
+
+impl Debug for LlamaSampler {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("LlamaSamplerChain").finish()
+    }
+}
+
+impl LlamaSampler {
+    pub fn new(params: params::LlamaSamplerChainParams) -> Result<Self, LlamaSamplerError> {
+        let sampler = unsafe {
+            NonNull::new(llama_cpp_sys_2::llama_sampler_chain_init(
+                params.sampler_chain_params,
+            ))
+            .ok_or(LlamaSamplerError::NullReturn)
+        }?;
+
+        Ok(Self { sampler })
+    }
+}
+
+impl Drop for LlamaSampler {
+    fn drop(&mut self) {
+        unsafe {
+            llama_cpp_sys_2::llama_sampler_free(self.sampler.as_ptr());
+        }
+    }
+}

llama-cpp-2/src/sampling/params.rs

@@ -0,0 +1,36 @@
+//! Safe wrapper around `llama_sampler`.
+
+use std::fmt::{Debug, Formatter};
+use std::ptr::NonNull;
+
+/// A safe wrapper around `llama_sampler`.
+pub struct LlamaSamplerChainParams {
+    pub(crate) sampler_chain_params: llama_cpp_sys_2::llama_sampler_chain_params,
+}
+
+impl Debug for LlamaSamplerChainParams {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("LlamaSamplerChainParams").finish()
+    }
+}
+
+impl Default for LlamaSamplerChainParams {
+    fn default() -> Self {
+        let sampler_chain_params = unsafe { llama_cpp_sys_2::llama_sampler_chain_default_params() };
+
+        Self {
+            sampler_chain_params,
+        }
+    }
+}
+
+impl LlamaSamplerChainParams {
+    pub fn with_no_perf(&mut self, no_perf: bool) -> &mut Self {
+        self.sampler_chain_params.no_perf = no_perf;
+        self
+    }
+
+    pub fn no_perf(&self) -> bool {
+        self.sampler_chain_params.no_perf
+    }
+}