use std::rc::Rc; use gpui::{ App, ElementId, FontWeight, IntoElement, ParentElement, RenderOnce, SharedString, Styled, Window, div, }; use crate::{ buttons::{Button, ButtonVariant}, charts::{LineChart, Series}, data_display::{Badge, Tone}, forms::Run, motion::ProgressBar, theme::{ActiveTheme, ControlSize, Radius, TextSize}, typography::{Ellipsis, format::percent, tabular}, }; /// Where a fine-tuning job stands. #[derive(Clone, Debug, PartialEq)] pub enum TunePhase { /// Waiting, with this many ahead of it. Queued(usize), /// Training: the epoch under way from one, of how many, and the share of it done. Running(usize, usize, f32), Done, Failed(SharedString), Canceled, } impl TunePhase { fn badge(&self) -> (&'static str, Tone) { match self { Self::Queued(_) => ("Queued", Tone::Neutral), Self::Running(..) => ("Training", Tone::Info), Self::Done => ("Done", Tone::Success), Self::Failed(_) => ("Failed", Tone::Danger), Self::Canceled => ("Canceled", Tone::Neutral), } } } /// A running job's bar, whole epochs and the share of the one under way, and the line under it. pub(crate) fn progress(epoch: usize, epochs: usize, share: f32) -> (f32, String) { let whole = ((epoch - 1) as f32 + share) / epochs as f32; let line = format!( "Epoch {epoch} of {epochs} · {}", percent(f64::from(share), 0, false) ); (whole, line) } /// A fine-tuning job: its name, the base model and data it learns from, where it stands, how far it has come by epoch, and its loss on training and held-out data. Cancel while it waits or trains; Open once done. #[derive(IntoElement)] pub struct FineTuneJobCard { id: ElementId, name: SharedString, base: SharedString, data: SharedString, phase: TunePhase, train: Vec, held: Vec, on_cancel: Option, on_open: Option, } impl FineTuneJobCard { pub fn new( id: impl Into, name: impl Into, base: impl Into, data: impl Into, phase: TunePhase, ) -> Self { if let TunePhase::Running(epoch, epochs, share) = phase { assert!( (1..=epochs).contains(&epoch) && (0.0..=1.0).contains(&share), "epoch {epoch} of {epochs} at {share}" ); } Self { id: id.into(), name: name.into(), base: base.into(), data: data.into(), phase, train: Vec::new(), held: Vec::new(), on_cancel: None, on_open: None, } } /// Loss at each checkpoint, on training data and on held-out data, one value each. pub fn loss(mut self, train: impl Into>, held: impl Into>) -> Self { let (train, held) = (train.into(), held.into()); assert_eq!( train.len(), held.len(), "a loss for each checkpoint on both sets" ); self.train = train; self.held = held; self } pub fn on_cancel(mut self, handler: impl Fn(&mut Window, &mut App) + 'static) -> Self { self.on_cancel = Some(Rc::new(handler)); self } pub fn on_open(mut self, handler: impl Fn(&mut Window, &mut App) + 'static) -> Self { self.on_open = Some(Rc::new(handler)); self } } impl RenderOnce for FineTuneJobCard { fn render(self, _: &mut Window, cx: &mut App) -> impl IntoElement { let theme = cx.theme(); let colors = theme.colors.clone(); let (word, tone) = self.phase.badge(); let live = matches!(self.phase, TunePhase::Queued(_) | TunePhase::Running(..)); let action = |key: &'static str, label: &'static str, run: Option| { run.map(|run| { Button::new((self.id.clone(), key), label) .variant(ButtonVariant::Ghost) .size(ControlSize::Sm) .on_click(move |_, window, cx| { log::info!("fine-tune job: {key}"); run(window, cx) }) }) }; let actions = match self.phase { TunePhase::Done => action("open", "Open", self.on_open), _ if live => action("cancel", "Cancel", self.on_cancel), _ => None, }; let progress = match &self.phase { TunePhase::Running(epoch, epochs, share) => { let (whole, line) = progress(*epoch, *epochs, *share); Some( div() .flex() .flex_col() .gap_1p5() .child(match *epochs { 1 => ProgressBar::new((self.id.clone(), "epochs"), whole), _ => ProgressBar::new((self.id.clone(), "epochs"), whole) .segments(*epochs), }) .child( tabular(div()) .text_size(theme.text_size(TextSize::Xs)) .text_color(colors.fg_muted) .child(line), ) .into_any_element(), ) } TunePhase::Queued(ahead) => Some( div() .text_size(theme.text_size(TextSize::Xs)) .text_color(colors.fg_muted) .child(format!("{ahead} ahead")) .into_any_element(), ), TunePhase::Failed(reason) => Some( div() .text_size(theme.text_size(TextSize::Xs)) .text_color(colors.danger) .child(reason.clone()) .into_any_element(), ), _ => None, }; let chart = (self.train.len() > 1).then(|| { let steps: Vec = (1..=self.train.len()) .map(|step| step.to_string().into()) .collect(); LineChart::new((self.id.clone(), "loss"), steps) .series(Series::new("Training", self.train)) .series(Series::new("Held out", self.held)) .format(|value| format!("{value:.2}")) .h(theme.generative().loss) }); div() .flex() .flex_col() .gap_3() .p_4() .rounded(theme.radius(Radius::Lg)) .border_1() .border_color(colors.border) .bg(colors.surface) .child( div() .flex() .flex_wrap() .items_start() .gap_x_3() .gap_y_2() .child( div() .flex_1() .min_w(theme.label_width()) .flex() .flex_col() .child( div() .text_size(theme.text_size(TextSize::Sm)) .font_weight(FontWeight::MEDIUM) .text_color(colors.fg) .child(Ellipsis::new(self.name)), ) .child( div() .text_size(theme.text_size(TextSize::Xs)) .text_color(colors.fg_muted) .child(Ellipsis::new(format!("{} · {}", self.base, self.data))), ), ) .child( div() .flex_none() .flex() .items_center() .gap_2() .child(Badge::new(word).tone(tone)) .children(actions), ), ) .children(progress) .children(chart) } } #[cfg(test)] mod tests { use super::progress; #[test] fn the_bar_counts_whole_epochs_and_the_share_of_the_one_under_way() { let (whole, line) = progress(2, 3, 0.64); assert!((whole - 0.5467).abs() < 1e-4, "{whole}"); assert_eq!(line, "Epoch 2 of 3 · 64%"); assert_eq!(progress(1, 1, 0.5), (0.5, "Epoch 1 of 1 · 50%".to_string())); assert_eq!(progress(3, 3, 1.0).0, 1.0); } }