From e1da61fd715793f209f52d4c23ea7919687c1108 Mon Sep 17 00:00:00 2001
From: Niels Geens <niels.geens@gmail.com>
Date: Wed, 9 Apr 2025 10:56:20 +0000
Subject: [PATCH] Fix import for xml_schema_validator

---
 lib/xml_schema_validator/pyproject.toml                       | 2 ++
 .../{ => python}/xml_schema_validator/__init__.py             | 4 ++--
 lib/xml_schema_validator/src/python.rs                        | 2 +-
 3 files changed, 5 insertions(+), 3 deletions(-)
 rename lib/xml_schema_validator/{ => python}/xml_schema_validator/__init__.py (95%)

diff --git a/lib/xml_schema_validator/pyproject.toml b/lib/xml_schema_validator/pyproject.toml
index 3026f43..a64b84b 100644
--- a/lib/xml_schema_validator/pyproject.toml
+++ b/lib/xml_schema_validator/pyproject.toml
@@ -13,4 +13,6 @@ requires = ["maturin>=1.0,<2.0"]
 build-backend = "maturin"
 
 [tool.maturin]
+python-source = "python"
+module-name = "xml_schema_validator._rs"
 features = ["pyo3/extension-module"]
diff --git a/lib/xml_schema_validator/xml_schema_validator/__init__.py b/lib/xml_schema_validator/python/xml_schema_validator/__init__.py
similarity index 95%
rename from lib/xml_schema_validator/xml_schema_validator/__init__.py
rename to lib/xml_schema_validator/python/xml_schema_validator/__init__.py
index 1586866..42182ca 100644
--- a/lib/xml_schema_validator/xml_schema_validator/__init__.py
+++ b/lib/xml_schema_validator/python/xml_schema_validator/__init__.py
@@ -2,7 +2,7 @@ import torch
 import sys
 from transformers import LogitsProcessor
 from transformers import AutoTokenizer
-import xml_schema_validator
+from xml_schema_validator._rs import XmlLogitsProcessorCore
 
 class XmlLogitsProcessor(LogitsProcessor):
     """
@@ -26,7 +26,7 @@ class XmlLogitsProcessor(LogitsProcessor):
         items = dict()
         for token, id in vocab.items():
             items[id] = token
-        self.core = xml_schema_validator.XmlLogitsProcessorCore(items, schema_text)
+        self.core = XmlLogitsProcessorCore(items, schema_text)
         self.prompt_length = None
         self.is_first_call = True
 
diff --git a/lib/xml_schema_validator/src/python.rs b/lib/xml_schema_validator/src/python.rs
index d8be3ea..dafeb47 100644
--- a/lib/xml_schema_validator/src/python.rs
+++ b/lib/xml_schema_validator/src/python.rs
@@ -2,7 +2,7 @@ use pyo3::prelude::*;
 
 /// Python module for XML Schema validation
 #[pymodule]
-fn xml_schema_validator(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
+fn _rs(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     crate::py_xml_logits_processor_core::register_module(py, m)?;
     crate::py_xml_schema_validator::register_module(py, m)?;
     Ok(())