Can index multiple files together

diff --git a/src/main.rs b/src/main.rs
index 0b56c69f3fda09fbbd5f9408f763604ae3296204..d0db01237ba6e441acf8dc899dcad79dc4530618 100644 (file)
--- a/src/main.rs
+++ b/src/main.rs
@@ -46,16 +46,17 @@
 //! The user for this operation requires only write access to the
 //! table (not superuser).
 //! 
+//! Specifying multiple files is more efficient than indexing one file
+//! in each invocation.
+//! 
 //! ```text
 //! $ export PASSWORD=$(gpg -d pw-cvmigrator.gpg)
-//! $ vecsearch index --file testdata/0
-//! indexing a file
-//! Loaded and encoded 59.479µs
-//! Took 14.982262ms
-//! $ vecsearch index --file testdata/1
-//! ...
-//! $ vecsearch index --file testdata/7
-//! ...
+//! $ vecsearch index --file testdata/0 --file testdata/1
+//! indexing file(s)
+//! Loaded and encoded 58.565µs
+//! Took 15.628167ms
+//! Loaded and encoded 55.513µs
+//! Took 8.018493ms
 //! ```
 //! 
 //! ## Search
@@ -82,7 +83,6 @@
 //! 
 //! ## TODO
 //! 
-//! - index multiple files
 //! - model from main not PR
 //! - env support for all args
 //! 
@@ -133,8 +133,8 @@ enum Action {
        password: String,
 
        #[arg(long)]
-       /// The file containing document contents
-       file: String,
+       /// The file containing document contents. Specify multiple
+       file: Vec<String>,
     },
 
     /// Search the database for documents matching --search
@@ -265,7 +265,7 @@ fn init_database(dbname: String, host: String, user: String, password: String) -
     Ok(())
 }
 
-fn get_embeddings(input: &String, model: BertModel, mut tokenizer: Tokenizer) -> Result<Vec<f32>> {
+fn get_embeddings(input: &String, model: &BertModel, mut tokenizer: Tokenizer) -> Result<Vec<f32>> {
     let start = std::time::Instant::now();
     let device = &model.device;
     let tokenizer = tokenizer
@@ -292,12 +292,9 @@ fn get_embeddings(input: &String, model: BertModel, mut tokenizer: Tokenizer) ->
 }
 
 fn index(dbname: String, host: String, user: String, password: String,
-        file: String, model: BertModel, tokenizer: Tokenizer) -> Result<()> {
-
-    println!("indexing a file");
+        files: Vec<String>, model: BertModel, tokenizer: Tokenizer) -> Result<()> {
 
-    let doc_content = std::fs::read_to_string(file)?;
-    let embeddings = get_embeddings(&doc_content,model,tokenizer)?;
+    println!("indexing file(s)");
 
     let mut client = postgres::Config::new()
        .dbname(&dbname)
@@ -306,11 +303,16 @@ fn index(dbname: String, host: String, user: String, password: String,
        .password(password)
        .connect(NoTls)?;
 
-    client.execute("INSERT INTO documents (content, embedding) \
-                   values ($1, $2) \
-                   ON CONFLICT (content) DO UPDATE SET embedding = $2",
-                  &[&doc_content,&Vector::from(embeddings)],
-    )?;
+    for file in files {
+       let doc_content = std::fs::read_to_string(file)?;
+       let embeddings = get_embeddings(&doc_content,&model,tokenizer.clone())?;
+
+       client.execute("INSERT INTO documents (content, embedding) \
+                       values ($1, $2) \
+                       ON CONFLICT (content) DO UPDATE SET embedding = $2",
+                      &[&doc_content,&Vector::from(embeddings)],
+       )?;
+    }
 
     let _ = client.close();
     Ok(())
@@ -320,7 +322,7 @@ fn search(dbname: String, host: String, user: String, password: String,
         search: String, model: BertModel, tokenizer: Tokenizer) -> Result<()> {
 
     println!("searching for document matches");
-    let embeddings = get_embeddings(&search,model,tokenizer)?;
+    let embeddings = get_embeddings(&search,&model,tokenizer.clone())?;
 
     let mut client = postgres::Config::new()
        .dbname(&dbname)