overbrainer 0.2.0

Distill a large LLM into a smaller open-weights model: generate questions, collect answers, fine-tune, visualize.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

//! The metrics plugin shipped with each run, and the lines it writes.

use serde::{Deserialize, Serialize};

/// Source of the Axolotl plugin that writes `metrics.jsonl`. It is uploaded as
/// [`PLUGIN_FILE`] and found through `PYTHONPATH`.
pub const METRICS_PLUGIN: &str = include_str!("overbrainer_metrics.py");

/// File name of the plugin module, inside the run's `plugin/` directory.
pub const PLUGIN_FILE: &str = "overbrainer_metrics.py";

/// The plugin as Axolotl's `plugins` list names it: module, then class.
pub const PLUGIN_CLASS: &str = "overbrainer_metrics.OverbrainerMetricsPlugin";

/// Env variable naming the file the plugin appends to.
pub const METRICS_ENV: &str = "OVERBRAINER_METRICS";

/// One training log line: the trainer's step and the values it logged.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TrainMetric {
    /// Unix time of the log, in seconds.
    pub time: f64,
    /// Optimizer steps done.
    pub step: u64,
    /// Fractional epoch, when known.
    #[serde(default)]
    pub epoch: Option<f64>,
    /// Total steps of the run, when known.
    #[serde(default)]
    pub max_steps: Option<u64>,
    /// Training loss.
    #[serde(default)]
    pub loss: Option<f64>,
    /// Loss on `data/eval.jsonl`, on evaluation lines.
    #[serde(default)]
    pub eval_loss: Option<f64>,
    /// Learning rate.
    #[serde(default)]
    pub learning_rate: Option<f64>,
    /// Gradient norm.
    #[serde(default)]
    pub grad_norm: Option<f64>,
}

/// A line of `metrics.jsonl`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "event", rename_all = "lowercase")]
pub enum MetricLine {
    /// Training started: the plugin is loaded.
    Begin {
        /// Unix time, in seconds.
        time: f64,
        /// Total steps of the run, when known.
        #[serde(default)]
        max_steps: Option<u64>,
    },
    /// A training or evaluation log.
    Log(TrainMetric),
}

/// Parses one line of `metrics.jsonl`.
///
/// # Errors
///
/// Returns the JSON error when the line is not a metrics record.
pub fn parse_line(line: &str) -> Result<MetricLine, serde_json::Error> {
    serde_json::from_str(line)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn log_lines_parse_with_missing_values() -> Result<(), serde_json::Error> {
        let line = r#"{"event": "log", "time": 1.5, "step": 3, "epoch": 0.25, "max_steps": 12, "loss": 1.2, "grad_norm": null, "extra": 1}"#;
        assert_eq!(
            parse_line(line)?,
            MetricLine::Log(TrainMetric {
                time: 1.5,
                step: 3,
                epoch: Some(0.25),
                max_steps: Some(12),
                loss: Some(1.2),
                eval_loss: None,
                learning_rate: None,
                grad_norm: None,
            })
        );
        Ok(())
    }

    #[test]
    fn begin_lines_parse() -> Result<(), serde_json::Error> {
        assert_eq!(
            parse_line(r#"{"event": "begin", "time": 1.0, "max_steps": 40}"#)?,
            MetricLine::Begin {
                time: 1.0,
                max_steps: Some(40)
            }
        );
        Ok(())
    }

    #[test]
    fn other_lines_are_errors() {
        assert!(parse_line(r#"{"event": "other"}"#).is_err());
        assert!(parse_line("not json").is_err());
    }
}