diff options
Diffstat (limited to 'src/ai_embedders.rs')
-rw-r--r-- | src/ai_embedders.rs | 25 |
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()?) } } |