mistralrs_core/utils/
progress.rs

1use indicatif::{MultiProgress, ProgressBar, ProgressBarIter, ProgressIterator, ProgressStyle};
2use tqdm::Iter;
3
4// Optionally display a progress bar via the `tqdm` crate:
5// Usage: `iter.with_progress(true)`
6// Similar to the `iter.tqdm()` feature except this supports opt-in via parameter.
7pub trait IterWithProgress<'a, T>: Iterator<Item = T> + 'a {
8    fn with_progress(self, is_silent: bool) -> Box<dyn Iterator<Item = T> + 'a>
9    where
10        Self: Sized,
11    {
12        // TODO: Should `is_silent` instead be referenced as a global read-only state? (`AtomicBool`)
13        if is_silent {
14            Box::new(self)
15        } else {
16            Box::new(self.tqdm())
17        }
18    }
19}
20
21impl<'a, T: Iterator + 'a> IterWithProgress<'a, T::Item> for T {}
22
23/// Nice progress bar with over an iterator and a message.
24/// COLOR is one of r,g,b
25pub struct NiceProgressBar<'a, T: ExactSizeIterator, const COLOR: char = 'b'>(
26    pub T,
27    pub &'static str,
28    pub &'a MultiProgress,
29);
30
31impl<T: ExactSizeIterator, const COLOR: char> IntoIterator for NiceProgressBar<'_, T, COLOR> {
32    type IntoIter = ProgressBarIter<T>;
33    type Item = T::Item;
34
35    fn into_iter(self) -> Self::IntoIter {
36        let color = match COLOR {
37            'b' => "blue",
38            'g' => "green",
39            'r' => "red",
40            other => panic!("Color char `{other}` not supported"),
41        };
42        let bar = ProgressBar::new(self.0.len() as u64);
43        bar.set_style(
44            ProgressStyle::default_bar()
45                .template(&format!(
46                    "{}: [{{elapsed_precise}}] [{{bar:40.{color}/{color}}}] {{pos}}/{{len}} ({{eta}})",
47                    self.1
48                ))
49                .unwrap()
50                .progress_chars("#>-"),
51        );
52
53        // Add to the multi progress
54        self.2.add(bar.clone());
55
56        self.0.progress_with(bar)
57    }
58}
mistralrs_core/utils/progress.rs

mistralrs_core/utils/
progress.rs
