Adding preliminary load

Author: GavinMendelGleason

Cargo.lock

@@ -30,6 +30,7 @@ urlencoding = "2.1"
 packed_simd = {version = "0.3.8", optional=true}
 aligned_box = "0.2"
 tiktoken-rs = "0.4"
+itertools = "0.10"
 
 [features]
 simd = ["packed_simd"]

Cargo.toml

@@ -1,7 +1,7 @@
 use crate::{
     openai::embeddings_for,
     server::Operation,
-    vecmath::Embedding,
+    vecmath::{self, Embedding},
     vectors::{LoadedVec, VectorStore},
 };
 use hnsw::{Hnsw, Searcher};
@@ -62,19 +62,14 @@ impl Metric<Point> for OpenAI {
     fn distance(&self, p1: &Point, p2: &Point) -> u32 {
         let a = p1.vec();
         let b = p2.vec();
-        let f = a
-            .iter()
-            .zip(b.iter())
-            .map(|(&a, &b)| (a - b).powi(2))
-            .sum::<f32>()
-            .sqrt();
+        let f = vecmath::normalized_cosine_distance(a, b);
         f.to_bits()
     }
 }
 
 impl Metric<IndexPoint> for OpenAI {
     type Unit = u32;
-    fn distance(&self, p1: &IndexPoint, p2: &IndexPoint) -> u32 {
+    fn distance(&self, _p1: &IndexPoint, _p2: &IndexPoint) -> u32 {
         unimplemented!()
     }
 }
@@ -86,8 +81,6 @@ pub enum PointOperation {
     Delete { id: String },
 }
 
-pub const OPENAI_API_KEY: &str = "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";
-
 enum Op {
     Insert,
     Changed,
@@ -97,6 +90,7 @@ pub async fn operations_to_point_operations(
     domain: &str,
     vector_store: &VectorStore,
     structs: Vec<Result<Operation, std::io::Error>>,
+    key: &str,
 ) -> Vec<PointOperation> {
     let ops: Vec<Operation> = structs.into_iter().map(|ro| ro.unwrap()).collect();
     let tuples: Vec<(Op, String, String)> = ops
@@ -108,7 +102,7 @@ pub async fn operations_to_point_operations(
         })
         .collect();
     let strings: Vec<String> = tuples.iter().map(|(_, s, _)| s.to_string()).collect();
-    let vecs: Vec<Embedding> = embeddings_for(OPENAI_API_KEY, &strings).await.unwrap();
+    let vecs: Vec<Embedding> = embeddings_for(key, &strings).await.unwrap();
     let domain = vector_store.get_domain(domain).unwrap();
     let loaded_vecs = vector_store
         .add_and_load_vecs(&domain, vecs.iter())
@@ -213,7 +207,7 @@ pub fn search(
 }
 
 pub fn serialize_index(mut path: PathBuf, name: &str, hnsw: HnswIndex) -> io::Result<()> {
-    let name = encode(name);
+    //let name = encode(name);
     path.push(format!("{name}.hnsw"));
     let write_file = File::options().write(true).create(true).open(&path)?;
 

src/indexer.rs

@@ -1,15 +1,27 @@
+use std::future;
+use std::io::ErrorKind;
+use std::path::Path;
+
 use clap::{Parser, Subcommand, ValueEnum};
+use hnsw::Hnsw;
+use indexer::start_indexing_from_operations;
 use indexer::Point;
+use indexer::{operations_to_point_operations, OpenAI};
+use server::Operation;
 use space::Metric;
-use terminusdb_semantic_indexer::vecmath::empty_embedding;
-
-use crate::indexer::OpenAI;
-
+use std::fs::File;
+use std::io::{self, BufRead};
+use {
+    indexer::{create_index_name, HnswIndex},
+    vecmath::empty_embedding,
+    vectors::VectorStore,
+};
 mod indexer;
 mod openai;
 mod server;
 mod vecmath;
 mod vectors;
+use itertools::Itertools;
 
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
@@ -21,13 +33,29 @@ struct Args {
 #[derive(Subcommand, Debug)]
 enum Commands {
     Serve {
+        #[arg(short, long)]
+        key: String,
         #[arg(short, long)]
         directory: String,
         #[arg(short, long, default_value_t = 8080)]
         port: u16,
         #[arg(short, long, default_value_t = 10000)]
         size: usize,
     },
+    Load {
+        #[arg(short, long)]
+        key: String,
+        #[arg(short, long)]
+        commit: String,
+        #[arg(short, long)]
+        domain: String,
+        #[arg(short, long)]
+        directory: String,
+        #[arg(short, long)]
+        input: String,
+        #[arg(short, long, default_value_t = 10000)]
+        size: usize,
+    },
     Embed {
         #[arg(short, long)]
         key: String,
@@ -70,10 +98,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     let args = Args::parse();
     match args.command {
         Commands::Serve {
+            key,
             directory,
             port,
             size,
-        } => server::serve(directory, port, size).await?,
+        } => server::serve(directory, port, size, key).await?,
         Commands::Embed { key, string } => {
             let v = openai::embeddings_for(&key, &[string]).await?;
             eprintln!("{:?}", v);
@@ -108,15 +137,56 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
             };
             println!("distance: {}", distance);
         }
-        Commands::Test {key } => {
-            let v = openai::embeddings_for(&key, &["king".to_string(), "man".to_string(), "woman".to_string(), "queen".to_string()]).await?;
+        Commands::Test { key } => {
+            let v = openai::embeddings_for(
+                &key,
+                &[
+                    "king".to_string(),
+                    "man".to_string(),
+                    "woman".to_string(),
+                    "queen".to_string(),
+                ],
+            )
+            .await?;
             let mut calculated = empty_embedding();
             for i in 0..calculated.len() {
                 calculated[i] = v[0][i] - v[1][i] + v[2][i];
             }
             let distance = vecmath::normalized_cosine_distance(&v[3], &calculated);
             eprintln!("{}", distance);
         }
+        Commands::Load {
+            key,
+            domain,
+            commit,
+            directory,
+            input,
+            size,
+        } => {
+            let path = Path::new(&input);
+            let dirpath = Path::new(&directory);
+            let mut hnsw: HnswIndex = Hnsw::new(OpenAI);
+            let store = VectorStore::new(dirpath, size);
+
+            let f = File::options().read(true).create(true).open(path)?;
+
+            let lines = io::BufReader::new(f).lines();
+
+            let opstream = &lines
+                .map(|l| {
+                    let ro: io::Result<Operation> = serde_json::from_str(&l.unwrap())
+                        .map_err(|e| std::io::Error::new(ErrorKind::Other, e));
+                    ro
+                })
+                .chunks(100);
+
+            for structs in opstream {
+                let structs: Vec<_> = structs.collect();
+                let new_ops =
+                    operations_to_point_operations(&domain.clone(), &store, structs, &key).await;
+                hnsw = start_indexing_from_operations(hnsw, new_ops).unwrap();
+            }
+        }
     }
 
     Ok(())

src/main.rs

@@ -87,8 +87,9 @@ pub enum EmbeddingError {
 }
 
 lazy_static! {
-    static ref ENCODER: CoreBPE  = cl100k_base().unwrap();
+    static ref ENCODER: CoreBPE = cl100k_base().unwrap();
 }
+
 fn tokens_for(s: &str) -> Vec<usize> {
     ENCODER.encode_with_special_tokens(s)
 }
@@ -100,7 +101,6 @@ fn truncated_tokens_for(s: &str) -> Vec<usize> {
         tokens.truncate(MAX_TOKEN_COUNT);
         let decoded = ENCODER.decode(tokens.clone()).unwrap();
         eprintln!("truncating to {decoded}");
-
     }
 
     tokens
@@ -115,7 +115,7 @@ pub async fn embeddings_for(
         static ref CLIENT: Client = Client::new();
     }
 
-    let token_lists: Vec<_> = strings.iter().map(|s|truncated_tokens_for(s)).collect();
+    let token_lists: Vec<_> = strings.iter().map(|s| truncated_tokens_for(s)).collect();
 
     let mut req = Request::new(Method::POST, ENDPOINT.clone());
     let headers = req.headers_mut();

src/openai.rs

@@ -37,7 +37,6 @@ use crate::indexer::search;
 use crate::indexer::serialize_index;
 use crate::indexer::Point;
 use crate::indexer::PointOperation;
-use crate::indexer::OPENAI_API_KEY;
 use crate::indexer::{start_indexing_from_operations, HnswIndex, IndexIdentifier, OpenAI};
 use crate::openai::embeddings_for;
 use crate::vectors::VectorStore;
@@ -182,6 +181,7 @@ pub struct QueryResult {
 }
 
 pub struct Service {
+    api_key: String,
     path: PathBuf,
     vector_store: VectorStore,
     pending: Mutex<HashSet<String>>,
@@ -271,9 +271,10 @@ impl Service {
         s
     }
 
-    fn new<P: Into<PathBuf>>(path: P, num_bufs: usize) -> Self {
+    fn new<P: Into<PathBuf>>(path: P, num_bufs: usize, key: String) -> Self {
         let path = path.into();
         Service {
+            api_key: key,
             path: path.clone(),
             vector_store: VectorStore::new(path, num_bufs),
             pending: Mutex::new(HashSet::new()),
@@ -306,37 +307,56 @@ impl Service {
         tokio::spawn(async move {
             let index_id = create_index_name(&domain, &commit);
             if self.test_and_set_pending(index_id.clone()).await {
-                let mut opstream = get_operations_from_terminusdb(
+                let opstream = get_operations_from_terminusdb(
                     domain.clone(),
                     commit.clone(),
                     previous.clone(),
                 )
                 .await
                 .unwrap()
                 .chunks(100);
-                let mut point_ops: Vec<PointOperation> = Vec::new();
-                while let Some(structs) = opstream.next().await {
-                    let mut new_ops =
-                        operations_to_point_operations(&domain, &self.vector_store, structs).await;
-                    point_ops.append(&mut new_ops)
-                }
-                let id = create_index_name(&domain, &commit);
-                let hnsw = self
-                    .load_hnsw_for_indexing(IndexIdentifier {
-                        domain,
-                        commit,
-                        previous,
-                    })
+                let (id, hnsw) = self
+                    .process_operation_chunks(opstream, domain, commit, previous, &index_id)
                     .await;
-                let hnsw = start_indexing_from_operations(hnsw, point_ops).unwrap();
-                let path = self.path.clone();
-                serialize_index(path, &index_id, hnsw.clone()).unwrap();
                 self.set_index(id, hnsw.into()).await;
                 self.clear_pending(&index_id).await;
             }
         });
     }
 
+    async fn process_operation_chunks(
+        self: &Arc<Self>,
+        mut opstream: futures::stream::Chunks<
+            impl Stream<Item = Result<Operation, io::Error>> + Unpin,
+        >,
+        domain: String,
+        commit: String,
+        previous: Option<String>,
+        index_id: &str,
+    ) -> (String, Hnsw<OpenAI, Point, rand_pcg::Lcg128Xsl64, 12, 24>) {
+        let id = create_index_name(&domain, &commit);
+        let mut hnsw = self
+            .load_hnsw_for_indexing(IndexIdentifier {
+                domain: domain.clone(),
+                commit,
+                previous,
+            })
+            .await;
+        while let Some(structs) = opstream.next().await {
+            let new_ops = operations_to_point_operations(
+                &domain.clone(),
+                &self.vector_store,
+                structs,
+                &self.api_key,
+            )
+            .await;
+            hnsw = start_indexing_from_operations(hnsw, new_ops).unwrap();
+        }
+        let path = self.path.clone();
+        serialize_index(path, index_id, hnsw.clone()).unwrap();
+        (id, hnsw)
+    }
+
     async fn get(self: Arc<Self>, req: Request<Body>) -> Result<Response<Body>, Infallible> {
         let uri = req.uri();
         match dbg!(uri_to_spec(uri)) {
@@ -395,7 +415,7 @@ impl Service {
             }) => {
                 let body_bytes = hyper::body::to_bytes(req.into_body()).await.unwrap();
                 let q = String::from_utf8(body_bytes.to_vec()).unwrap();
-                let vec = Box::new((embeddings_for(OPENAI_API_KEY, &[q]).await.unwrap())[0]);
+                let vec = Box::new((embeddings_for(&self.api_key, &[q]).await.unwrap())[0]);
                 let qp = Point::Mem { vec };
                 let index_id = create_index_name(&domain, &commit);
                 // if None, then return 404
@@ -421,9 +441,10 @@ pub async fn serve<P: Into<PathBuf>>(
     directory: P,
     port: u16,
     num_bufs: usize,
+    key: String,
 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     let addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), port);
-    let service = Arc::new(Service::new(directory, num_bufs));
+    let service = Arc::new(Service::new(directory, num_bufs, key));
     let make_svc = make_service_fn(move |_conn| {
         let s = service.clone();
         async {