diff --git a/Cargo.lock b/Cargo.lock index fe6468f..e085d96 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -103,26 +103,6 @@ version = "1.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "55248b47b0caf0546f7988906588779981c43bb1bc9d0c44087278f80cdb44ba" -[[package]] -name = "bincode" -version = "2.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "36eaf5d7b090263e8150820482d5d93cd964a81e4019913c972f4edcc6edb740" -dependencies = [ - "bincode_derive", - "serde", - "unty", -] - -[[package]] -name = "bincode_derive" -version = "2.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bf95709a440f45e986983918d0e8a1f30a9b1df04918fc828670606804ac3c09" -dependencies = [ - "virtue", -] - [[package]] name = "bitflags" version = "2.9.1" @@ -143,9 +123,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.10.1" +version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a" +checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" [[package]] name = "cc" @@ -236,8 +216,18 @@ version = "0.20.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" dependencies = [ - "darling_core", - "darling_macro", + "darling_core 0.20.11", + "darling_macro 0.20.11", +] + +[[package]] +name = "darling" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9cdf337090841a411e2a7f3deb9187445851f91b309c0c0a29e05f74a00a48c0" +dependencies = [ + "darling_core 0.21.3", + "darling_macro 0.21.3", ] [[package]] @@ -254,13 +244,38 @@ dependencies = [ "syn", ] +[[package]] +name = "darling_core" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1247195ecd7e3c85f83c8d2a366e4210d588e802133e1e355180a9870b517ea4" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn", +] + [[package]] name = "darling_macro" version = "0.20.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" dependencies = [ - "darling_core", + "darling_core 0.20.11", + "quote", + "syn", +] + +[[package]] +name = "darling_macro" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d38308df82d1080de0afee5d069fa14b0326a88c14f15c5ccda35b4a6c414c81" +dependencies = [ + "darling_core 0.21.3", "quote", "syn", ] @@ -290,7 +305,7 @@ version = "0.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" dependencies = [ - "darling", + "darling 0.20.11", "proc-macro2", "quote", "syn", @@ -561,7 +576,7 @@ dependencies = [ "noisy_float", "num-traits", "order-stat", - "thiserror 2.0.12", + "thiserror 2.0.18", ] [[package]] @@ -918,6 +933,12 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pastey" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b867cad97c0791bbd3aaa6472142568c6c9e8f71937e98379f584cfb0cf35bec" + [[package]] name = "pem-rfc7468" version = "0.7.0" @@ -971,18 +992,18 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.95" +version = "1.0.106" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02b3e5e68a3a1a02aad3ec490a98007cbc13c37cbe84a3cd7b8e406d76e7f778" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" dependencies = [ "unicode-ident", ] [[package]] name = "quote" -version = "1.0.40" +version = "1.0.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1885c039570dc00dcb4ff087a89e185fd56bae234ddc7f056a945bf36467248d" +checksum = "21b2ebcf727b7760c461f091f9f0f539b77b8e87f2fd88131e7f1b433b3cece4" dependencies = [ "proc-macro2", ] @@ -1271,9 +1292,9 @@ checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" [[package]] name = "syn" -version = "2.0.101" +version = "2.0.114" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ce2b7fc941b3a24138a0a7cf8e858bfc6a992e7978a068a5c760deb0ed43caf" +checksum = "d4d107df263a3013ef9b1879b0df87d706ff80f65a86ea879bd9c31f9b307c2a" dependencies = [ "proc-macro2", "quote", @@ -1318,11 +1339,11 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.12" +version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "567b8a2dae586314f7be2a752ec7474332959c6460e02bde30d702a66d488708" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" dependencies = [ - "thiserror-impl 2.0.12", + "thiserror-impl 2.0.18", ] [[package]] @@ -1338,9 +1359,9 @@ dependencies = [ [[package]] name = "thiserror-impl" -version = "2.0.12" +version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f7cf42b4507d8ea322120659672cf1b9dbb93f8f2d4ecfd6e51350ff5b17a1d" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" dependencies = [ "proc-macro2", "quote", @@ -1372,7 +1393,7 @@ dependencies = [ "serde", "serde_json", "spm_precompiled", - "thiserror 2.0.12", + "thiserror 2.0.18", "unicode-normalization-alignments", "unicode-segmentation", "unicode_categories", @@ -1424,12 +1445,6 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" -[[package]] -name = "unty" -version = "0.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6d49784317cd0d1ee7ec5c716dd598ec5b4483ea832a2dced265471cc0f690ae" - [[package]] name = "ureq" version = "3.1.4" @@ -1485,9 +1500,8 @@ dependencies = [ [[package]] name = "valentinus" -version = "1.1.3" +version = "1.2.0" dependencies = [ - "bincode", "csv", "env_logger", "kn0sys-lmdb-rs", @@ -1500,9 +1514,10 @@ dependencies = [ "serde", "serde_json", "sysinfo", - "thiserror 2.0.12", + "thiserror 2.0.18", "tokenizers", "uuid", + "wincode", ] [[package]] @@ -1511,12 +1526,6 @@ version = "0.2.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" -[[package]] -name = "virtue" -version = "0.0.18" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "051eb1abcf10076295e815102942cc58f9d5e3b4560e46e53c21e8ff6f3af7b1" - [[package]] name = "wasi" version = "0.11.0+wasi-snapshot-preview1" @@ -1621,6 +1630,31 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "wincode" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd358c35ea3fbf8590e8b9d9e7fe6450701c520c8b58c320aea0b8b75f8d9866" +dependencies = [ + "pastey", + "proc-macro2", + "quote", + "thiserror 2.0.18", + "wincode-derive", +] + +[[package]] +name = "wincode-derive" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6505f603ab2302ff300837c3c96e5b1c6e4b65a66b756e3eb07376c935ff1907" +dependencies = [ + "darling 0.21.3", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "windows" version = "0.61.1" diff --git a/Cargo.toml b/Cargo.toml index 33ecbd8..20c496b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "valentinus" -version = "1.1.3" +version = "1.2.0" edition = "2024" authors = ["n12n "] documentation = "https://docs.rs/valentinus" @@ -25,7 +25,7 @@ crate-type = ["lib"] # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] -bincode = { version = "2.0.1", features = ["serde"]} +wincode = { version = ">=0.4.0", features = ["derive"] } kn0sys-nn = "0.9.2" kn0sys-lmdb-rs = "0.1.6" log = "0.4" diff --git a/src/embeddings.rs b/src/embeddings.rs index 866b4c4..00d3186 100644 --- a/src/embeddings.rs +++ b/src/embeddings.rs @@ -82,7 +82,6 @@ //! ``` use crate::{database::*, md2f::filter_where, onnx::*}; -use bincode::{config, Decode, Encode}; use kn0sys_lmdb_rs as lmdb; use kn0sys_lmdb_rs::MdbError; use kn0sys_nn::distance::L2Dist; @@ -95,6 +94,7 @@ use std::collections::HashMap; use std::sync::{Arc, LazyLock, RwLock}; use thiserror::Error; use uuid::Uuid; +use wincode::{SchemaRead, SchemaWrite}; // --- Public Structs and Enums --- @@ -114,12 +114,13 @@ pub struct Valentinus { /// This struct holds all the data related to a collection, including documents, /// metadata, and the vector embeddings themselves. It is designed to be immutable /// once created and cached in memory. -#[derive(Clone, Debug, Default, Deserialize, Serialize)] +#[derive(Clone, Debug, Default, Deserialize, Serialize, SchemaWrite, SchemaRead)] pub struct EmbeddingCollection { /// The original text documents. documents: Vec, /// The vector embeddings generated from the documents. - embeddings: Array2, + data: Vec, + shape: (usize, usize), /// Metadata associated with each document, matched by index. metadata: Vec>, /// Path to the ONNX model files used for this collection. @@ -135,7 +136,7 @@ pub struct EmbeddingCollection { } /// Identifier for the model used with the collection. -#[derive(Clone, Debug, Default, serde::Deserialize, Serialize)] +#[derive(Clone, Debug, Default, Deserialize, Serialize, SchemaWrite, SchemaRead)] pub enum ModelType { /// AllMiniLmL12V2 model. AllMiniLmL12V2, @@ -158,9 +159,12 @@ pub struct CosineQueryResult { /// Error handling enum for all operations. #[derive(Debug, Error)] pub enum ValentinusError { - /// Bincode serialization/deserialization failure. + /// Cache read error + #[error("Cache error: {0}")] + CacheError(String), + /// Wincode serialization/deserialization failure. #[error("Serialization/deserialization error: {0}")] - BincodeError(String), + WincodeError(String), /// A collection with the given name was not found. #[error("Collection '{0}' not found")] CollectionNotFound(String), @@ -192,20 +196,18 @@ pub enum ValentinusError { // --- Internal Serialization Structs (for backward compatibility) --- -#[derive(Decode, Encode)] +#[derive(SchemaWrite, SchemaRead)] struct PreCollection { - #[bincode(with_serde)] serde: EmbeddingCollection, } -#[derive(Debug, Default, Deserialize, Serialize)] +#[derive(Debug, Default, Deserialize, Serialize, SchemaWrite, SchemaRead)] struct KeyViewIndexer { values: Vec, } -#[derive(Default, Decode, Encode)] +#[derive(Default, SchemaWrite, SchemaRead)] struct KVIndexer { - #[bincode(with_serde)] serde: KeyViewIndexer, } @@ -257,15 +259,17 @@ impl Valentinus { // --- 2. Generate Embeddings --- info!("Generating embeddings for new collection '{}'", name); - let embeddings = batch_embeddings(&model_path, &documents) - .map_err(ValentinusError::OnnxError)?; - + let array_embeddings: Array2 = + batch_embeddings(&model_path, &documents).map_err(ValentinusError::OnnxError)?; + let shape = (array_embeddings.nrows(), array_embeddings.ncols()); + let data = array_embeddings.into_raw_vec_and_offset().0; // --- 3. Prepare Collection Struct --- let key = format!("{}-{}", VALENTINUS_KEY, Uuid::new_v4()); let view = format!("{}-{}", VALENTINUS_VIEW, name); let collection = EmbeddingCollection { documents, - embeddings, + data, + shape, metadata, model_path, model_type, @@ -306,8 +310,8 @@ impl Valentinus { let pre_collection = PreCollection { serde: collection.clone(), }; - let encoded_collection = bincode::encode_to_vec(&pre_collection, config::standard()) - .map_err(|e| ValentinusError::BincodeError(e.to_string()))?; + let encoded_collection = wincode::serialize(&pre_collection) + .map_err(|e| ValentinusError::WincodeError(e.to_string()))?; write_chunks_in_txn( &txn, @@ -322,10 +326,16 @@ impl Valentinus { } /// Retrieves a collection, loading it from the database and caching it if necessary. - pub fn get_collection(&self, view_name: &str) -> Result, ValentinusError> { + pub fn get_collection( + &self, + view_name: &str, + ) -> Result, ValentinusError> { // --- 1. Check cache with a read lock --- { - let cache = self.collections.read().unwrap(); + let cache = self + .collections + .read() + .map_err(|e| ValentinusError::CacheError(e.to_string()))?; if let Some(collection) = cache.values().find(|c| c.view.ends_with(view_name)) { info!("Cache hit for collection '{}'", view_name); return Ok(Arc::clone(collection)); @@ -342,14 +352,16 @@ impl Valentinus { } // --- 3. Load from DB --- - info!("Cache miss. Loading collection '{}' from database.", view_name); + info!( + "Cache miss. Loading collection '{}' from database.", + view_name + ); let key = self.get_key_for_view(view_name)?; let collection_data = read(&self.db.env, &self.db.handle, &key.as_bytes().to_vec())? .ok_or_else(|| ValentinusError::CollectionNotFound(view_name.to_string()))?; - let (pre_collection, _): (PreCollection, usize) = - bincode::decode_from_slice(&collection_data, config::standard()) - .map_err(|e| ValentinusError::BincodeError(e.to_string()))?; + let pre_collection: PreCollection = wincode::deserialize(&collection_data) + .map_err(|e| ValentinusError::WincodeError(e.to_string()))?; let collection = Arc::new(pre_collection.serde); cache.insert(key, Arc::clone(&collection)); @@ -421,9 +433,11 @@ impl Valentinus { let mut results: Vec<(f32, String, Vec)> = Vec::new(); + // Consume the flattened data back to Array2 + let collection_embeddings = + Array2::from_shape_vec(collection.shape, collection.data.clone()).unwrap_or_default(); // --- Iterate safely using enumerate to get a reliable index --- - for (index, (cv, sentence)) in collection - .embeddings + for (index, (cv, sentence)) in collection_embeddings .axis_iter(Axis(0)) .zip(collection.documents.iter()) .enumerate() @@ -431,8 +445,14 @@ impl Valentinus { let metadata = &collection.metadata[index]; let raw_f = f_where.as_deref().unwrap_or(&[]); - if !is_filtering || filter_where(raw_f, metadata).map_err(|_| ValentinusError::Md2fsError)? { - let dot_product: f32 = query_embedding.iter().zip(cv.iter()).map(|(a, b)| a * b).sum(); + if !is_filtering + || filter_where(raw_f, metadata).map_err(|_| ValentinusError::Md2fsError)? + { + let dot_product: f32 = query_embedding + .iter() + .zip(cv.iter()) + .map(|(a, b)| a * b) + .sum(); results.push((dot_product, sentence.clone(), metadata.clone())); } } @@ -476,9 +496,10 @@ impl Valentinus { let qv = batch_embeddings(&collection.model_path, &qv_string) .map_err(ValentinusError::OnnxError)?; let query_embedding = qv.index_axis(Axis(0), 0); - + let collection_embeddings = + Array2::from_shape_vec(collection.shape, collection.data.clone()).unwrap_or_default(); let nn = CommonNearestNeighbour::KdTree - .batch(&collection.embeddings, L2Dist) + .batch(&collection_embeddings, L2Dist) .map_err(|e| ValentinusError::NearestError(e.to_string()))?; let nearest = nn @@ -486,12 +507,13 @@ impl Valentinus { .map_err(|e| ValentinusError::NearestError(e.to_string()))?; if nearest.is_empty() { - return Err(ValentinusError::NotFound("No nearest neighbor found.".to_string())); + return Err(ValentinusError::NotFound( + "No nearest neighbor found.".to_string(), + )); } let nearest_embedding = nearest[0].0.to_vec(); - let position = collection - .embeddings + let position = collection_embeddings .axis_iter(Axis(0)) .position(|x| x.to_vec() == nearest_embedding); @@ -523,9 +545,8 @@ impl Valentinus { indexer_name: &str, ) -> Result { match txn.bind(db_handle).get::>(&indexer_name.as_bytes()) { - Ok(bytes) => Ok(bincode::decode_from_slice(&bytes, config::standard()) - .map_err(|e| ValentinusError::BincodeError(e.to_string()))? - .0), + Ok(bytes) => Ok(wincode::deserialize(&bytes) + .map_err(|e| ValentinusError::WincodeError(e.to_string()))?), Err(MdbError::NotFound) => Ok(KVIndexer::default()), // Return empty if not found Err(e) => Err(ValentinusError::DatabaseError(e)), } @@ -537,8 +558,8 @@ impl Valentinus { indexer_name: &str, indexer: &KVIndexer, ) -> Result<(), ValentinusError> { - let encoded = bincode::encode_to_vec(indexer, config::standard()) - .map_err(|e| ValentinusError::BincodeError(e.to_string()))?; + let encoded = wincode::serialize(indexer) + .map_err(|e| ValentinusError::WincodeError(e.to_string()))?; txn.bind(db_handle) .set(&indexer_name.as_bytes(), &encoded)?; Ok(()) @@ -571,7 +592,7 @@ mod tests { use std::{fs, fs::File, path::Path}; /// Test data structure for CSV parsing. - #[derive(Default, Deserialize)] + #[derive(Default, Deserialize, SchemaWrite, SchemaRead)] struct Review { review: Option, rating: Option, @@ -591,6 +612,7 @@ mod tests { #[test] fn test_full_etl_and_query_workflow() -> Result<(), ValentinusError> { + env_logger::init(); let valentinus = setup_test_env("full_workflow_test"); let collection_name = "tesla_reviews".to_string(); @@ -609,7 +631,7 @@ mod tests { // --- 2. Verify creation by getting the collection --- let collection = valentinus.get_collection(&collection_name)?; assert_eq!(collection.documents, expected_docs); - assert!(!collection.embeddings.is_empty()); + assert!(!collection.data.is_empty()); // --- 3. Test Cosine Query with Filters --- let query_string = "Find the best reviews.".to_string(); @@ -676,7 +698,11 @@ mod tests { for result in rdr.deserialize() { let record: Review = result.unwrap_or_default(); documents.push(record.review.unwrap_or_default()); - let rating = record.rating.unwrap_or_default().parse::().unwrap_or(0); + let rating = record + .rating + .unwrap_or_default() + .parse::() + .unwrap_or(0); let year_str = record.vehicle_title.unwrap_or_default(); let year = if year_str.len() >= 4 { year_str[0..4].to_string()