entrenar_inspect/
lib.rs

1//! SafeTensors model inspection and format conversion.
2//!
3//! This crate provides tools for:
4//! - Inspecting model architectures and layer structures
5//! - Analyzing memory requirements
6//! - Converting between formats (SafeTensors, GGUF, APR)
7//!
8//! # Toyota Way Principles
9//!
10//! - **Andon**: Integrity checker surfaces data problems immediately
11//! - **Genchi Genbutsu**: Direct inspection of actual model weights
12//! - **SMED**: Quick format conversion for deployment changeover
13
14pub mod architecture;
15pub mod convert;
16pub mod inspect;
17pub mod validate;
18
19pub use architecture::{Architecture, ArchitectureDetector};
20pub use convert::{ConversionResult, FormatConverter};
21pub use inspect::{ModelInfo, TensorInfo};
22pub use validate::{IntegrityChecker, ValidationResult};
23
24use entrenar_common::Result;
25use std::path::Path;
26
27/// Inspect a model file and return detailed information.
28pub fn inspect(path: impl AsRef<Path>) -> Result<ModelInfo> {
29    inspect::inspect_model(path)
30}
31
32/// Validate a model file's integrity.
33pub fn validate(path: impl AsRef<Path>) -> Result<ValidationResult> {
34    validate::validate_model(path)
35}
36
37/// Convert a model between formats.
38pub fn convert(
39    input: impl AsRef<Path>,
40    output: impl AsRef<Path>,
41    format: OutputFormat,
42) -> Result<ConversionResult> {
43    convert::convert_model(input, output, format)
44}
45
46/// Output format for model conversion.
47#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub enum OutputFormat {
49    /// SafeTensors format (secure, recommended)
50    SafeTensors,
51    /// GGUF format (llama.cpp compatible)
52    Gguf,
53    /// APR format (JSON metadata)
54    Apr,
55}
56
57impl std::str::FromStr for OutputFormat {
58    type Err = String;
59
60    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
61        match s.to_lowercase().as_str() {
62            "safetensors" | "st" => Ok(Self::SafeTensors),
63            "gguf" => Ok(Self::Gguf),
64            "apr" => Ok(Self::Apr),
65            _ => Err(format!(
66                "Unknown format: {s}. Use: safetensors, gguf, apr"
67            )),
68        }
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75
76    #[test]
77    fn test_output_format_parsing() {
78        assert_eq!(
79            "safetensors".parse::<OutputFormat>().unwrap(),
80            OutputFormat::SafeTensors
81        );
82        assert_eq!("GGUF".parse::<OutputFormat>().unwrap(), OutputFormat::Gguf);
83        assert_eq!("apr".parse::<OutputFormat>().unwrap(), OutputFormat::Apr);
84    }
85
86    #[test]
87    fn test_output_format_st_alias() {
88        assert_eq!("st".parse::<OutputFormat>().unwrap(), OutputFormat::SafeTensors);
89        assert_eq!("ST".parse::<OutputFormat>().unwrap(), OutputFormat::SafeTensors);
90    }
91
92    #[test]
93    fn test_output_format_invalid() {
94        let result = "pickle".parse::<OutputFormat>();
95        assert!(result.is_err());
96        let err = result.unwrap_err();
97        assert!(err.contains("Unknown format"));
98        assert!(err.contains("pickle"));
99    }
100
101    #[test]
102    fn test_output_format_equality() {
103        assert_eq!(OutputFormat::SafeTensors, OutputFormat::SafeTensors);
104        assert_ne!(OutputFormat::SafeTensors, OutputFormat::Gguf);
105        assert_ne!(OutputFormat::Gguf, OutputFormat::Apr);
106    }
107}