thrust_rl/inference/
weights.rs

1//! Weight export/import utilities for trained models
2
3use std::{
4    fs::File,
5    io::{Read, Write},
6    path::Path,
7};
8
9use anyhow::Result;
10
11use super::ExportedModel;
12
13impl ExportedModel {
14    /// Save model to JSON file
15    pub fn save_json<P: AsRef<Path>>(&self, path: P) -> Result<()> {
16        let json = serde_json::to_string_pretty(self)?;
17        let mut file = File::create(path)?;
18        file.write_all(json.as_bytes())?;
19        Ok(())
20    }
21
22    /// Load model from JSON file
23    pub fn load_json<P: AsRef<Path>>(path: P) -> Result<Self> {
24        let mut file = File::open(path)?;
25        let mut contents = String::new();
26        file.read_to_string(&mut contents)?;
27        let model = serde_json::from_str(&contents)?;
28        Ok(model)
29    }
30}
31
32#[cfg(test)]
33mod tests {
34    use tempfile::NamedTempFile;
35
36    use super::*;
37    use crate::inference::LayerWeights;
38
39    fn create_test_model() -> ExportedModel {
40        let feature_weights = LayerWeights::new(vec![1.0, 0.0, 0.0, 1.0], vec![0.0, 0.0], 2, 2);
41        let policy_weights = LayerWeights::new(vec![1.0, -1.0, -1.0, 1.0], vec![0.0, 0.0], 2, 2);
42        ExportedModel::new(2, 2, 2, feature_weights, policy_weights, None)
43    }
44
45    #[test]
46    fn test_json_roundtrip() -> Result<()> {
47        let model = create_test_model();
48        let temp_file = NamedTempFile::new()?;
49
50        model.save_json(temp_file.path())?;
51        let loaded = ExportedModel::load_json(temp_file.path())?;
52
53        assert_eq!(model.input_dim, loaded.input_dim);
54        assert_eq!(model.output_dim, loaded.output_dim);
55        assert_eq!(model.hidden_dim, loaded.hidden_dim);
56
57        Ok(())
58    }
59}
thrust_rl/inference/weights.rs

thrust_rl/inference/
weights.rs
