fix(python): dispatch vindex infer through inference weights

crates/larql-python/src/vindex.rs

@@ -393,25 +393,25 @@ impl PyVindex {
         })
     }
 
-    /// Run a closure with a reference to the lazily-loaded walk FFN state.
+    /// Run a closure with a mutable reference to the lazily-loaded walk FFN state.
     /// Loads on first call; subsequent calls reuse the mmap'd weights.
     fn with_walk_model<F, R>(&self, f: F) -> PyResult<R>
     where
-        F: FnOnce(&crate::walk::InferState) -> PyResult<R>,
+        F: FnOnce(&mut crate::walk::InferState) -> PyResult<R>,
     {
         {
             let mut state = self.walk_model.borrow_mut();
             if state.is_none() {
                 let dir = std::path::Path::new(&self.path);
-                *state = Some(crate::walk::InferState::load(dir).map_err(|e| {
+                *state = Some(crate::walk::InferState::load(dir, &self.config).map_err(|e| {
                     pyo3::exceptions::PyRuntimeError::new_err(format!(
                         "Failed to load model weights: {e}"
                     ))
                 })?);
             }
         }
-        let state = self.walk_model.borrow();
-        f(state.as_ref().unwrap())
+        let mut state = self.walk_model.borrow_mut();
+        f(state.as_mut().unwrap())
     }
 
     /// Compute scaled embedding for entity text. Multi-token entities are averaged.
@@ -1201,8 +1201,7 @@ impl PyVindex {
                 .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
             let token_ids: Vec<u32> = encoding.get_ids().to_vec();
 
-            let result = larql_inference::infer_patched(
-                &infer_state.weights,
+            let result = infer_state.inference.infer_patched(
                 &self.tokenizer,
                 &self.index,
                 self.knn_store.as_ref(),
@@ -1303,7 +1302,7 @@ impl PyVindex {
                 normalise: false,
             };
             let result = larql_inference::forward::target_delta::optimise_target_delta(
-                &infer_state.weights,
+                infer_state.inference.as_weights(),
                 &prompt_ids,
                 target_id,
                 install_layer,
@@ -1351,8 +1350,7 @@ impl PyVindex {
                 .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
             let token_ids: Vec<u32> = encoding.get_ids().to_vec();
 
-            let result = larql_inference::infer_patched(
-                &infer_state.weights,
+            let result = infer_state.inference.infer_patched(
                 &self.tokenizer,
                 &self.index,
                 self.knn_store.as_ref(),
@@ -1385,7 +1383,7 @@ impl PyVindex {
         top_k: usize,
     ) -> PyResult<Vec<(usize, usize, f32, String)>> {
         self.with_walk_model(|infer_state| {
-            let weights = &infer_state.weights;
+            let weights = infer_state.inference.as_weights();
 
             let encoding = self
                 .tokenizer

crates/larql-python/src/walk.rs

@@ -283,22 +283,20 @@ fn load_mmap_weights(dir: &Path) -> Result<(ModelWeights, Vec<WeightMmap>), Stri
     Ok((weights, mmaps))
 }
 
-// ── InferState: lazy-loaded mmap'd weights for vindex.infer() ──
+// ── InferState: lazy-loaded inference weights for vindex.infer() ──
 
-/// Mmap'd model weights, reusable across infer() calls.
+/// Format-aware model weights, reusable across infer() calls.
 /// Created lazily on first infer(), held by PyVindex.
 pub struct InferState {
-    pub weights: ModelWeights,
-    _mmaps: Vec<WeightMmap>,
+    pub inference: larql_inference::InferenceWeights,
 }
 
 impl InferState {
-    pub fn load(dir: &Path) -> Result<Self, String> {
-        let (weights, mmaps) = load_mmap_weights(dir)?;
-        Ok(Self {
-            weights,
-            _mmaps: mmaps,
-        })
+    pub fn load(dir: &Path, config: &larql_vindex::VindexConfig) -> Result<Self, String> {
+        let mut cb = larql_vindex::SilentLoadCallbacks;
+        let inference = larql_inference::InferenceWeights::load(dir, config, &mut cb)
+            .map_err(|e| e.to_string())?;
+        Ok(Self { inference })
     }
 }