Skip to main content

torsh_cli/commands/model/
mod.rs

1//! Model operation commands
2//!
3//! This module provides a comprehensive set of model operations including:
4//! - Conversion between different model formats
5//! - Optimization for deployment
6//! - Quantization for efficiency
7//! - Pruning for size reduction
8//! - Analysis and inspection
9//! - Benchmarking performance
10//!
11//! The implementation is modularized for better maintainability:
12//! - `analysis`: Model inspection and validation
13//! - `benchmarking`: Performance testing
14//! - `conversion`: Format conversion and utility operations
15//! - `optimization`: Model optimization, quantization, and pruning
16//! - `types`: Shared data structures
17//! - `args`: Command-line argument definitions
18
19use anyhow::Result;
20use clap::Subcommand;
21
22use crate::config::Config;
23
24// Import all sub-modules
25pub mod analysis;
26pub mod args;
27pub mod benchmarking;
28pub mod conversion;
29pub mod enhanced_profiling;
30pub mod enhanced_serialization;
31pub mod optimization;
32pub mod profiling;
33pub mod pytorch_parser;
34pub mod real_benchmarking;
35pub mod serialization;
36pub mod tensor_integration;
37pub mod types;
38pub mod validation;
39
40// Re-export commonly used types for convenience
41pub use args::*;
42
43/// Model operation subcommands
44#[derive(Subcommand)]
45pub enum ModelCommands {
46    /// Convert model between different formats
47    Convert(ConvertArgs),
48
49    /// Optimize model for deployment
50    Optimize(OptimizeArgs),
51
52    /// Quantize model to reduce size and improve performance
53    Quantize(QuantizeArgs),
54
55    /// Prune model to remove unnecessary parameters
56    Prune(PruneArgs),
57
58    /// Inspect model architecture and properties
59    Inspect(InspectArgs),
60
61    /// Validate model functionality and accuracy
62    Validate(ValidateArgs),
63
64    /// Benchmark model performance
65    Benchmark(BenchmarkArgs),
66
67    /// Compress model using various techniques
68    Compress(CompressArgs),
69
70    /// Extract model components (weights, architecture, etc.)
71    Extract(ExtractArgs),
72
73    /// Merge multiple models
74    Merge(MergeArgs),
75}
76
77/// Execute model operation based on subcommand
78pub async fn execute(cmd: ModelCommands, config: &Config, output_format: &str) -> Result<()> {
79    match cmd {
80        ModelCommands::Convert(args) => {
81            conversion::convert_model(args, config, output_format).await
82        }
83        ModelCommands::Optimize(args) => {
84            optimization::optimize_model(args, config, output_format).await
85        }
86        ModelCommands::Quantize(args) => {
87            optimization::quantize_model(args, config, output_format).await
88        }
89        ModelCommands::Prune(args) => optimization::prune_model(args, config, output_format).await,
90        ModelCommands::Inspect(args) => analysis::inspect_model(args, config, output_format).await,
91        ModelCommands::Validate(args) => {
92            analysis::validate_model(args, config, output_format).await
93        }
94        ModelCommands::Benchmark(args) => {
95            benchmarking::benchmark_model(args, config, output_format).await
96        }
97        ModelCommands::Compress(args) => {
98            conversion::compress_model(args, config, output_format).await
99        }
100        ModelCommands::Extract(args) => {
101            conversion::extract_model(args, config, output_format).await
102        }
103        ModelCommands::Merge(args) => conversion::merge_model(args, config, output_format).await,
104    }
105}