ravenrs_ai/
manager.rs

1//! The screen for choosing models: offers what fits, and does the work away from
2//! the UI thread.
3
4use std::cell::RefCell;
5use std::rc::Rc;
6use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
7use std::sync::{Arc, Mutex};
8use std::time::Instant;
9
10use ravenrs_app_i18n::{Language, Texts};
11use ravenrs_reactive::State;
12use ravenrs_widgets::{
13    BuildCx, ModelChoice, ModelPicker, ModelPickerLabels, ModelStatus, NodeId, UiWaker, View,
14};
15
16use crate::{
17    HEARING_RATE, LocalAssistant, LocalEngine, LocalModel, Machine, ModelError, ModelFamily, ModelFolder,
18    ModelState, Transcriber,
19};
20
21/// How a model that hears does on this machine, on a quiet tone it has to listen to whole.
22fn hearing_test(folder: &ModelFolder, model: &LocalModel, started: Instant) -> Option<Tested> {
23    let mut transcriber = Transcriber::load(folder, model).ok()?;
24    let tone: Vec<f32> = (0..TEST_HEARING_SECONDS * HEARING_RATE)
25        .map(|i| (i as f32 * 440.0 * std::f32::consts::TAU / HEARING_RATE as f32).sin() * 0.01)
26        .collect();
27    transcriber.transcribe(&tone, "pt").ok()?;
28    let seconds = started.elapsed().as_secs_f64();
29    let heard = TEST_HEARING_SECONDS as f64;
30    let tokens_per_second = if seconds <= heard / 5.0 {
31        FAST
32    } else if seconds <= heard {
33        USABLE
34    } else {
35        0.0
36    };
37    Some(Tested { seconds, tokens_per_second })
38}
39
40/// How many times a download tells the UI it moved: one per cent is plenty.
41const STEPS: u64 = 100;
42const GIGABYTE: f32 = 1_000_000_000.0;
43/// The file in the folder that says which model is in use.
44const IN_USE_FILE: &str = "in-use";
45/// What the speed test asks, and how long an answer it wants.
46const TEST_SYSTEM: &str = "Answer in one short sentence.";
47const TEST_USER: &str = "Why keep a copy of important files?";
48const TEST_TOKENS: usize = 40;
49/// Tokens per second from which an answer feels instant, and from which it can be
50/// read as it arrives.
51const FAST: f64 = 15.0;
52const USABLE: f64 = 5.0;
53/// What a model that hears is tested with: ten seconds, heard in two feels instant, and
54/// in ten it keeps up with whoever speaks.
55const TEST_HEARING_SECONDS: usize = 10;
56
57/// The words a [`ModelManager`] shows. English by default.
58#[derive(Clone, Debug)]
59#[non_exhaustive]
60pub struct ModelManagerLabels {
61    /// The picker's own words.
62    pub picker: ModelPickerLabels,
63    /// Before how many requests a model understood.
64    pub understood: String,
65    /// When that was not measured.
66    pub not_measured: String,
67    /// Before a model's license.
68    pub license: String,
69    /// Before how much memory a model takes.
70    pub memory: String,
71    /// Writes numbers with a decimal comma.
72    pub decimal_comma: bool,
73    /// When the connection failed.
74    pub network_error: String,
75    /// When the disk refused.
76    pub disk_error: String,
77    /// When what arrived is not the published file.
78    pub wrong_file: String,
79    /// Before how long the speed test took, in seconds.
80    pub answered_in: String,
81    /// A fast result.
82    pub fast: String,
83    /// A result that can be lived with.
84    pub usable: String,
85    /// A slow result, suggesting what to do.
86    pub slow: String,
87    /// When the model did not run.
88    pub did_not_run: String,
89}
90
91impl Default for ModelManagerLabels {
92    fn default() -> Self {
93        Self::from_texts(&Texts::new(Language::English))
94    }
95}
96
97impl ModelManagerLabels {
98    /// The texts of a catalog, in its language. What it does not have comes from English.
99    pub fn from_texts(texts: &Texts) -> Self {
100        Self {
101            picker: ModelPickerLabels::from_texts(texts),
102            understood: texts.get("ai.models.understood"),
103            not_measured: texts.get("ai.models.not_measured"),
104            license: texts.get("ai.models.license"),
105            memory: texts.get("ai.models.memory"),
106            decimal_comma: texts.char("format.decimal", '.') == ',',
107            network_error: texts.get("ai.models.network_error"),
108            disk_error: texts.get("ai.models.disk_error"),
109            wrong_file: texts.get("ai.models.wrong_file"),
110            answered_in: texts.get("ai.models.answered_in"),
111            fast: texts.get("ai.models.fast"),
112            usable: texts.get("ai.models.usable"),
113            slow: texts.get("ai.models.slow"),
114            did_not_run: texts.get("ai.models.did_not_run"),
115        }
116    }
117
118    /// The words in Brazilian Portuguese.
119    pub fn pt_br() -> Self {
120        Self::from_texts(&Texts::new(Language::BrazilianPortuguese))
121    }
122
123    fn details(&self, model: &LocalModel) -> String {
124        let gigabytes = |bytes: u64| {
125            let size = format!("{:.1} GB", bytes as f32 / GIGABYTE);
126            if self.decimal_comma { size.replace('.', ",") } else { size }
127        };
128        let understood =
129            model.understood.map_or_else(|| self.not_measured.clone(), |percent| format!("{percent}%"));
130        format!(
131            "{} · {}: {} · {}: {understood} · {}: {}",
132            gigabytes(model.bytes()),
133            self.memory,
134            gigabytes(model.memory),
135            self.understood,
136            self.license,
137            model.license
138        )
139    }
140
141    fn result(&self, outcome: Option<&Tested>) -> Option<String> {
142        let tested = outcome?;
143        let verdict = if tested.tokens_per_second >= FAST {
144            &self.fast
145        } else if tested.tokens_per_second >= USABLE {
146            &self.usable
147        } else {
148            &self.slow
149        };
150        let seconds = format!("{:.1} s", tested.seconds);
151        let seconds = if self.decimal_comma { seconds.replace('.', ",") } else { seconds };
152        Some(format!("{} {seconds}. {verdict}", self.answered_in))
153    }
154
155    fn error(&self, error: &ModelError) -> Option<String> {
156        match error {
157            ModelError::Cancelled => None,
158            ModelError::Network(_) => Some(self.network_error.clone()),
159            ModelError::WrongFile => Some(self.wrong_file.clone()),
160            _ => Some(self.disk_error.clone()),
161        }
162    }
163}
164
165/// A download in progress, shared with its thread.
166#[derive(Debug, Default)]
167struct Download {
168    received: AtomicU64,
169    cancelled: AtomicBool,
170    finished: Mutex<Option<Result<(), ModelError>>>,
171}
172
173/// How a speed test went.
174#[derive(Clone, Copy, Debug, PartialEq)]
175struct Tested {
176    seconds: f64,
177    tokens_per_second: f64,
178}
179
180/// A speed test in progress: what its thread leaves when it ends.
181type TestRun = Arc<Mutex<Option<Option<Tested>>>>;
182
183/// One model on offer.
184#[derive(Debug)]
185struct Entry {
186    model: LocalModel,
187    installed: bool,
188    download: Option<Arc<Download>>,
189    test: Option<TestRun>,
190    tested: Option<Tested>,
191    did_not_run: bool,
192    error: Option<ModelError>,
193}
194
195/// The models of a list that fit this machine, to download, try, use and remove. A
196/// model that does not fit is not offered, unless it is already installed and still
197/// fits in memory. Build it once and keep it: downloads belong to it.
198///
199/// ```no_run
200/// use ravenrs_ai::{CATALOG, ModelFolder, ModelManager};
201///
202/// let manager = ModelManager::new(CATALOG, ModelFolder::new("models"));
203/// ```
204#[derive(Clone, Debug)]
205pub struct ModelManager {
206    folder: ModelFolder,
207    entries: Rc<RefCell<Vec<Entry>>>,
208    in_use: Rc<RefCell<Option<(LocalModel, LocalAssistant)>>>,
209    changes: State<u64>,
210    labels: ModelManagerLabels,
211}
212
213impl ModelManager {
214    /// The models of `catalog` that fit this machine, kept in `folder`.
215    pub fn new(catalog: &[LocalModel], folder: ModelFolder) -> Self {
216        let machine = Machine::detect(folder.root());
217        Self::for_machine(catalog, folder, &machine)
218    }
219
220    /// The models of `catalog` that fit `machine`.
221    pub fn for_machine(catalog: &[LocalModel], folder: ModelFolder, machine: &Machine) -> Self {
222        let entries = catalog
223            .iter()
224            .filter_map(|model| {
225                let installed = folder.state(model) == ModelState::Installed;
226                model.fits(machine, installed).then_some(Entry {
227                    model: *model,
228                    installed,
229                    download: None,
230                    test: None,
231                    tested: None,
232                    did_not_run: false,
233                    error: None,
234                })
235            })
236            .collect::<Vec<Entry>>();
237        // What was in use last time, if it is still here and still fits.
238        let saved = std::fs::read_to_string(folder.root().join(IN_USE_FILE)).unwrap_or_default();
239        let in_use = entries
240            .iter()
241            .find(|entry| entry.installed && entry.model.id == saved.trim())
242            .map(|entry| (entry.model, LocalAssistant::new(folder.clone(), entry.model)));
243        Self {
244            folder,
245            entries: Rc::new(RefCell::new(entries)),
246            in_use: Rc::new(RefCell::new(in_use)),
247            changes: State::new(0),
248            labels: ModelManagerLabels::default(),
249        }
250    }
251
252    /// The words it shows.
253    #[must_use]
254    pub fn labels(mut self, labels: ModelManagerLabels) -> Self {
255        self.labels = labels;
256        self
257    }
258
259    /// The model in use, ready to be asked, if there is one.
260    pub fn assistant(&self) -> Option<LocalAssistant> {
261        self.in_use.borrow().as_ref().map(|(_, assistant)| assistant.clone())
262    }
263
264    /// The model in use, if any.
265    pub fn model_in_use(&self) -> Option<LocalModel> {
266        self.in_use.borrow().as_ref().map(|(model, _)| *model)
267    }
268
269    fn is_in_use(&self, model: &LocalModel) -> bool {
270        self.in_use.borrow().as_ref().is_some_and(|(used, _)| used.id == model.id)
271    }
272
273    /// Starts or stops using a model, and remembers it. Letting go of the assistant
274    /// frees its memory.
275    fn toggle_use(&self, index: usize) {
276        let Some(model) = self.entries.borrow().get(index).map(|entry| entry.model) else { return };
277        let next =
278            (!self.is_in_use(&model)).then(|| (model, LocalAssistant::new(self.folder.clone(), model)));
279        self.set_in_use(next);
280    }
281
282    fn set_in_use(&self, next: Option<(LocalModel, LocalAssistant)>) {
283        let saved = next.as_ref().map_or("", |(model, _)| model.id);
284        let _ = std::fs::create_dir_all(self.folder.root());
285        // Remembering is a convenience: without it, the next run starts with no model.
286        let _ = std::fs::write(self.folder.root().join(IN_USE_FILE), saved);
287        *self.in_use.borrow_mut() = next;
288    }
289
290    fn test(&self, index: usize, waker: UiWaker) {
291        let mut entries = self.entries.borrow_mut();
292        let Some(entry) = entries.get_mut(index) else { return };
293        if entry.test.is_some() {
294            return;
295        }
296        let run: TestRun = Arc::new(Mutex::new(None));
297        entry.test = Some(run.clone());
298        entry.tested = None;
299        entry.did_not_run = false;
300        let (model, folder) = (entry.model, self.folder.clone());
301        std::thread::spawn(move || {
302            let started = Instant::now();
303            if model.family == ModelFamily::Whisper {
304                let outcome = hearing_test(&folder, &model, started);
305                if let Ok(mut done) = run.lock() {
306                    *done = Some(outcome);
307                }
308                waker.wake();
309                return;
310            }
311            let outcome = LocalEngine::load(&folder, &model).ok().and_then(|mut engine| {
312                let generated = engine.generate(TEST_SYSTEM, TEST_USER, TEST_TOKENS, &mut |_| true).ok()?;
313                (generated.tokens > 0).then(|| Tested {
314                    seconds: started.elapsed().as_secs_f64(),
315                    tokens_per_second: generated.tokens_per_second(),
316                })
317            });
318            if let Ok(mut done) = run.lock() {
319                *done = Some(outcome);
320            }
321            waker.wake();
322        });
323    }
324
325    fn changed(&self) {
326        self.changes.update(|n| *n += 1);
327    }
328
329    fn start(&self, index: usize, waker: UiWaker) {
330        let mut entries = self.entries.borrow_mut();
331        let Some(entry) = entries.get_mut(index) else { return };
332        if entry.download.is_some() {
333            return;
334        }
335        let download = Arc::new(Download::default());
336        entry.download = Some(download.clone());
337        entry.error = None;
338        let (model, folder) = (entry.model, self.folder.clone());
339        std::thread::spawn(move || {
340            let step = (model.bytes() / STEPS).max(1);
341            let told = AtomicU64::new(0);
342            let progress = |received: u64, _total: u64| {
343                download.received.store(received, Ordering::Relaxed);
344                if received / step != told.swap(received / step, Ordering::Relaxed) {
345                    waker.wake();
346                }
347            };
348            let result = folder.download(&model, &progress, &download.cancelled);
349            if let Ok(mut finished) = download.finished.lock() {
350                *finished = Some(result);
351            }
352            waker.wake();
353        });
354    }
355
356    fn cancel(&self, index: usize) {
357        if let Some(download) = self.entries.borrow().get(index).and_then(|entry| entry.download.as_ref()) {
358            download.cancelled.store(true, Ordering::Relaxed);
359        }
360    }
361
362    fn remove(&self, index: usize) {
363        if let Some(entry) = self.entries.borrow_mut().get_mut(index) {
364            if self.is_in_use(&entry.model) {
365                self.set_in_use(None);
366            }
367            match self.folder.remove(&entry.model) {
368                Ok(()) => {
369                    entry.installed = false;
370                    entry.tested = None;
371                }
372                Err(error) => entry.error = Some(error),
373            }
374        }
375    }
376
377    /// Takes in the downloads that ended, and says how each model stands.
378    fn choices(&self) -> Vec<ModelChoice> {
379        let mut entries = self.entries.borrow_mut();
380        entries
381            .iter_mut()
382            .map(|entry| {
383                let ended = entry
384                    .download
385                    .as_ref()
386                    .and_then(|download| download.finished.lock().ok().and_then(|mut done| done.take()));
387                if let Some(result) = ended {
388                    entry.download = None;
389                    match result {
390                        Ok(()) => entry.installed = true,
391                        Err(ModelError::Cancelled) => {}
392                        Err(error) => entry.error = Some(error),
393                    }
394                }
395                let test_ended =
396                    entry.test.as_ref().and_then(|run| run.lock().ok().and_then(|mut done| done.take()));
397                if let Some(outcome) = test_ended {
398                    entry.test = None;
399                    entry.tested = outcome;
400                    entry.did_not_run = outcome.is_none();
401                }
402                let status = if entry.test.is_some() {
403                    ModelStatus::Testing
404                } else if let Some(download) = &entry.download {
405                    let received = download.received.load(Ordering::Relaxed);
406                    ModelStatus::Downloading(received as f32 / entry.model.bytes().max(1) as f32)
407                } else if entry.installed {
408                    ModelStatus::Installed
409                } else {
410                    ModelStatus::Available
411                };
412                ModelChoice {
413                    name: entry.model.name.to_owned(),
414                    details: self.labels.details(&entry.model),
415                    status,
416                    in_use: self.is_in_use(&entry.model),
417                    result: self.labels.result(entry.tested.as_ref()),
418                    error: if entry.did_not_run {
419                        Some(self.labels.did_not_run.clone())
420                    } else {
421                        entry.error.as_ref().and_then(|error| self.labels.error(error))
422                    },
423                }
424            })
425            .collect()
426    }
427}
428
429impl View for ModelManager {
430    fn build(&self, cx: &mut BuildCx<'_>) -> NodeId {
431        // Read, so that a click rebuilds this view. A thread wakes it the other way.
432        let _ = self.changes.get();
433        let waker = cx.ui_waker();
434        let (starting, cancelling, removing) = (self.clone(), self.clone(), self.clone());
435        let (testing, using, test_waker) = (self.clone(), self.clone(), waker.clone());
436        let picker = ModelPicker::new(self.choices())
437            .labels(self.labels.picker.clone())
438            .on_download(move |index| {
439                starting.start(index, waker.clone());
440                starting.changed();
441            })
442            .on_cancel(move |index| {
443                cancelling.cancel(index);
444                cancelling.changed();
445            })
446            .on_test(move |index| {
447                testing.test(index, test_waker.clone());
448                testing.changed();
449            })
450            .on_use(move |index| {
451                using.toggle_use(index);
452                using.changed();
453            })
454            .on_remove(move |index| {
455                removing.remove(index);
456                removing.changed();
457            });
458        cx.child(&picker)
459    }
460}
ravenrs_ai/manager.rs

ravenrs_ai/
manager.rs
