Gemma inference with llama.cpp and logits processor

lib/xml_schema_validator/src/py_xml_logits_processor_core.rs

@@ -9,6 +9,7 @@ pub struct XmlLogitsProcessorCore {
     xml_schema_validator: XmlSchemaValidator,
     trie: char_trie::CharTrie,
     tokens: Vec<(u32, String)>,
+    invalid_tokens: Vec<u32>,
 }
 
 #[pymethods]
@@ -27,13 +28,20 @@ impl XmlLogitsProcessorCore {
             })
             .collect();
         let mut trie = char_trie::CharTrie::default();
+        let mut invalid_tokens = Vec::new();
         for (token_id, token_text) in tokens.iter() {
-            trie.insert(token_text, token_id.clone());
+            let valid = token_text.chars().all(|c| c.is_ascii_graphic() || c.is_ascii_whitespace());
+            if valid {
+                trie.insert(token_text, token_id.clone());
+            } else {
+                invalid_tokens.push(*token_id);
+            }
         }
         Ok(Self {
             xml_schema_validator,
             trie,
             tokens,
+            invalid_tokens,
         })
     }
 
@@ -47,7 +55,10 @@ impl XmlLogitsProcessorCore {
     fn get_invalid_tokens(&mut self, py: Python<'_>) -> PyResult<Vec<u32>> {
         Ok(self
             .trie
-            .find_invalid_tokens(&self.xml_schema_validator))
+            .find_invalid_tokens(&self.xml_schema_validator)
+            .into_iter()
+            .chain(self.invalid_tokens.iter().cloned())
+            .collect())
     }
 
     /// Check if the validator has reached the end of the XML

lib/xml_schema_validator/src/token/mod.rs

@@ -49,9 +49,6 @@ pub enum Token {
 
 impl Token {
     pub fn append(self, c: char) -> Vec<Token> {
-        if !c.is_ascii_graphic() && ! c.is_ascii_whitespace() {
-            return vec![];
-        }
         match self {
             Token::AttributeEquals(attribute_equals) => attribute_equals.append(c),
             Token::AttributeName(attribute_name) => attribute_name.append(c),