Skip to main content

yscv_model/
training_log.rs

1use std::collections::HashMap;
2
3/// Records per-epoch training metrics.
4#[derive(Debug, Clone, Default)]
5pub struct TrainingLog {
6    entries: Vec<HashMap<String, f32>>,
7}
8
9impl TrainingLog {
10    /// Create a new empty training log.
11    pub fn new() -> Self {
12        Self::default()
13    }
14
15    /// Record one epoch's metrics.
16    pub fn log_epoch(&mut self, metrics: HashMap<String, f32>) {
17        self.entries.push(metrics);
18    }
19
20    /// All logged entries.
21    pub fn entries(&self) -> &[HashMap<String, f32>] {
22        &self.entries
23    }
24
25    /// Extract the history of a single metric across all epochs.
26    ///
27    /// Epochs that did not record the metric are skipped.
28    pub fn get_metric_history(&self, name: &str) -> Vec<f32> {
29        self.entries
30            .iter()
31            .filter_map(|e| e.get(name).copied())
32            .collect()
33    }
34
35    /// Number of epochs logged so far.
36    pub fn num_epochs(&self) -> usize {
37        self.entries.len()
38    }
39
40    /// Export the log as a CSV string.
41    ///
42    /// The header row contains the union of all metric names across every epoch,
43    /// sorted alphabetically for deterministic output. Missing values are empty.
44    pub fn to_csv(&self) -> String {
45        if self.entries.is_empty() {
46            return String::new();
47        }
48
49        // Collect all unique keys, sorted.
50        let mut keys: Vec<String> = self
51            .entries
52            .iter()
53            .flat_map(|e| e.keys().cloned())
54            .collect();
55        keys.sort();
56        keys.dedup();
57
58        let mut out = keys.join(",");
59        out.push('\n');
60
61        for entry in &self.entries {
62            let row: Vec<String> = keys
63                .iter()
64                .map(|k| match entry.get(k) {
65                    Some(v) => v.to_string(),
66                    None => String::new(),
67                })
68                .collect();
69            out.push_str(&row.join(","));
70            out.push('\n');
71        }
72
73        out
74    }
75
76    /// Export the log as JSONL (one JSON object per epoch).
77    ///
78    /// Compatible with TensorBoard's `--logdir` when saved as `.jsonl`.
79    /// Each line: `{"epoch": N, "loss": 0.5, "accuracy": 0.9, ...}`
80    pub fn to_jsonl(&self) -> String {
81        let mut out = String::new();
82        for (epoch, entry) in self.entries.iter().enumerate() {
83            out.push_str("{\"epoch\":");
84            out.push_str(&epoch.to_string());
85            let mut keys: Vec<&String> = entry.keys().collect();
86            keys.sort();
87            for key in keys {
88                let val = entry[key];
89                out.push_str(",\"");
90                out.push_str(key);
91                out.push_str("\":");
92                if val.is_finite() {
93                    out.push_str(&format!("{val:.6}"));
94                } else {
95                    out.push_str("null");
96                }
97            }
98            out.push_str("}\n");
99        }
100        out
101    }
102
103    /// Write the log to a JSONL file (append-friendly).
104    pub fn save_jsonl(&self, path: &std::path::Path) -> std::io::Result<()> {
105        std::fs::write(path, self.to_jsonl())
106    }
107
108    /// Append a single epoch's metrics to a JSONL file (streaming mode).
109    pub fn append_epoch_jsonl(
110        path: &std::path::Path,
111        epoch: usize,
112        metrics: &HashMap<String, f32>,
113    ) -> std::io::Result<()> {
114        use std::io::Write;
115        let mut file = std::fs::OpenOptions::new()
116            .create(true)
117            .append(true)
118            .open(path)?;
119        write!(file, "{{\"epoch\":{epoch}")?;
120        let mut keys: Vec<&String> = metrics.keys().collect();
121        keys.sort();
122        for key in keys {
123            let val = metrics[key];
124            if val.is_finite() {
125                write!(file, ",\"{key}\":{val:.6}")?;
126            } else {
127                write!(file, ",\"{key}\":null")?;
128            }
129        }
130        writeln!(file, "}}")?;
131        Ok(())
132    }
133}