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 pytorch_reader;
35pub mod real_benchmarking;
36pub mod serialization;
37pub mod tensor_integration;
38pub mod types;
39pub mod validation;
40
41// Re-export commonly used types for convenience
42pub use args::*;
43
44/// Model operation subcommands
45#[derive(Subcommand)]
46pub enum ModelCommands {
47    /// Convert model between different formats
48    Convert(ConvertArgs),
49
50    /// Optimize model for deployment
51    Optimize(OptimizeArgs),
52
53    /// Quantize model to reduce size and improve performance
54    Quantize(QuantizeArgs),
55
56    /// Prune model to remove unnecessary parameters
57    Prune(PruneArgs),
58
59    /// Inspect model architecture and properties
60    Inspect(InspectArgs),
61
62    /// Validate model functionality and accuracy
63    Validate(ValidateArgs),
64
65    /// Benchmark model performance
66    Benchmark(BenchmarkArgs),
67
68    /// Compress model using various techniques
69    Compress(CompressArgs),
70
71    /// Extract model components (weights, architecture, etc.)
72    Extract(ExtractArgs),
73
74    /// Merge multiple models
75    Merge(MergeArgs),
76}
77
78/// Execute model operation based on subcommand
79pub async fn execute(cmd: ModelCommands, config: &Config, output_format: &str) -> Result<()> {
80    match cmd {
81        ModelCommands::Convert(args) => {
82            conversion::convert_model(args, config, output_format).await
83        }
84        ModelCommands::Optimize(args) => {
85            optimization::optimize_model(args, config, output_format).await
86        }
87        ModelCommands::Quantize(args) => {
88            optimization::quantize_model(args, config, output_format).await
89        }
90        ModelCommands::Prune(args) => optimization::prune_model(args, config, output_format).await,
91        ModelCommands::Inspect(args) => analysis::inspect_model(args, config, output_format).await,
92        ModelCommands::Validate(args) => {
93            analysis::validate_model(args, config, output_format).await
94        }
95        ModelCommands::Benchmark(args) => {
96            benchmarking::benchmark_model(args, config, output_format).await
97        }
98        ModelCommands::Compress(args) => {
99            conversion::compress_model(args, config, output_format).await
100        }
101        ModelCommands::Extract(args) => {
102            conversion::extract_model(args, config, output_format).await
103        }
104        ModelCommands::Merge(args) => conversion::merge_model(args, config, output_format).await,
105    }
106}