aboutsummaryrefslogtreecommitdiff
path: root/src/ai_embedders.rs
diff options
context:
space:
mode:
Diffstat (limited to 'src/ai_embedders.rs')
-rw-r--r--src/ai_embedders.rs25
1 files changed, 21 insertions, 4 deletions
diff --git a/src/ai_embedders.rs b/src/ai_embedders.rs
index 14e1d8c..b30d890 100644
--- a/src/ai_embedders.rs
+++ b/src/ai_embedders.rs
@@ -1,6 +1,7 @@
use anyhow::Result;
+use indicatif::{ProgressBar, ProgressIterator, ProgressStyle};
use serde::{Deserialize, Serialize};
-use std::{io::{copy, Cursor}, path::PathBuf, process::Command};
+use std::{io::{BufRead, BufReader, copy, Cursor}, path::PathBuf, process::{Command, Stdio}};
use crate::{Config, BatchEmbedder, MetricElem};
@@ -46,20 +47,36 @@ impl BatchEmbedder for ContentEmbedder<'_> {
let api_prog = include_bytes!("imgbeddings-api.py");
copy(&mut Cursor::new(api_prog), &mut std::fs::File::create(&script_file)?)?;
+ let bar = ProgressBar::new_spinner();
+ bar.set_style(ProgressStyle::with_template("{spinner} {msg}")?);
+ bar.enable_steady_tick(std::time::Duration::from_millis(100));
+
+ bar.set_message("Creating venv...");
Command::new("python3")
.args(["-m", "venv", venv_dir.to_str().unwrap()])
+ .stdout(Stdio::null())
.spawn()?
.wait()?;
+
+ bar.set_message("Installing/checking packages...");
Command::new(venv_dir.join("bin/pip3"))
.args(["install", "imgbeddings"])
+ .stdout(Stdio::null())
.spawn()?
.wait()?;
+ bar.finish();
- let output = Command::new(venv_dir.join("bin/python3"))
+ let child = Command::new(venv_dir.join("bin/python3"))
.arg(script_file)
.args(paths)
- .output()?;
+ .stderr(Stdio::null())
+ .stdout(Stdio::piped())
+ .spawn()?;
- Ok(serde_json::from_slice(&output.stdout)?)
+ Ok(BufReader::new(child.stdout.unwrap())
+ .lines()
+ .progress_count(paths.len().try_into().unwrap())
+ .map(|l| Ok::<_, anyhow::Error>(serde_json::from_str(&l?)?))
+ .try_collect()?)
}
}