diff --git a/lib/xml_schema_validator/src/py_xml_logits_processor_core.rs b/lib/xml_schema_validator/src/py_xml_logits_processor_core.rs index 1e46659..9107bbe 100644 --- a/lib/xml_schema_validator/src/py_xml_logits_processor_core.rs +++ b/lib/xml_schema_validator/src/py_xml_logits_processor_core.rs @@ -31,7 +31,12 @@ impl XmlLogitsProcessorCore { /// Get a list of tokens that would lead to invalid XML fn get_invalid_tokens(&mut self, py: Python<'_>) -> PyResult> { + let mut chars = [false; 127]; + for i in 0..chars.len() { + chars[i] = self.xml_schema_validator.append_char(i as u8 as char).is_ok(); + } Ok(self.tokens.iter() + .filter(|token| !chars[token.1.chars().next().unwrap() as usize]) .filter_map(|token| { if self.xml_schema_validator.clone().append(&token.1).is_ok() { None @@ -41,19 +46,6 @@ impl XmlLogitsProcessorCore { }) .collect::>()) } - - /// Get a list of tokens that would lead to valid XML - fn get_valid_tokens(&mut self, py: Python<'_>) -> PyResult> { - Ok(self.tokens.iter() - .filter_map(|token| { - if self.xml_schema_validator.clone().append(&token.1).is_ok() { - Some(token.0.clone_ref(py)) - } else { - None - } - }) - .collect::>()) - } /// Check if the validator has reached the end of the XML fn eof(&self) -> PyResult { @@ -81,31 +73,6 @@ impl XmlLogitsProcessorCore { Ok(false) } } - - /// Reset the validator to its initial state - fn reset(&mut self) -> PyResult<()> { - // Create a new validator with the same schema elements - let schema_text = r#" - - - - - - - - -"#; - - match crate::XmlSchemaValidator::new(schema_text) { - Ok(validator) => { - self.xml_schema_validator = validator; - Ok(()) - }, - Err(e) => Err(pyo3::exceptions::PyRuntimeError::new_err(format!( - "Failed to reset validator: {}", e - ))), - } - } } /// Register the class with the Python module @@ -185,43 +152,6 @@ mod tests { }); } - #[test] - fn test_get_valid_tokens() { - Python::with_gil(|py| { - let mut processor = create_test_processor(py).unwrap(); - - // At start, only "<" is valid - let valid_tokens = processor.get_valid_tokens(py).unwrap(); - - let contains_0 = valid_tokens.iter().any(|obj| { - let obj_id: Option = obj.extract(py).ok(); - obj_id == Some(0) - }); - - assert!(contains_0, "Token '<' should be valid at start"); - assert_eq!(valid_tokens.len(), 1, "Only one token should be valid at start"); - - // After , any character should be valid in mixed content - let mut processor = create_test_processor(py).unwrap(); - assert!(processor.append(""), "Should be able to append "); - - let valid_tokens = processor.get_valid_tokens(py).unwrap(); - - let contains_3 = valid_tokens.iter().any(|obj| { - let obj_id: Option = obj.extract(py).ok(); - obj_id == Some(3) - }); - - let contains_4 = valid_tokens.iter().any(|obj| { - let obj_id: Option = obj.extract(py).ok(); - obj_id == Some(4) - }); - - assert!(contains_3, "Token ' ' should be valid inside mixed content"); - assert!(contains_4, "Token 'a' should be valid inside mixed content"); - }); - } - #[test] fn test_eof_detection() { Python::with_gil(|py| { @@ -299,21 +229,19 @@ mod tests { "Should contain 'reasoning' element"); assert_eq!(names.len(), 1, "Should have exactly one element"); - // Test with real schema from disk if available - if let Ok(schema_text) = std::fs::read_to_string("../../action_schema.xsd") { - let tokens = PyDict::new(py); - tokens.set_item(0, "<").unwrap(); - - let processor = XmlLogitsProcessorCore::new(tokens.into(), &schema_text).unwrap(); - let names = processor.get_element_names().unwrap(); - - // Print out all elements from schema for debugging - println!("Elements in schema: {:?}", names); - - // Check if reasoning exists - assert!(names.contains(&"reasoning".to_string()), - "Real schema should contain 'reasoning'"); - } + let tokens = PyDict::new(py); + tokens.set_item(0, "<").unwrap(); + + let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap(); + let processor = XmlLogitsProcessorCore::new(tokens.into(), &schema_text).unwrap(); + let names = processor.get_element_names().unwrap(); + + // Print out all elements from schema for debugging + println!("Elements in schema: {:?}", names); + + // Check if reasoning exists + assert!(names.contains(&"reasoning".to_string()), + "Real schema should contain 'reasoning'"); }); } @@ -403,50 +331,4 @@ mod tests { } }); } - - #[test] - fn test_debug_token_processing() { - Python::with_gil(|py| { - if let Ok(schema_text) = std::fs::read_to_string("../../action_schema.xsd") { - // Create a dictionary mapping token IDs to their string values - // Simulating how the tokenizer vocabulary works - let tokens = PyDict::new(py); - tokens.set_item(1000, "<").unwrap(); - tokens.set_item(1001, "reasoning").unwrap(); - tokens.set_item(1002, ">").unwrap(); - tokens.set_item(1003, " ").unwrap(); - tokens.set_item(1004, "test").unwrap(); - tokens.set_item(1005, "").unwrap(); - - let mut processor = XmlLogitsProcessorCore::new(tokens.clone().into(), &schema_text).unwrap(); - - // Test which tokens are valid at the start - let invalid_tokens = processor.get_invalid_tokens(py).unwrap(); - let valid_tokens = processor.get_valid_tokens(py).unwrap(); - - println!("At start - valid tokens count: {}", valid_tokens.len()); - println!("At start - invalid tokens count: {}", invalid_tokens.len()); - - // Create a new processor and progress to after "" - let mut processor = XmlLogitsProcessorCore::new(tokens.into(), &schema_text).unwrap(); - let result = processor.append(""); - println!("Append '': {}", result); - - if result { - // Now get valid/invalid tokens after "" - let invalid_tokens = processor.get_invalid_tokens(py).unwrap(); - let valid_tokens = processor.get_valid_tokens(py).unwrap(); - - println!("After '' - valid tokens count: {}", valid_tokens.len()); - println!("After '' - invalid tokens count: {}", invalid_tokens.len()); - - // At this point we'd expect to have more valid tokens than at the start - // since mixed content allows any text - } - } else { - println!("Warning: Couldn't find schema file for testing"); - } - }); - } } \ No newline at end of file