trustformers_debug/
netron_export.rs

1//! Netron export functionality for model visualization
2//!
3//! This module provides tools to export TrustformeRS models to formats compatible with
4//! Netron (<https://netron.app/>), a powerful neural network visualizer.
5
6use anyhow::Result;
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9use std::fs;
10use std::path::Path;
11
12/// ONNX-like model representation for Netron visualization
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct NetronModel {
15    /// Model metadata
16    pub metadata: ModelMetadata,
17    /// Graph definition
18    pub graph: ModelGraph,
19    /// Model version
20    pub version: String,
21}
22
23/// Model metadata
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct ModelMetadata {
26    /// Model name
27    pub name: String,
28    /// Model description
29    pub description: String,
30    /// Model author
31    pub author: Option<String>,
32    /// Model version
33    pub version: Option<String>,
34    /// License information
35    pub license: Option<String>,
36    /// Additional properties
37    pub properties: HashMap<String, String>,
38}
39
40/// Model graph containing nodes and edges
41#[derive(Debug, Clone, Serialize, Deserialize)]
42pub struct ModelGraph {
43    /// Graph name
44    pub name: String,
45    /// Input tensors
46    pub inputs: Vec<TensorInfo>,
47    /// Output tensors
48    pub outputs: Vec<TensorInfo>,
49    /// Graph nodes (layers/operations)
50    pub nodes: Vec<GraphNode>,
51    /// Initializers (weights and biases)
52    pub initializers: Vec<TensorData>,
53}
54
55/// Tensor information
56#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct TensorInfo {
58    /// Tensor name
59    pub name: String,
60    /// Data type (e.g., "float32", "int64")
61    pub dtype: String,
62    /// Tensor shape
63    pub shape: Vec<i64>,
64    /// Optional documentation
65    pub doc_string: Option<String>,
66}
67
68/// Graph node representing a layer or operation
69#[derive(Debug, Clone, Serialize, Deserialize)]
70pub struct GraphNode {
71    /// Node name
72    pub name: String,
73    /// Operation type (e.g., "Linear", "Conv2d", "Softmax")
74    pub op_type: String,
75    /// Input tensor names
76    pub inputs: Vec<String>,
77    /// Output tensor names
78    pub outputs: Vec<String>,
79    /// Node attributes
80    pub attributes: HashMap<String, AttributeValue>,
81    /// Optional documentation
82    pub doc_string: Option<String>,
83}
84
85/// Attribute value types
86#[derive(Debug, Clone, Serialize, Deserialize)]
87#[serde(untagged)]
88pub enum AttributeValue {
89    /// Integer value
90    Int(i64),
91    /// Float value
92    Float(f64),
93    /// String value
94    String(String),
95    /// Boolean value
96    Bool(bool),
97    /// Array of integers
98    Ints(Vec<i64>),
99    /// Array of floats
100    Floats(Vec<f64>),
101    /// Array of strings
102    Strings(Vec<String>),
103}
104
105/// Tensor data for weights and biases
106#[derive(Debug, Clone, Serialize, Deserialize)]
107pub struct TensorData {
108    /// Tensor name
109    pub name: String,
110    /// Data type
111    pub dtype: String,
112    /// Tensor shape
113    pub shape: Vec<i64>,
114    /// Raw data (encoded as base64 for binary data)
115    #[serde(skip_serializing_if = "Option::is_none")]
116    pub data: Option<Vec<f32>>,
117    /// Data location (for external data)
118    #[serde(skip_serializing_if = "Option::is_none")]
119    pub data_location: Option<String>,
120}
121
122/// Netron exporter for model visualization
123pub struct NetronExporter {
124    model: NetronModel,
125    output_format: ExportFormat,
126}
127
128/// Export format options
129#[derive(Debug, Clone, Copy, PartialEq, Eq)]
130pub enum ExportFormat {
131    /// JSON format (human-readable)
132    Json,
133    /// Real ONNX protobuf.
134    ///
135    /// **Not implemented**: [`NetronExporter::export`] returns a structured
136    /// error for this variant. It is kept in the enum so the intent stays
137    /// expressible (and so selecting it is a compile-time-visible choice rather
138    /// than a silent fallback), but nothing in this crate can encode the ONNX
139    /// `ModelProto` protobuf schema.
140    Onnx,
141}
142
143impl NetronExporter {
144    /// Create a new Netron exporter
145    ///
146    /// # Arguments
147    ///
148    /// * `model_name` - Name of the model
149    /// * `description` - Model description
150    ///
151    /// # Example
152    ///
153    /// ```
154    /// use trustformers_debug::NetronExporter;
155    ///
156    /// let exporter = NetronExporter::new("bert-base", "BERT base model");
157    /// ```
158    pub fn new(model_name: &str, description: &str) -> Self {
159        let metadata = ModelMetadata {
160            name: model_name.to_string(),
161            description: description.to_string(),
162            author: None,
163            version: None,
164            license: None,
165            properties: HashMap::new(),
166        };
167
168        let graph = ModelGraph {
169            name: format!("{}_graph", model_name),
170            inputs: Vec::new(),
171            outputs: Vec::new(),
172            nodes: Vec::new(),
173            initializers: Vec::new(),
174        };
175
176        let model = NetronModel {
177            metadata,
178            graph,
179            version: "1.0".to_string(),
180        };
181
182        Self {
183            model,
184            output_format: ExportFormat::Json,
185        }
186    }
187
188    /// Set the export format
189    pub fn with_format(mut self, format: ExportFormat) -> Self {
190        self.output_format = format;
191        self
192    }
193
194    /// Set model metadata
195    pub fn set_metadata(&mut self, metadata: ModelMetadata) {
196        self.model.metadata = metadata;
197    }
198
199    /// Add model author
200    pub fn set_author(&mut self, author: &str) {
201        self.model.metadata.author = Some(author.to_string());
202    }
203
204    /// Add model version
205    pub fn set_version(&mut self, version: &str) {
206        self.model.metadata.version = Some(version.to_string());
207    }
208
209    /// Add a custom property to metadata
210    pub fn add_property(&mut self, key: &str, value: &str) {
211        self.model.metadata.properties.insert(key.to_string(), value.to_string());
212    }
213
214    /// Add an input tensor
215    pub fn add_input(&mut self, name: &str, dtype: &str, shape: Vec<i64>) {
216        self.model.graph.inputs.push(TensorInfo {
217            name: name.to_string(),
218            dtype: dtype.to_string(),
219            shape,
220            doc_string: None,
221        });
222    }
223
224    /// Add an output tensor
225    pub fn add_output(&mut self, name: &str, dtype: &str, shape: Vec<i64>) {
226        self.model.graph.outputs.push(TensorInfo {
227            name: name.to_string(),
228            dtype: dtype.to_string(),
229            shape,
230            doc_string: None,
231        });
232    }
233
234    /// Add a graph node (layer/operation)
235    ///
236    /// # Example
237    ///
238    /// ```
239    /// # use trustformers_debug::NetronExporter;
240    /// # use std::collections::HashMap;
241    /// let mut exporter = NetronExporter::new("model", "test model");
242    ///
243    /// let mut attrs = HashMap::new();
244    /// attrs.insert("in_features".to_string(),
245    ///              trustformers_debug::netron_export::AttributeValue::Int(768));
246    /// attrs.insert("out_features".to_string(),
247    ///              trustformers_debug::netron_export::AttributeValue::Int(3072));
248    ///
249    /// exporter.add_node(
250    ///     "fc1",
251    ///     "Linear",
252    ///     vec!["input".to_string()],
253    ///     vec!["hidden".to_string()],
254    ///     attrs,
255    /// );
256    /// ```
257    pub fn add_node(
258        &mut self,
259        name: &str,
260        op_type: &str,
261        inputs: Vec<String>,
262        outputs: Vec<String>,
263        attributes: HashMap<String, AttributeValue>,
264    ) {
265        self.model.graph.nodes.push(GraphNode {
266            name: name.to_string(),
267            op_type: op_type.to_string(),
268            inputs,
269            outputs,
270            attributes,
271            doc_string: None,
272        });
273    }
274
275    /// Add a node with documentation
276    pub fn add_node_with_doc(
277        &mut self,
278        name: &str,
279        op_type: &str,
280        inputs: Vec<String>,
281        outputs: Vec<String>,
282        attributes: HashMap<String, AttributeValue>,
283        doc_string: &str,
284    ) {
285        self.model.graph.nodes.push(GraphNode {
286            name: name.to_string(),
287            op_type: op_type.to_string(),
288            inputs,
289            outputs,
290            attributes,
291            doc_string: Some(doc_string.to_string()),
292        });
293    }
294
295    /// Add tensor data (weights/biases)
296    pub fn add_tensor_data(
297        &mut self,
298        name: &str,
299        dtype: &str,
300        shape: Vec<i64>,
301        data: Option<Vec<f32>>,
302    ) {
303        self.model.graph.initializers.push(TensorData {
304            name: name.to_string(),
305            dtype: dtype.to_string(),
306            shape,
307            data,
308            data_location: None,
309        });
310    }
311
312    /// Export the model to a file
313    ///
314    /// # Arguments
315    ///
316    /// * `path` - Output file path
317    ///
318    /// # Example
319    ///
320    /// ```no_run
321    /// # use trustformers_debug::NetronExporter;
322    /// # let exporter = NetronExporter::new("model", "test");
323    /// exporter.export("model.json").unwrap();
324    /// ```
325    pub fn export<P: AsRef<Path>>(&self, path: P) -> Result<()> {
326        let path = path.as_ref();
327
328        // Create parent directory if needed
329        if let Some(parent) = path.parent() {
330            fs::create_dir_all(parent)?;
331        }
332
333        match self.output_format {
334            ExportFormat::Json => {
335                let json = serde_json::to_string_pretty(&self.model)?;
336                fs::write(path, json)?;
337            },
338            ExportFormat::Onnx => {
339                // Refuse rather than write JSON bytes under an ONNX name. The
340                // previous implementation serialised `self.model` as JSON and
341                // wrote it to the caller's `.onnx` path, producing a file that
342                // no ONNX runtime, checker or Netron ONNX importer can open,
343                // while reporting success.
344                return Err(anyhow::anyhow!(
345                    "ExportFormat::Onnx is not implemented: writing a real .onnx file needs a \
346                     protobuf encoder for the ONNX ModelProto schema, which trustformers-debug \
347                     does not link. Use ExportFormat::Json -- Netron opens the JSON graph \
348                     directly."
349                ));
350            },
351        }
352
353        Ok(())
354    }
355
356    /// Get a reference to the model
357    pub fn model(&self) -> &NetronModel {
358        &self.model
359    }
360
361    /// Get a mutable reference to the model
362    pub fn model_mut(&mut self) -> &mut NetronModel {
363        &mut self.model
364    }
365
366    /// Export model to a string (JSON format)
367    pub fn to_json_string(&self) -> Result<String> {
368        Ok(serde_json::to_string_pretty(&self.model)?)
369    }
370
371    /// Create a simple linear layer node
372    pub fn create_linear_node(
373        name: &str,
374        input_name: &str,
375        output_name: &str,
376        in_features: i64,
377        out_features: i64,
378        has_bias: bool,
379    ) -> GraphNode {
380        let mut attributes = HashMap::new();
381        attributes.insert("in_features".to_string(), AttributeValue::Int(in_features));
382        attributes.insert(
383            "out_features".to_string(),
384            AttributeValue::Int(out_features),
385        );
386        attributes.insert("bias".to_string(), AttributeValue::Bool(has_bias));
387
388        GraphNode {
389            name: name.to_string(),
390            op_type: "Linear".to_string(),
391            inputs: vec![input_name.to_string()],
392            outputs: vec![output_name.to_string()],
393            attributes,
394            doc_string: None,
395        }
396    }
397
398    /// Create a transformer attention node
399    pub fn create_attention_node(
400        name: &str,
401        input_name: &str,
402        output_name: &str,
403        num_heads: i64,
404        head_dim: i64,
405    ) -> GraphNode {
406        let mut attributes = HashMap::new();
407        attributes.insert("num_heads".to_string(), AttributeValue::Int(num_heads));
408        attributes.insert("head_dim".to_string(), AttributeValue::Int(head_dim));
409
410        GraphNode {
411            name: name.to_string(),
412            op_type: "MultiHeadAttention".to_string(),
413            inputs: vec![input_name.to_string()],
414            outputs: vec![output_name.to_string()],
415            attributes,
416            doc_string: Some("Multi-head self-attention layer".to_string()),
417        }
418    }
419
420    /// Create a layer normalization node
421    pub fn create_layernorm_node(
422        name: &str,
423        input_name: &str,
424        output_name: &str,
425        normalized_shape: Vec<i64>,
426        eps: f64,
427    ) -> GraphNode {
428        let mut attributes = HashMap::new();
429        attributes.insert(
430            "normalized_shape".to_string(),
431            AttributeValue::Ints(normalized_shape),
432        );
433        attributes.insert("eps".to_string(), AttributeValue::Float(eps));
434
435        GraphNode {
436            name: name.to_string(),
437            op_type: "LayerNorm".to_string(),
438            inputs: vec![input_name.to_string()],
439            outputs: vec![output_name.to_string()],
440            attributes,
441            doc_string: None,
442        }
443    }
444}
445
446#[cfg(test)]
447mod tests {
448    use super::*;
449
450    #[test]
451    fn onnx_export_refuses_instead_of_writing_json_under_an_onnx_name() {
452        let exporter = NetronExporter::new("m", "test model").with_format(ExportFormat::Onnx);
453        let path = std::env::temp_dir().join(format!("tfdbg_netron_{}.onnx", std::process::id()));
454        let _ = std::fs::remove_file(&path);
455        let err = exporter.export(&path).expect_err("ONNX must be refused");
456        let msg = err.to_string();
457        assert!(msg.contains("not implemented"), "{msg}");
458        assert!(msg.contains("protobuf"), "must name what is missing: {msg}");
459        assert!(
460            !path.exists(),
461            "no file may be written for a refused format"
462        );
463    }
464
465    #[test]
466    fn json_export_still_writes_a_real_graph() {
467        let exporter = NetronExporter::new("m", "test model").with_format(ExportFormat::Json);
468        let path = std::env::temp_dir().join(format!("tfdbg_netron_{}.json", std::process::id()));
469        exporter.export(&path).expect("json export works");
470        let text = std::fs::read_to_string(&path).expect("written");
471        let parsed: serde_json::Value = serde_json::from_str(&text).expect("valid json");
472        assert!(parsed.is_object());
473        let _ = std::fs::remove_file(&path);
474    }
475    use std::env;
476
477    #[test]
478    fn test_netron_exporter_creation() {
479        let exporter = NetronExporter::new("test_model", "A test model");
480        assert_eq!(exporter.model.metadata.name, "test_model");
481        assert_eq!(exporter.model.metadata.description, "A test model");
482    }
483
484    #[test]
485    fn test_add_input_output() {
486        let mut exporter = NetronExporter::new("test", "test");
487
488        exporter.add_input("input_ids", "int64", vec![1, 128]);
489        exporter.add_output("logits", "float32", vec![1, 128, 30522]);
490
491        assert_eq!(exporter.model.graph.inputs.len(), 1);
492        assert_eq!(exporter.model.graph.outputs.len(), 1);
493        assert_eq!(exporter.model.graph.inputs[0].name, "input_ids");
494    }
495
496    #[test]
497    fn test_add_node() {
498        let mut exporter = NetronExporter::new("test", "test");
499
500        let mut attrs = HashMap::new();
501        attrs.insert("in_features".to_string(), AttributeValue::Int(768));
502        attrs.insert("out_features".to_string(), AttributeValue::Int(3072));
503
504        exporter.add_node(
505            "fc1",
506            "Linear",
507            vec!["input".to_string()],
508            vec!["output".to_string()],
509            attrs,
510        );
511
512        assert_eq!(exporter.model.graph.nodes.len(), 1);
513        assert_eq!(exporter.model.graph.nodes[0].name, "fc1");
514        assert_eq!(exporter.model.graph.nodes[0].op_type, "Linear");
515    }
516
517    #[test]
518    fn test_export_json() {
519        let temp_dir = env::temp_dir();
520        let output_path = temp_dir.join("test_model.json");
521
522        let mut exporter = NetronExporter::new("test_model", "Test model");
523        exporter.add_input("input", "float32", vec![1, 10]);
524        exporter.add_output("output", "float32", vec![1, 5]);
525
526        exporter.export(&output_path).expect("operation failed in test");
527        assert!(output_path.exists());
528
529        // Clean up
530        let _ = fs::remove_file(output_path);
531    }
532
533    #[test]
534    fn test_create_linear_node() {
535        let node = NetronExporter::create_linear_node("fc1", "input", "output", 768, 3072, true);
536
537        assert_eq!(node.name, "fc1");
538        assert_eq!(node.op_type, "Linear");
539        assert!(node.attributes.contains_key("in_features"));
540        assert!(node.attributes.contains_key("bias"));
541    }
542
543    #[test]
544    fn test_create_attention_node() {
545        let node = NetronExporter::create_attention_node("attn", "input", "output", 12, 64);
546
547        assert_eq!(node.op_type, "MultiHeadAttention");
548        assert!(node.doc_string.is_some());
549    }
550
551    #[test]
552    fn test_metadata_setters() {
553        let mut exporter = NetronExporter::new("test", "test");
554
555        exporter.set_author("Test Author");
556        exporter.set_version("1.0.0");
557        exporter.add_property("framework", "TrustformeRS");
558
559        assert_eq!(
560            exporter.model.metadata.author,
561            Some("Test Author".to_string())
562        );
563        assert_eq!(exporter.model.metadata.version, Some("1.0.0".to_string()));
564        assert_eq!(
565            exporter.model.metadata.properties.get("framework"),
566            Some(&"TrustformeRS".to_string())
567        );
568    }
569
570    #[test]
571    fn test_to_json_string() {
572        let mut exporter = NetronExporter::new("test", "test");
573        exporter.add_input("input", "float32", vec![1, 10]);
574
575        let json = exporter.to_json_string().expect("operation failed in test");
576        assert!(json.contains("test"));
577        assert!(json.contains("input"));
578    }
579
580    #[test]
581    fn test_add_tensor_data() {
582        let mut exporter = NetronExporter::new("test", "test");
583
584        let weights = vec![0.1, 0.2, 0.3, 0.4];
585        exporter.add_tensor_data("layer.weight", "float32", vec![2, 2], Some(weights));
586
587        assert_eq!(exporter.model.graph.initializers.len(), 1);
588        assert_eq!(exporter.model.graph.initializers[0].name, "layer.weight");
589    }
590}
trustformers_debug/netron_export.rs

trustformers_debug/
netron_export.rs
