From a56d5f2ddfcfb290c308f2b1a3464c3bebc340de Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Sun, 8 Sep 2024 19:50:02 +0800
Subject: feat: save rag in YAML instead of bin (#848)

---
 src/rag/mod.rs           | 23 +++++++++-------
 src/rag/serde_vectors.rs | 69 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 82 insertions(+), 10 deletions(-)
 create mode 100644 src/rag/serde_vectors.rs

(limited to 'src/rag')

diff --git a/src/rag/mod.rs b/src/rag/mod.rs
index f6a511d..e03b3f1 100644
--- a/src/rag/mod.rs
+++ b/src/rag/mod.rs
@@ -8,18 +8,17 @@ use crate::utils::*;
 
 mod bm25;
 mod loader;
+mod serde_vectors;
 mod splitter;
 
-use anyhow::bail;
-use anyhow::{anyhow, Context, Result};
+use anyhow::{anyhow, bail, Context, Result};
 use hnsw_rs::prelude::*;
 use indexmap::{IndexMap, IndexSet};
 use inquire::{required, validator::Validation, Confirm, Select, Text};
 use path_absolutize::Absolutize;
 use serde::{Deserialize, Serialize};
 use serde_json::json;
-use std::collections::HashMap;
-use std::{fmt::Debug, io::BufReader, path::Path};
+use std::{collections::HashMap, fmt::Debug, fs, path::Path};
 
 pub struct Rag {
     config: GlobalConfig,
@@ -103,9 +102,8 @@ impl Rag {
 
     pub fn load(config: &GlobalConfig, name: &str, path: &Path) -> Result<Self> {
         let err = || format!("Failed to load rag '{name}' at '{}'", path.display());
-        let file = std::fs::File::open(path).with_context(err)?;
-        let reader = BufReader::new(file);
-        let data: RagData = bincode::deserialize_from(reader).with_context(err)?;
+        let content = fs::read_to_string(path).with_context(err)?;
+        let data: RagData = serde_yaml::from_str(&content).with_context(err)?;
         Self::create(config, name, path, data)
     }
 
@@ -236,9 +234,13 @@ impl Rag {
         }
         let path = Path::new(&self.path);
         ensure_parent_exists(path)?;
-        let mut file = std::fs::File::create(path)?;
-        bincode::serialize_into(&mut file, &self.data)
-            .with_context(|| format!("Failed to save rag '{}'", self.name))?;
+
+        let content = serde_yaml::to_string(&self.data)
+            .with_context(|| format!("Failed to serde rag '{}'", self.name))?;
+        fs::write(path, content).with_context(|| {
+            format!("Failed to save rag '{}' to '{}'", self.name, path.display())
+        })?;
+
         Ok(true)
     }
 
@@ -576,6 +578,7 @@ pub struct RagData {
     pub next_file_id: FileId,
     pub document_paths: Vec<String>,
     pub files: IndexMap<FileId, RagFile>,
+    #[serde(with = "serde_vectors")]
     pub vectors: IndexMap<DocumentId, Vec<f32>>,
 }
 
diff --git a/src/rag/serde_vectors.rs b/src/rag/serde_vectors.rs
new file mode 100644
index 0000000..894c22c
--- /dev/null
+++ b/src/rag/serde_vectors.rs
@@ -0,0 +1,69 @@
+use super::*;
+
+use base64::{engine::general_purpose::STANDARD, Engine};
+use serde::{de, Deserializer, Serializer};
+
+pub fn serialize<S>(
+    vectors: &IndexMap<DocumentId, Vec<f32>>,
+    serializer: S,
+) -> Result<S::Ok, S::Error>
+where
+    S: Serializer,
+{
+    let encoded_map: IndexMap<String, String> = vectors
+        .iter()
+        .map(|(key, vec)| {
+            let (h, l) = split_document_id(*key);
+            let byte_slice = unsafe {
+                std::slice::from_raw_parts(
+                    vec.as_ptr() as *const u8,
+                    vec.len() * std::mem::size_of::<f32>(),
+                )
+            };
+            (format!("{h}-{l}"), STANDARD.encode(byte_slice))
+        })
+        .collect();
+
+    encoded_map.serialize(serializer)
+}
+
+pub fn deserialize<'de, D>(deserializer: D) -> Result<IndexMap<DocumentId, Vec<f32>>, D::Error>
+where
+    D: Deserializer<'de>,
+{
+    let encoded_map: IndexMap<String, String> =
+        IndexMap::<String, String>::deserialize(deserializer)?;
+
+    let mut decoded_map = IndexMap::new();
+    for (key, base64_str) in encoded_map {
+        let decoded_key: DocumentId = key
+            .split_once('-')
+            .and_then(|(h, l)| {
+                let h = h.parse::<usize>().ok()?;
+                let l = l.parse::<usize>().ok()?;
+                Some(combine_document_id(h, l))
+            })
+            .ok_or_else(|| de::Error::custom(format!("Invalid key '{key}'")))?;
+
+        let decoded_data = STANDARD.decode(&base64_str).map_err(de::Error::custom)?;
+
+        if decoded_data.len() % std::mem::size_of::<f32>() != 0 {
+            return Err(de::Error::custom(format!("Invalid vector at '{key}'")));
+        }
+
+        let num_f32s = decoded_data.len() / std::mem::size_of::<f32>();
+
+        let mut vec_f32 = vec![0.0f32; num_f32s];
+        unsafe {
+            std::ptr::copy_nonoverlapping(
+                decoded_data.as_ptr(),
+                vec_f32.as_mut_ptr() as *mut u8,
+                decoded_data.len(),
+            );
+        }
+
+        decoded_map.insert(decoded_key, vec_f32);
+    }
+
+    Ok(decoded_map)
+}
-- 
cgit v1.2.3