diff --git a/crates/larql-python/src/vindex.rs b/crates/larql-python/src/vindex.rs index 371d503ef..b590355f6 100644 --- a/crates/larql-python/src/vindex.rs +++ b/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(&self, f: F) -> PyResult where - F: FnOnce(&crate::walk::InferState) -> PyResult, + F: FnOnce(&mut crate::walk::InferState) -> PyResult, { { 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 = 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 = 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> { self.with_walk_model(|infer_state| { - let weights = &infer_state.weights; + let weights = infer_state.inference.as_weights(); let encoding = self .tokenizer diff --git a/crates/larql-python/src/walk.rs b/crates/larql-python/src/walk.rs index 4fc125cda..77902dee8 100644 --- a/crates/larql-python/src/walk.rs +++ b/crates/larql-python/src/walk.rs @@ -283,22 +283,20 @@ fn load_mmap_weights(dir: &Path) -> Result<(ModelWeights, Vec), 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, + pub inference: larql_inference::InferenceWeights, } impl InferState { - pub fn load(dir: &Path) -> Result { - let (weights, mmaps) = load_mmap_weights(dir)?; - Ok(Self { - weights, - _mmaps: mmaps, - }) + pub fn load(dir: &Path, config: &larql_vindex::VindexConfig) -> Result { + 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 }) } }