feat(cli): Add information to progress bar

emmyoh · Apr 30, 2024 · 004cd9a · 004cd9a
1 parent 5ad3ea7
commit 004cd9a
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 7 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -34,7 +34,7 @@ space = "0.17.0"
 lz4_flex = { version = "0.11.3", default-features = false, features = ["frame"] }
 ticky = { version = "1.0.2", optional = true }
 pretty-duration = { version = "0.1.1", optional = true }
-indicatif = { version = "0.17.8", optional = true }
+indicatif = { version = "0.17.8", optional = true, features = ["rayon", "improved_unicode"] }
 distances = "1.6.3"
 candle-examples = "0.4.1"
 candle-core = "0.4.1"

diff --git a/src/main.rs b/src/main.rs
@@ -1,5 +1,6 @@
 use clap::{command, Parser, Subcommand};
 use fastembed::TextEmbedding;
+use indicatif::ProgressStyle;
 use indicatif::{ProgressBar, ProgressDrawTarget};
 use pretty_duration::pretty_duration;
 use rodio::{Decoder, OutputStream, Sink};
@@ -138,8 +139,9 @@ fn main() -> Result<(), Box<dyn Error>> {
                     let progress_bar = ProgressBar::with_draw_target(
                         Some(num_texts.try_into()?),
                         ProgressDrawTarget::hidden(),
-                    )
-                    .with_message(format!("Inserting texts from {} file(s).", num_texts));
+                    );
+                    // .with_message(format!("Inserting texts from {} file(s).", num_texts));
+                    progress_bar.set_style(progress_bar_style()?);
                     let mut i = 0;
                     let mut texts = Vec::new();
                     // Insert texts in batches of INSERT_BATCH_SIZE.
@@ -228,8 +230,9 @@ fn main() -> Result<(), Box<dyn Error>> {
                 let progress_bar = ProgressBar::with_draw_target(
                     Some(num_images.try_into()?),
                     ProgressDrawTarget::hidden(),
-                )
-                .with_message(format!("Inserting images from {} file(s).", num_images));
+                );
+                // .with_message(format!("Inserting images from {} file(s).", num_images));
+                progress_bar.set_style(progress_bar_style()?);
                 let images: Vec<String> = file_paths
                     .into_iter()
                     .map(|x| x.to_str().unwrap().to_string())
@@ -316,8 +319,9 @@ fn main() -> Result<(), Box<dyn Error>> {
                 let progress_bar = ProgressBar::with_draw_target(
                     Some(num_sounds.try_into()?),
                     ProgressDrawTarget::hidden(),
-                )
-                .with_message(format!("Inserting sounds from {} file(s).", num_sounds));
+                );
+                // .with_message(format!("Inserting sounds from {} file(s).", num_sounds));
+                progress_bar.set_style(progress_bar_style()?);
                 let sounds: Vec<String> = file_paths
                     .into_iter()
                     .map(|x| x.to_str().unwrap().to_string())
@@ -391,3 +395,7 @@ fn main() -> Result<(), Box<dyn Error>> {
     }
     Ok(())
 }
+
+fn progress_bar_style() -> Result<ProgressStyle, Box<dyn Error>> {
+    Ok(ProgressStyle::with_template("[{elapsed} elapsed, {eta} remaining ({duration} total)] {wide_bar:.cyan/blue} {human_pos} of {human_len} ({percent}%) {msg}")?)
+}