Refuse a sampler application that excludes every candidate

.github/workflows/unit-tests.yml

@@ -57,7 +57,7 @@ jobs:
 
       - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830
         with:
-          path: target/cppcheck-src
+          path: lint-tools/cppcheck-src
           key: cppcheck-${{ runner.os }}-${{ hashFiles('Makefile') }}
 
       - uses: ./.github/actions/install-rust-toolchain
@@ -75,6 +75,6 @@ jobs:
 
       - uses: ./.github/actions/install-rust-toolchain
 
-      - run: make target/lint-venv/bin/clang-tidy
+      - run: make lint-tools/clang-tidy-venv/bin/clang-tidy
 
-      - run: make lint.cpp.clang-tidy CLANG_TIDY=target/lint-venv/bin/clang-tidy
+      - run: make lint.cpp.clang-tidy CLANG_TIDY=lint-tools/clang-tidy-venv/bin/clang-tidy

.gitignore

@@ -1,4 +1,5 @@
 .DS_Store
 **/*.rs.bk
+lint-tools/
 node_modules/
 target/

Makefile

@@ -8,7 +8,7 @@ CLANG_TIDY ?= clang-tidy
 
 COMPILE_COMMANDS = target/compile_commands.json
 
-CPPCHECK = target/cppcheck-src/cppcheck
+CPPCHECK = lint-tools/cppcheck-src/cppcheck
 
 CPPCHECK_COMMIT = f726f98ed2c3277780ee133ec5338f6352d7b43b
 
@@ -33,6 +33,19 @@ LLAMA_CPP_AND_GSL_SUPPRESSIONS = \
 	--suppress='*:*llama-cpp-bindings-sys/llama.cpp/*' \
 	--suppress='*:*llama-cpp-bindings-sys/GSL/*'
 
+lint-tools/clang-tidy-venv/bin/clang-tidy: requirements-lint.txt
+	python3 -m venv lint-tools/clang-tidy-venv
+	lint-tools/clang-tidy-venv/bin/pip install --quiet --requirement requirements-lint.txt
+
+lint-tools/cppcheck-src/Makefile:
+	git init --quiet lint-tools/cppcheck-src
+	git -C lint-tools/cppcheck-src fetch --quiet --depth 1 $(CPPCHECK_SOURCE) $(CPPCHECK_COMMIT)
+	git -C lint-tools/cppcheck-src checkout --quiet FETCH_HEAD
+
+$(CPPCHECK): lint-tools/cppcheck-src/Makefile
+	$(MAKE) --directory=lint-tools/cppcheck-src cppcheck \
+		FILESDIR=$(CURDIR)/lint-tools/cppcheck-src
+
 node_modules: package-lock.json
 	npm ci
 	touch node_modules
@@ -43,19 +56,6 @@ package-lock.json: package.json
 $(COMPILE_COMMANDS): $(WRAPPER_SOURCES_CRATE_FILES)
 	$(EMIT_WRAPPER_BUILD_INPUTS)
 
-target/cppcheck-src/Makefile:
-	git init --quiet target/cppcheck-src
-	git -C target/cppcheck-src fetch --quiet --depth 1 $(CPPCHECK_SOURCE) $(CPPCHECK_COMMIT)
-	git -C target/cppcheck-src checkout --quiet FETCH_HEAD
-
-$(CPPCHECK): target/cppcheck-src/Makefile
-	$(MAKE) --directory=target/cppcheck-src cppcheck \
-		FILESDIR=$(CURDIR)/target/cppcheck-src
-
-target/lint-venv/bin/clang-tidy: requirements-lint.txt
-	python3 -m venv target/lint-venv
-	target/lint-venv/bin/pip install --quiet --requirement requirements-lint.txt
-
 $(WRAPPER_SOURCES_RESPONSE_FILE): $(WRAPPER_SOURCES_CRATE_FILES)
 	$(EMIT_WRAPPER_BUILD_INPUTS)
 

llama-cpp-bindings/src/error/sampler_apply_error.rs

@@ -6,6 +6,8 @@ pub enum SamplerApplyError {
     FfiContract(#[from] crate::FfiContractError),
     #[error("the sampler pointer was null when applying to the token data array")]
     NullSampler,
+    #[error("the sampler excluded every candidate, leaving no token to select")]
+    EveryCandidateExcluded,
     #[error("the sampler ran out of memory while applying to the token data array")]
     NotEnoughMemory,
     #[error("the llama.cpp sampler ran out of memory")]

llama-cpp-bindings/src/sampling.rs

@@ -291,12 +291,16 @@ impl LlamaSampler {
 
     /// # Errors
     ///
-    /// Returns [`SampleError`] if the grammar sampler callback recorded a failure during application.
+    /// Returns [`SampleError::GrammarCallbackFailed`] if the grammar sampler callback recorded a
+    /// failure during application, which takes precedence over any other failure, or
+    /// [`SampleError::SamplerApply`] if applying the sampler fails, including when it excludes
+    /// every candidate.
     pub fn apply(&self, data_array: &mut LlamaTokenDataArray) -> Result<(), SampleError> {
         let scope = ErrorScope::enter();
-        data_array.apply_sampler(self)?;
+        let applied = data_array.apply_sampler(self);
+        grammar_callback_error_to_result(scope.take())?;
 
-        grammar_callback_error_to_result(scope.take())
+        Ok(applied?)
     }
 
     /// # Errors
@@ -722,6 +726,7 @@ mod tests {
     use std::ffi::CString;
 
     use llama_cpp_error_recorder::RecordedError;
+    use llama_cpp_error_recorder::record;
 
     use super::LlamaSampler;
     use super::grammar_callback_error_to_accept_result;
@@ -739,6 +744,70 @@ mod tests {
             .expect("penalties sampler must initialize")
     }
 
+    unsafe extern "C" fn record_a_failure_and_exclude_every_candidate(
+        _sampler: *mut llama_cpp_bindings_sys::llama_sampler,
+        candidates: *mut llama_cpp_bindings_sys::llama_token_data_array,
+    ) {
+        record(RecordedError::new("mask failed".to_string()));
+
+        let candidates = unsafe { &mut *candidates };
+        let candidate_data =
+            unsafe { std::slice::from_raw_parts_mut(candidates.data, candidates.size) };
+        for candidate in candidate_data {
+            candidate.logit = f32::NEG_INFINITY;
+        }
+    }
+
+    static mut FAILING_EXCLUDING_SAMPLER_I: llama_cpp_bindings_sys::llama_sampler_i =
+        llama_cpp_bindings_sys::llama_sampler_i {
+            name: None,
+            accept: None,
+            apply: Some(record_a_failure_and_exclude_every_candidate),
+            reset: None,
+            clone: None,
+            free: None,
+            backend_init: None,
+            backend_accept: None,
+            backend_apply: None,
+            backend_set_input: None,
+            backend_reset: None,
+            copy_state: None,
+        };
+
+    fn failing_excluding_sampler() -> LlamaSampler {
+        let sampler = unsafe {
+            llama_cpp_bindings_sys::llama_sampler_init(
+                &raw mut FAILING_EXCLUDING_SAMPLER_I,
+                std::ptr::null_mut(),
+            )
+        };
+
+        LlamaSampler::from_raw(sampler, "failing excluding sampler")
+            .expect("the failing excluding sampler must initialize")
+    }
+
+    #[test]
+    fn a_grammar_callback_failure_takes_precedence_over_a_failed_application() {
+        use crate::token::LlamaToken;
+        use crate::token::data::LlamaTokenData;
+        use crate::token::data_array::LlamaTokenDataArray;
+
+        let mut data_array = LlamaTokenDataArray::new(
+            vec![
+                LlamaTokenData::new(LlamaToken::new(0), 1.0, 0.0),
+                LlamaTokenData::new(LlamaToken::new(1), 2.0, 0.0),
+            ],
+            false,
+        );
+
+        assert_eq!(
+            failing_excluding_sampler().apply(&mut data_array),
+            Err(SampleError::GrammarCallbackFailed {
+                message: "mask failed".to_string()
+            })
+        );
+    }
+
     #[test]
     fn null_native_sampler_is_initialization_error() {
         let result = LlamaSampler::from_raw(std::ptr::null_mut(), "test");
@@ -880,6 +949,38 @@ mod tests {
         assert_eq!(data_array.selected_token(), Some(LlamaToken::new(1)));
     }
 
+    #[test]
+    fn apply_surfaces_a_sampler_that_excludes_every_candidate() {
+        use crate::error::sampler_apply_error::SamplerApplyError;
+        use crate::token::LlamaToken;
+        use crate::token::data::LlamaTokenData;
+        use crate::token::data_array::LlamaTokenDataArray;
+        use crate::token::logit_bias::LlamaLogitBias;
+
+        let excluding_sampler = LlamaSampler::logit_bias(
+            2,
+            &[
+                LlamaLogitBias::new(LlamaToken::new(0), f32::NEG_INFINITY),
+                LlamaLogitBias::new(LlamaToken::new(1), f32::NEG_INFINITY),
+            ],
+        )
+        .expect("the logit bias sampler must initialize");
+        let mut data_array = LlamaTokenDataArray::new(
+            vec![
+                LlamaTokenData::new(LlamaToken::new(0), 1.0, 0.0),
+                LlamaTokenData::new(LlamaToken::new(1), 2.0, 0.0),
+            ],
+            false,
+        );
+
+        assert_eq!(
+            excluding_sampler.apply(&mut data_array),
+            Err(SampleError::SamplerApply(
+                SamplerApplyError::EveryCandidateExcluded
+            ))
+        );
+    }
+
     #[test]
     fn accept_succeeds() {
         let mut sampler = LlamaSampler::chain_simple([penalties_sampler(), greedy_sampler()])

llama-cpp-bindings/src/token/data_array.rs

@@ -160,9 +160,10 @@ impl LlamaTokenDataArray {
     /// # Errors
     ///
     /// Returns [`SamplerApplyError`] if the sampler pointer is null, the llama.cpp
-    /// sampler runs out of memory, or it throws a C++ exception while applying.
+    /// sampler runs out of memory, it throws a C++ exception while applying, or it
+    /// excludes every candidate so that no token is left to select.
     pub fn apply_sampler(&mut self, sampler: &LlamaSampler) -> Result<(), SamplerApplyError> {
-        unsafe {
+        let applied = unsafe {
             self.modify_as_c_llama_token_data_array(|c_llama_token_data_array| {
                 let mut out_error: *mut std::os::raw::c_char = ptr::null_mut();
                 let status = llama_cpp_bindings_sys::llama_rs_sampler_apply(
@@ -172,6 +173,20 @@ impl LlamaTokenDataArray {
                 );
                 sampler_apply_status_to_result(status, out_error)
             })?
+        };
+
+        applied.and_then(|()| self.ensure_a_candidate_remains())
+    }
+
+    fn ensure_a_candidate_remains(&self) -> Result<(), SamplerApplyError> {
+        if self
+            .data
+            .iter()
+            .any(|candidate| candidate.logit() > f32::NEG_INFINITY)
+        {
+            Ok(())
+        } else {
+            Err(SamplerApplyError::EveryCandidateExcluded)
         }
     }
 
@@ -251,6 +266,33 @@ mod tests {
         );
     }
 
+    #[test]
+    fn a_sampler_excluding_every_candidate_leaves_nothing_to_select() {
+        use crate::sampling::LlamaSampler;
+        use crate::token::logit_bias::LlamaLogitBias;
+
+        let mut array = LlamaTokenDataArray::new(
+            vec![
+                LlamaTokenData::new(LlamaToken::new(0), 1.0, 0.0),
+                LlamaTokenData::new(LlamaToken::new(1), 2.0, 0.0),
+            ],
+            false,
+        );
+        let excluding_sampler = LlamaSampler::logit_bias(
+            2,
+            &[
+                LlamaLogitBias::new(LlamaToken::new(0), f32::NEG_INFINITY),
+                LlamaLogitBias::new(LlamaToken::new(1), f32::NEG_INFINITY),
+            ],
+        )
+        .expect("the logit bias sampler must initialize");
+
+        assert_eq!(
+            array.apply_sampler(&excluding_sampler),
+            Err(SamplerApplyError::EveryCandidateExcluded)
+        );
+    }
+
     #[test]
     fn apply_greedy_sampler_selects_highest_logit() {
         use crate::sampling::LlamaSampler;