use super::types::TensorBoardError;
pub type Result<T> = std::result::Result<T, TensorBoardError>;
pub(crate) const EVENT_FILE_VERSION: &str = "brain.Event:2";
#[cfg(test)]
mod tests {
use crate::tensorboard::{PluginConfig, SummaryWriter, TensorBoardWriter};
use serde_json::json;
use std::collections::HashMap;
use tempfile::TempDir;
use torsh_tensor::Tensor;
#[test]
fn test_tensorboard_writer() {
let temp_dir = TempDir::new().expect("Temp Dir should succeed");
let mut writer =
TensorBoardWriter::new(temp_dir.path(), None).expect("operation should succeed");
writer
.log_scalar("loss", 0.5, Some(0))
.expect("operation should succeed");
writer
.log_scalar("loss", 0.3, Some(1))
.expect("operation should succeed");
writer
.log_scalar("loss", 0.1, Some(2))
.expect("operation should succeed");
let mut metrics = HashMap::new();
metrics.insert("accuracy".to_string(), 0.95);
metrics.insert("f1_score".to_string(), 0.92);
writer
.log_scalars(metrics, Some(3))
.expect("operation should succeed");
assert_eq!(writer.get_step(), 3);
}
#[test]
fn test_summary_writer() {
let temp_dir = TempDir::new().expect("Temp Dir should succeed");
let mut writer = SummaryWriter::new(temp_dir.path()).expect("operation should succeed");
writer
.add_scalar("train/loss", 0.5, Some(0))
.expect("operation should succeed");
writer
.add_scalar("train/loss", 0.3, Some(1))
.expect("operation should succeed");
let mut scalars = HashMap::new();
scalars.insert("accuracy".to_string(), 0.95);
scalars.insert("precision".to_string(), 0.93);
writer
.add_scalars("eval", scalars, Some(2))
.expect("operation should succeed");
writer.flush().expect("flush should succeed");
}
#[test]
fn test_image_logging() {
let temp_dir = TempDir::new().expect("Temp Dir should succeed");
let mut writer = SummaryWriter::new(temp_dir.path()).expect("operation should succeed");
let shape = vec![3, 32, 32];
let data: Vec<f32> = (0..3072).map(|i| (i as f32) / 3072.0).collect();
let image_tensor = Tensor::from_vec(data, &shape).expect("Tensor should succeed");
writer
.add_image("test_image", &image_tensor, Some(0), "CHW")
.expect("operation should succeed");
writer.flush().expect("flush should succeed");
let image_file = temp_dir.path().join("test_image_step_0.json");
assert!(image_file.exists());
}
#[test]
fn test_audio_logging() {
let temp_dir = TempDir::new().expect("Temp Dir should succeed");
let mut writer = SummaryWriter::new(temp_dir.path()).expect("operation should succeed");
let sample_rate = 44100;
let duration = 1.0;
let samples = (sample_rate as f32 * duration) as usize;
let freq = 440.0;
let audio_data: Vec<f32> = (0..samples)
.map(|i| {
(2.0 * std::f32::consts::PI * freq * i as f32 / sample_rate as f32).sin() * 0.5
})
.collect();
let shape = vec![samples];
let audio_tensor = Tensor::from_vec(audio_data, &shape).expect("Tensor should succeed");
writer
.add_audio("test_audio", &audio_tensor, sample_rate, Some(0))
.expect("operation should succeed");
writer.flush().expect("flush should succeed");
let audio_file = temp_dir.path().join("test_audio_step_0_audio.json");
assert!(audio_file.exists());
}
#[test]
fn test_embedding_logging() {
let temp_dir = TempDir::new().expect("Temp Dir should succeed");
let mut writer = SummaryWriter::new(temp_dir.path()).expect("operation should succeed");
let shape = vec![10, 128];
let data: Vec<f32> = (0..1280).map(|i| (i as f32) / 1280.0 - 0.5).collect();
let embedding_tensor = Tensor::from_vec(data, &shape).expect("Tensor should succeed");
let metadata = Some(vec![
"item_0".to_string(),
"item_1".to_string(),
"item_2".to_string(),
"item_3".to_string(),
"item_4".to_string(),
"item_5".to_string(),
"item_6".to_string(),
"item_7".to_string(),
"item_8".to_string(),
"item_9".to_string(),
]);
writer
.add_embedding(
&embedding_tensor,
metadata,
None,
Some(0),
"test_embeddings",
)
.expect("operation should succeed");
writer.flush().expect("flush should succeed");
let embedding_file = temp_dir
.path()
.join("test_embeddings_step_0_embeddings.json");
assert!(embedding_file.exists());
}
#[test]
fn test_plugin_system() {
let temp_dir = TempDir::new().expect("Temp Dir should succeed");
let mut writer = SummaryWriter::new(temp_dir.path()).expect("operation should succeed");
let mut plugin_config = HashMap::new();
plugin_config.insert("threshold".to_string(), json!(0.5));
plugin_config.insert("color_scheme".to_string(), json!("viridis"));
let plugin = PluginConfig {
name: "custom_visualizer".to_string(),
version: "1.0.0".to_string(),
entry_point: "visualizer.js".to_string(),
config: plugin_config,
enabled: true,
};
writer
.install_plugin(plugin)
.expect("plugin installation should succeed");
writer.flush().expect("flush should succeed");
let plugin_file = temp_dir
.path()
.join("plugins")
.join("custom_visualizer.json");
assert!(plugin_file.exists());
let manifest_file = temp_dir.path().join("plugins").join("manifest.json");
assert!(manifest_file.exists());
}
#[test]
fn test_custom_dashboard() {
let temp_dir = TempDir::new().expect("Temp Dir should succeed");
let mut writer = SummaryWriter::new(temp_dir.path()).expect("operation should succeed");
let layout = json!(
{ "type" : "grid", "columns" : 2, "rows" : 2, "widgets" : [{ "type" :
"scalar", "tag" : "loss", "position" : [0, 0] }, { "type" : "scalar", "tag" :
"accuracy", "position" : [0, 1] }, { "type" : "image", "tag" : "samples",
"position" : [1, 0] }, { "type" : "histogram", "tag" : "weights", "position"
: [1, 1] }] }
);
writer
.create_dashboard("training_dashboard", layout)
.expect("operation should succeed");
writer.flush().expect("flush should succeed");
let dashboard_file = temp_dir
.path()
.join("dashboards")
.join("training_dashboard.json");
assert!(dashboard_file.exists());
}
}