1use pyo3::exceptions::PyValueError;
21use pyo3::prelude::*;
22use pyo3::types::{PyDict, PyList};
23
24use xlog_core::RelId;
25use xlog_cuda::CudaBuffer;
26use xlog_induce::{
27 induce_exact as induce_exact_engine, ExactInductionConfig, ExactInductionResult,
28 InduceExactRequest,
29};
30
31use super::{dlpack_from_py, types, CompiledIlpProgram};
32
33#[pymethods]
34impl CompiledIlpProgram {
35 #[pyo3(signature = (
59 head_relation,
60 candidate_relations,
61 positive_arg0,
62 positive_arg1,
63 negative_arg0 = None,
64 negative_arg1 = None,
65 k_per_topology = 2,
66 deterministic = true,
67 ))]
68 #[expect(
69 clippy::too_many_arguments,
70 reason = "the Python binding exposes each exact-induction input as a named argument"
71 )]
72 pub fn induce_exact_native<'py>(
73 &mut self,
74 py: Python<'py>,
75 head_relation: String,
76 candidate_relations: Vec<String>,
77 positive_arg0: &Bound<'py, PyAny>,
78 positive_arg1: &Bound<'py, PyAny>,
79 negative_arg0: Option<&Bound<'py, PyAny>>,
80 negative_arg1: Option<&Bound<'py, PyAny>>,
81 k_per_topology: u32,
82 deterministic: bool,
83 ) -> PyResult<Py<PyAny>> {
84 let head_rid = match self.rel_index.iter().find(|(_, n)| n == &head_relation) {
88 Some((rid, _)) => *rid,
89 None => return empty_result_dict(py),
90 };
91
92 let filtered_candidates: Vec<(RelId, String)> = candidate_relations
96 .into_iter()
97 .filter_map(|name| {
98 self.rel_index
99 .iter()
100 .find(|(_, n)| n == &name)
101 .map(|(rid, _)| (*rid, name))
102 })
103 .collect();
104
105 let head_schema = self.schemas.get(&head_relation).cloned().ok_or_else(|| {
109 PyValueError::new_err(format!(
110 "induce_exact_native: head_relation {:?} has no schema",
111 head_relation,
112 ))
113 })?;
114
115 let pos_dmt0 = dlpack_from_py(positive_arg0)?;
116 let pos_dmt1 = dlpack_from_py(positive_arg1)?;
117 let pos_buf = self
118 .provider
119 .from_dlpack_tensors_with_schema(head_schema.clone(), vec![pos_dmt0, pos_dmt1])
120 .map_err(types::xlog_err)?;
121
122 let neg_buf = match (negative_arg0, negative_arg1) {
123 (Some(a0), Some(a1)) => {
124 let dmt0 = dlpack_from_py(a0)?;
125 let dmt1 = dlpack_from_py(a1)?;
126 Some(
127 self.provider
128 .from_dlpack_tensors_with_schema(head_schema.clone(), vec![dmt0, dmt1])
129 .map_err(types::xlog_err)?,
130 )
131 }
132 (None, None) => None,
133 _ => {
134 return Err(PyValueError::new_err(
135 "induce_exact_native: negative_arg0 and negative_arg1 must both be set or both None",
136 ));
137 }
138 };
139
140 let store = self.executor.store();
146 let candidates_with_bufs: Vec<(RelId, &CudaBuffer)> = filtered_candidates
147 .iter()
148 .filter_map(|(rid, name)| store.get(name).map(|buf| (*rid, buf)))
149 .collect();
150
151 let request = InduceExactRequest {
153 head_rel_idx: head_rid,
154 candidates: &candidates_with_bufs,
155 positives: &pos_buf,
156 negatives: neg_buf.as_ref(),
157 config: ExactInductionConfig {
158 k_per_topology,
159 deterministic,
160 },
161 };
162 let result = induce_exact_engine(&self.provider, &request).map_err(types::xlog_err)?;
163
164 result_to_py_dict(py, &result, &head_relation, &self.rel_index)
167 }
168}
169
170fn empty_result_dict(py: Python<'_>) -> PyResult<Py<PyAny>> {
171 let d = PyDict::new(py);
172 let empty_candidates = PyList::empty(py);
173 d.set_item("candidates", empty_candidates)?;
174 d.set_item("total_scored", 0u32)?;
175 d.set_item("candidate_count", 0u32)?;
176 d.set_item("positive_count", 0u32)?;
177 d.set_item("negative_count", 0u32)?;
178 Ok(d.into())
179}
180
181fn result_to_py_dict(
182 py: Python<'_>,
183 result: &ExactInductionResult,
184 head_relation: &str,
185 rel_index: &[(RelId, String)],
186) -> PyResult<Py<PyAny>> {
187 let name_of = |rid: RelId| -> PyResult<&str> {
188 rel_index
189 .iter()
190 .find(|(r, _)| *r == rid)
191 .map(|(_, n)| n.as_str())
192 .ok_or_else(|| {
193 PyValueError::new_err(format!(
194 "induce_exact_native: result references unknown RelId {:?}",
195 rid,
196 ))
197 })
198 };
199
200 let candidates_list = PyList::empty(py);
201 for c in &result.candidates {
202 let entry = PyDict::new(py);
203 entry.set_item("topology", c.topology.as_str())?;
204 entry.set_item("head_relation", head_relation)?;
205 entry.set_item("left_relation", name_of(c.left_rel_idx)?)?;
206 entry.set_item("right_relation", name_of(c.right_rel_idx)?)?;
207 entry.set_item("positives_covered", c.positives_covered)?;
208 entry.set_item("negatives_covered", c.negatives_covered)?;
209 entry.set_item("local_rank", c.local_rank)?;
210 entry.set_item("next_positives_covered", c.next_positives_covered)?;
211 entry.set_item("next_negatives_covered", c.next_negatives_covered)?;
212 entry.set_item("tie_class_size", c.tie_class_size)?;
213 candidates_list.append(entry)?;
214 }
215
216 let d = PyDict::new(py);
217 d.set_item("candidates", candidates_list)?;
218 d.set_item("total_scored", result.total_scored)?;
219 d.set_item("candidate_count", result.candidate_count)?;
220 d.set_item("positive_count", result.positive_count)?;
221 d.set_item("negative_count", result.negative_count)?;
222 Ok(d.into())
223}