langvision 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

langvision/cli/export.py

@@ -0,0 +1,177 @@
+#!/usr/bin/env python3
+"""
+Export CLI for Langvision models.
+"""
+
+import argparse
+import logging
+import os
+import sys
+from pathlib import Path
+
+import torch
+import torch.onnx
+
+# Import from langvision modules
+try:
+    from langvision.models.vision_transformer import VisionTransformer
+except ImportError as e:
+    print(f"❌ Error importing langvision modules: {e}")
+    print("Please ensure langvision is properly installed:")
+    print("  pip install langvision")
+    sys.exit(1)
+
+
+def parse_args():
+    """Parse command-line arguments for model export."""
+    parser = argparse.ArgumentParser(
+        description='Export a trained VisionTransformer model to various formats',
+        epilog='''\nExamples:\n  langvision export --checkpoint model.pth --format onnx --output model.onnx\n  langvision export --checkpoint model.pth --format torchscript --output model.pt\n''',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    
+    # Model
+    model_group = parser.add_argument_group('Model')
+    model_group.add_argument('--checkpoint', type=str, required=True, help='Path to model checkpoint')
+    model_group.add_argument('--img_size', type=int, default=224, help='Input image size')
+    model_group.add_argument('--patch_size', type=int, default=16, help='Patch size for ViT')
+    model_group.add_argument('--num_classes', type=int, default=10, help='Number of classes')
+    model_group.add_argument('--embed_dim', type=int, default=768, help='Embedding dimension')
+    model_group.add_argument('--depth', type=int, default=12, help='Number of transformer layers')
+    model_group.add_argument('--num_heads', type=int, default=12, help='Number of attention heads')
+    model_group.add_argument('--mlp_ratio', type=float, default=4.0, help='MLP hidden dim ratio')
+    
+    # Export
+    export_group = parser.add_argument_group('Export')
+    export_group.add_argument('--format', type=str, required=True, choices=['onnx', 'torchscript', 'state_dict'], help='Export format')
+    export_group.add_argument('--output', type=str, required=True, help='Output file path')
+    export_group.add_argument('--batch_size', type=int, default=1, help='Batch size for export (ONNX/TorchScript)')
+    export_group.add_argument('--opset_version', type=int, default=11, help='ONNX opset version')
+    
+    # Device
+    device_group = parser.add_argument_group('Device')
+    device_group.add_argument('--device', type=str, default='cpu', help='Device to use for export (use CPU for ONNX)')
+    
+    # Misc
+    misc_group = parser.add_argument_group('Misc')
+    misc_group.add_argument('--log_level', type=str, default='info', help='Logging level')
+    
+    return parser.parse_args()
+
+
+def setup_logging(log_level: str) -> None:
+    """Set up logging with the specified log level."""
+    numeric_level = getattr(logging, log_level.upper(), None)
+    if not isinstance(numeric_level, int):
+        numeric_level = logging.INFO
+    logging.basicConfig(level=numeric_level, format='[%(levelname)s] %(message)s')
+
+
+def main():
+    """Main function for model export."""
+    args = parse_args()
+    setup_logging(args.log_level)
+    logger = logging.getLogger(__name__)
+    
+    logger.info("🚀 Starting model export...")
+    
+    # Load model
+    logger.info("🤖 Loading model...")
+    try:
+        model = VisionTransformer(
+            img_size=args.img_size,
+            patch_size=args.patch_size,
+            in_chans=3,
+            num_classes=args.num_classes,
+            embed_dim=args.embed_dim,
+            depth=args.depth,
+            num_heads=args.num_heads,
+            mlp_ratio=args.mlp_ratio,
+        ).to(args.device)
+        
+        # Load checkpoint
+        if not os.path.isfile(args.checkpoint):
+            logger.error(f"❌ Checkpoint file not found: {args.checkpoint}")
+            return 1
+            
+        checkpoint = torch.load(args.checkpoint, map_location=args.device)
+        if 'model' in checkpoint:
+            model.load_state_dict(checkpoint['model'])
+        else:
+            model.load_state_dict(checkpoint)
+        
+        model.eval()
+        logger.info("✅ Model loaded successfully")
+        
+    except Exception as e:
+        logger.error(f"❌ Failed to load model: {e}")
+        return 1
+    
+    # Create output directory
+    os.makedirs(os.path.dirname(args.output), exist_ok=True)
+    
+    # Export model
+    logger.info(f"📦 Exporting model to {args.format.upper()} format...")
+    try:
+        if args.format == 'onnx':
+            # Create dummy input
+            dummy_input = torch.randn(args.batch_size, 3, args.img_size, args.img_size).to(args.device)
+            
+            # Export to ONNX
+            torch.onnx.export(
+                model,
+                dummy_input,
+                args.output,
+                export_params=True,
+                opset_version=args.opset_version,
+                do_constant_folding=True,
+                input_names=['input'],
+                output_names=['output'],
+                dynamic_axes={
+                    'input': {0: 'batch_size'},
+                    'output': {0: 'batch_size'}
+                }
+            )
+            logger.info(f"✅ ONNX model exported to {args.output}")
+            
+        elif args.format == 'torchscript':
+            # Create dummy input
+            dummy_input = torch.randn(args.batch_size, 3, args.img_size, args.img_size).to(args.device)
+            
+            # Export to TorchScript
+            traced_model = torch.jit.trace(model, dummy_input)
+            traced_model.save(args.output)
+            logger.info(f"✅ TorchScript model exported to {args.output}")
+            
+        elif args.format == 'state_dict':
+            # Export state dict
+            torch.save(model.state_dict(), args.output)
+            logger.info(f"✅ State dict exported to {args.output}")
+            
+    except Exception as e:
+        logger.error(f"❌ Export failed: {e}")
+        return 1
+    
+    # Verify export
+    if args.format in ['onnx', 'torchscript']:
+        logger.info("🔍 Verifying exported model...")
+        try:
+            if args.format == 'onnx':
+                import onnx
+                onnx_model = onnx.load(args.output)
+                onnx.checker.check_model(onnx_model)
+                logger.info("✅ ONNX model verification passed")
+                
+            elif args.format == 'torchscript':
+                loaded_model = torch.jit.load(args.output)
+                logger.info("✅ TorchScript model verification passed")
+                
+        except Exception as e:
+            logger.warning(f"⚠️ Model verification failed: {e}")
+    
+    logger.info("✅ Model export completed successfully!")
+    return 0
+
+
+if __name__ == '__main__':
+    sys.exit(main())

langvision/cli/finetune.py

@@ -23,47 +23,67 @@ def set_seed(seed: int) -> None:
 
 def parse_args() -> argparse.Namespace:
     """Parse command-line arguments for fine-tuning."""
-    parser = argparse.ArgumentParser(description='Fine-tune VisionTransformer with LoRA')
+    parser = argparse.ArgumentParser(
+        description='Fine-tune VisionTransformer with LoRA and advanced LLM concepts',
+        epilog='''\nExamples:\n  langvision finetune --dataset cifar10 --epochs 10\n  langvision finetune --dataset cifar100 --lora_r 8 --rlhf\n''',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
     # Data
-    parser.add_argument('--dataset', type=str, default='cifar10', choices=['cifar10', 'cifar100'], help='Dataset to use')
-    parser.add_argument('--data_dir', type=str, default='./data', help='Dataset directory')
-    parser.add_argument('--num_classes', type=int, default=10, help='Number of classes')
-    parser.add_argument('--img_size', type=int, default=224, help='Input image size')
-    parser.add_argument('--patch_size', type=int, default=16, help='Patch size for ViT')
-    parser.add_argument('--num_workers', type=int, default=2, help='Number of data loader workers')
+    data_group = parser.add_argument_group('Data')
+    data_group.add_argument('--dataset', type=str, default='cifar10', choices=['cifar10', 'cifar100'], help='Dataset to use')
+    data_group.add_argument('--data_dir', type=str, default='./data', help='Dataset directory')
+    data_group.add_argument('--num_classes', type=int, default=10, help='Number of classes in the dataset')
+    data_group.add_argument('--img_size', type=int, default=224, help='Input image size (pixels)')
+    data_group.add_argument('--patch_size', type=int, default=16, help='Patch size for Vision Transformer')
+    data_group.add_argument('--num_workers', type=int, default=2, help='Number of data loader workers')
     # Model
-    parser.add_argument('--embed_dim', type=int, default=768, help='Embedding dimension')
-    parser.add_argument('--depth', type=int, default=12, help='Number of transformer layers')
-    parser.add_argument('--num_heads', type=int, default=12, help='Number of attention heads')
-    parser.add_argument('--mlp_ratio', type=float, default=4.0, help='MLP ratio')
-    parser.add_argument('--lora_r', type=int, default=4, help='LoRA rank')
-    parser.add_argument('--lora_alpha', type=float, default=1.0, help='LoRA alpha')
-    parser.add_argument('--lora_dropout', type=float, default=0.1, help='LoRA dropout')
+    model_group = parser.add_argument_group('Model')
+    model_group.add_argument('--embed_dim', type=int, default=768, help='Embedding dimension for ViT')
+    model_group.add_argument('--depth', type=int, default=12, help='Number of transformer layers')
+    model_group.add_argument('--num_heads', type=int, default=12, help='Number of attention heads')
+    model_group.add_argument('--mlp_ratio', type=float, default=4.0, help='MLP hidden dim ratio')
+    model_group.add_argument('--lora_r', type=int, default=4, help='LoRA rank (low-rank adaptation)')
+    model_group.add_argument('--lora_alpha', type=float, default=1.0, help='LoRA alpha scaling')
+    model_group.add_argument('--lora_dropout', type=float, default=0.1, help='LoRA dropout rate')
     # Training
-    parser.add_argument('--batch_size', type=int, default=64, help='Batch size for training')
-    parser.add_argument('--epochs', type=int, default=10, help='Number of epochs')
-    parser.add_argument('--lr', type=float, default=1e-3, help='Learning rate')
-    parser.add_argument('--optimizer', type=str, default='adam', choices=['adam', 'adamw', 'sgd'], help='Optimizer')
-    parser.add_argument('--weight_decay', type=float, default=0.01, help='Weight decay (L2 regularization)')
-    parser.add_argument('--scheduler', type=str, default='cosine', choices=['cosine', 'step'], help='LR scheduler')
-    parser.add_argument('--step_size', type=int, default=5, help='StepLR step size')
-    parser.add_argument('--gamma', type=float, default=0.5, help='StepLR gamma')
-    parser.add_argument('--resume', type=str, default=None, help='Path to checkpoint to resume from')
-    parser.add_argument('--eval_only', action='store_true', help='Only run evaluation')
+    train_group = parser.add_argument_group('Training')
+    train_group.add_argument('--batch_size', type=int, default=64, help='Batch size for training')
+    train_group.add_argument('--epochs', type=int, default=10, help='Number of training epochs')
+    train_group.add_argument('--lr', type=float, default=1e-3, help='Learning rate')
+    train_group.add_argument('--optimizer', type=str, default='adam', choices=['adam', 'adamw', 'sgd'], help='Optimizer type')
+    train_group.add_argument('--weight_decay', type=float, default=0.01, help='Weight decay (L2 regularization)')
+    train_group.add_argument('--scheduler', type=str, default='cosine', choices=['cosine', 'step'], help='Learning rate scheduler')
+    train_group.add_argument('--step_size', type=int, default=5, help='StepLR: step size')
+    train_group.add_argument('--gamma', type=float, default=0.5, help='StepLR: gamma')
+    train_group.add_argument('--resume', type=str, default=None, help='Path to checkpoint to resume from')
+    train_group.add_argument('--eval_only', action='store_true', help='Only run evaluation (no training)')
     # Output
-    parser.add_argument('--output_dir', type=str, default='outputs', help='Directory to save outputs and checkpoints')
-    parser.add_argument('--save_name', type=str, default='vit_lora_best.pth', help='Checkpoint file name')
+    output_group = parser.add_argument_group('Output')
+    output_group.add_argument('--output_dir', type=str, default='outputs', help='Directory to save outputs and checkpoints')
+    output_group.add_argument('--save_name', type=str, default='vit_lora_best.pth', help='Checkpoint file name')
     # Callbacks
-    parser.add_argument('--early_stopping', action='store_true', help='Enable early stopping')
-    parser.add_argument('--patience', type=int, default=5, help='Early stopping patience')
+    callback_group = parser.add_argument_group('Callbacks')
+    callback_group.add_argument('--early_stopping', action='store_true', help='Enable early stopping')
+    callback_group.add_argument('--patience', type=int, default=5, help='Early stopping patience')
     # CUDA
-    parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu', help='Device to use')
-    parser.add_argument('--cuda_deterministic', action='store_true', help='Enable deterministic CUDA (reproducible, slower)')
-    parser.add_argument('--cuda_benchmark', action='store_true', default=True, help='Enable cudnn.benchmark for fast training (default: True)')
-    parser.add_argument('--cuda_max_split_size_mb', type=int, default=None, help='Set CUDA max split size in MB (for large models, PyTorch >=1.10)')
+    cuda_group = parser.add_argument_group('CUDA')
+    cuda_group.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu', help='Device to use (cuda or cpu)')
+    cuda_group.add_argument('--cuda_deterministic', action='store_true', help='Enable deterministic CUDA (reproducible, slower)')
+    cuda_group.add_argument('--cuda_benchmark', action='store_true', default=True, help='Enable cudnn.benchmark for fast training')
+    cuda_group.add_argument('--cuda_max_split_size_mb', type=int, default=None, help='Set CUDA max split size in MB (for large models, PyTorch >=1.10)')
     # Misc
-    parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility')
-    parser.add_argument('--log_level', type=str, default='info', help='Logging level (debug, info, warning, error)')
+    misc_group = parser.add_argument_group('Misc')
+    misc_group.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility')
+    misc_group.add_argument('--log_level', type=str, default='info', help='Logging level (debug, info, warning, error)')
+    # Advanced LLM concepts
+    llm_group = parser.add_argument_group('Advanced LLM Concepts')
+    llm_group.add_argument('--rlhf', action='store_true', help='Use RLHF (Reinforcement Learning from Human Feedback)')
+    llm_group.add_argument('--ppo', action='store_true', help='Use PPO (Proximal Policy Optimization)')
+    llm_group.add_argument('--dpo', action='store_true', help='Use DPO (Direct Preference Optimization)')
+    llm_group.add_argument('--lime', action='store_true', help='Use LIME for model explainability')
+    llm_group.add_argument('--shap', action='store_true', help='Use SHAP for model explainability')
+    llm_group.add_argument('--cot', action='store_true', help='Use Chain-of-Thought (CoT) prompt generation')
+    llm_group.add_argument('--ccot', action='store_true', help='Use Contrastive Chain-of-Thought (CCoT)')
     return parser.parse_args()
 
 
@@ -76,12 +96,11 @@ def setup_logging(log_level: str) -> None:
 
 
 def main() -> None:
-    """Main function for fine-tuning VisionTransformer with LoRA."""
+    """Main function for fine-tuning VisionTransformer with LoRA and advanced LLM concepts."""
     args = parse_args()
     setup_logging(args.log_level)
     logger = logging.getLogger(__name__)
-    logger.info(f"Using device: {args.device}")
-
+    logger.info("[STEP 1] Loading dataset...")
     setup_cuda(seed=args.seed, deterministic=args.cuda_deterministic, benchmark=args.cuda_benchmark, max_split_size_mb=args.cuda_max_split_size_mb)
     set_seed(args.seed)
 
@@ -95,6 +114,7 @@ def main() -> None:
     train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers)
     val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers)
 
+    logger.info("[STEP 2] Initializing model...")
     # Model
     try:
         model = VisionTransformer(
@@ -116,6 +136,7 @@ def main() -> None:
         logger.error(f"Failed to initialize model: {e}")
         return
 
+    logger.info("[STEP 3] Setting up optimizer and scheduler...")
     # Optimizer
     lora_params = [p for n, p in model.named_parameters() if 'lora' in n and p.requires_grad]
     if args.optimizer == 'adam':
@@ -125,7 +146,6 @@ def main() -> None:
     else:
         optimizer = torch.optim.SGD(lora_params, lr=args.lr, momentum=0.9, weight_decay=args.weight_decay)
 
-    # Scheduler
     if args.scheduler == 'cosine':
         scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
     else:
@@ -134,11 +154,64 @@ def main() -> None:
     criterion = torch.nn.CrossEntropyLoss()
     scaler = torch.cuda.amp.GradScaler() if args.device == 'cuda' else None
 
+    logger.info("[STEP 4] Setting up trainer...")
     # Callbacks
     callbacks = []
     if args.early_stopping:
         callbacks.append(EarlyStopping(patience=args.patience))
 
+    # Advanced LLM concept objects
+    rlhf = None
+    ppo = None
+    dpo = None
+    lime = None
+    shap = None
+    cot = None
+    ccot = None
+    if args.rlhf:
+        from langvision.concepts.rlhf import RLHF
+        class SimpleRLHF(RLHF):
+            def train(self, model, data, feedback_fn, optimizer):
+                super().train(model, data, feedback_fn, optimizer)
+        rlhf = SimpleRLHF()
+    if args.ppo:
+        from langvision.concepts.ppo import PPO
+        class SimplePPO(PPO):
+            def step(self, policy, old_log_probs, states, actions, rewards, optimizer):
+                super().step(policy, old_log_probs, states, actions, rewards, optimizer)
+        ppo = SimplePPO()
+    if args.dpo:
+        from langvision.concepts.dpo import DPO
+        class SimpleDPO(DPO):
+            def optimize_with_preferences(self, model, preferences, optimizer):
+                super().optimize_with_preferences(model, preferences, optimizer)
+        dpo = SimpleDPO()
+    if args.lime:
+        from langvision.concepts.lime import LIME
+        class SimpleLIME(LIME):
+            def explain(self, model, input_data):
+                return super().explain(model, input_data)
+        lime = SimpleLIME()
+    if args.shap:
+        from langvision.concepts.shap import SHAP
+        class SimpleSHAP(SHAP):
+            def explain(self, model, input_data):
+                return super().explain(model, input_data)
+        shap = SimpleSHAP()
+    if args.cot:
+        from langvision.concepts.cot import CoT
+        class SimpleCoT(CoT):
+            def generate_chain(self, prompt):
+                return super().generate_chain(prompt)
+        cot = SimpleCoT()
+    if args.ccot:
+        from langvision.concepts.ccot import CCoT
+        class SimpleCCoT(CCoT):
+            def contrastive_train(self, positive_chains, negative_chains):
+                super().contrastive_train(positive_chains, negative_chains)
+        ccot = SimpleCCoT()
+
+    logger.info("[STEP 5] (Optional) Loading checkpoint if provided...")
     # Trainer
     trainer = Trainer(
         model=model,
@@ -148,6 +221,9 @@ def main() -> None:
         scaler=scaler,
         callbacks=callbacks,
         device=args.device,
+        rlhf=rlhf,
+        ppo=ppo,
+        dpo=dpo,
     )
 
     # Optionally resume
@@ -162,20 +238,61 @@ def main() -> None:
         logger.info(f"Resumed from {args.resume} at epoch {start_epoch}")
 
     if args.eval_only:
+        logger.info("[STEP 6] Running evaluation only...")
         val_loss, val_acc = trainer.evaluate(val_loader)
         logger.info(f"Eval Loss: {val_loss:.4f}, Eval Acc: {val_acc:.4f}")
         return
 
-    # Training
-    best_acc = trainer.fit(
-        train_loader,
-        val_loader,
-        epochs=args.epochs,
-        start_epoch=start_epoch,
-        best_acc=best_acc,
-        checkpoint_path=os.path.join(args.output_dir, args.save_name),
-    )
-    logger.info(f"Best validation accuracy: {best_acc:.4f}")
+    logger.info("[STEP 6] Starting training...")
+    # Training with per-epoch progress
+    for epoch in range(start_epoch, args.epochs):
+        logger.info(f"  [Epoch {epoch+1}/{args.epochs}] Starting epoch...")
+        best_acc = trainer.fit(
+            train_loader,
+            val_loader,
+            epochs=epoch+1,
+            start_epoch=epoch,
+            best_acc=best_acc,
+            checkpoint_path=os.path.join(args.output_dir, args.save_name),
+        )
+        logger.info(f"  [Epoch {epoch+1}/{args.epochs}] Done. Best validation accuracy so far: {best_acc:.4f}")
+    logger.info(f"[STEP 7] Training complete. Best validation accuracy: {best_acc:.4f}")
+
+    # Example: Use RLHF, PPO, DPO, LIME, SHAP, CoT, CCoT in training
+    if rlhf is not None:
+        logger.info("[STEP 8] Applying RLHF...")
+        def feedback_fn(output):
+            return 1.0 if output.sum().item() > 0 else -1.0
+        rlhf.train(model, [torch.randn(3) for _ in range(10)], feedback_fn, optimizer)
+    if ppo is not None:
+        logger.info("[STEP 9] Applying PPO...")
+        import torch
+        policy = model
+        old_log_probs = torch.zeros(10)
+        states = torch.randn(10, 3)
+        actions = torch.randint(0, args.num_classes, (10,))
+        rewards = torch.randn(10)
+        ppo.step(policy, old_log_probs, states, actions, rewards, optimizer)
+    if dpo is not None:
+        logger.info("[STEP 10] Applying DPO...")
+        preferences = [(torch.randn(3), 1.0), (torch.randn(3), -1.0)]
+        dpo.optimize_with_preferences(model, preferences, optimizer)
+    if lime is not None:
+        logger.info("[STEP 11] Running LIME explainability...")
+        lime_explanation = lime.explain(model, [[0.5, 1.0, 2.0], [1.0, 2.0, 3.0]])
+        logger.info(f"LIME explanation: {lime_explanation}")
+    if shap is not None:
+        logger.info("[STEP 12] Running SHAP explainability...")
+        shap_explanation = shap.explain(model, [[0.5, 1.0, 2.0], [1.0, 2.0, 3.0]])
+        logger.info(f"SHAP explanation: {shap_explanation}")
+    if cot is not None:
+        logger.info("[STEP 13] Generating Chain-of-Thought...")
+        chain = cot.generate_chain("What is 2 + 2?")
+        logger.info(f"CoT chain: {chain}")
+    if ccot is not None:
+        logger.info("[STEP 14] Running Contrastive Chain-of-Thought...")
+        ccot.contrastive_train([['Step 1: Think', 'Step 2: Solve']], [['Step 1: Guess', 'Step 2: Wrong']])
+    logger.info("[COMPLETE] All steps finished.")
 
 if __name__ == '__main__':
     main() 

langvision/cli/model_zoo.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python3
+"""
+Model Zoo CLI for Langvision.
+"""
+
+import argparse
+import json
+import logging
+import sys
+from pathlib import Path
+
+# Import from langvision modules
+try:
+    from langvision.model_zoo import get_available_models, get_model_info, download_model
+except ImportError as e:
+    print(f"❌ Error importing langvision modules: {e}")
+    print("Please ensure langvision is properly installed:")
+    print("  pip install langvision")
+    sys.exit(1)
+
+
+def parse_args():
+    """Parse command-line arguments for model zoo operations."""
+    parser = argparse.ArgumentParser(
+        description='Langvision Model Zoo - Browse, download, and manage pre-trained models',
+        epilog='''\nExamples:\n  langvision model-zoo list\n  langvision model-zoo info vit_base_patch16_224\n  langvision model-zoo download vit_base_patch16_224\n''',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    
+    subparsers = parser.add_subparsers(dest='command', required=True, help='Model zoo commands')
+    
+    # List command
+    list_parser = subparsers.add_parser('list', help='List all available models')
+    list_parser.add_argument('--format', type=str, default='table', choices=['table', 'json'], help='Output format')
+    list_parser.add_argument('--filter', type=str, help='Filter models by name or type')
+    
+    # Info command
+    info_parser = subparsers.add_parser('info', help='Get detailed information about a model')
+    info_parser.add_argument('model_name', type=str, help='Name of the model')
+    info_parser.add_argument('--format', type=str, default='table', choices=['table', 'json'], help='Output format')
+    
+    # Download command
+    download_parser = subparsers.add_parser('download', help='Download a pre-trained model')
+    download_parser.add_argument('model_name', type=str, help='Name of the model to download')
+    download_parser.add_argument('--output_dir', type=str, default='./models', help='Directory to save the model')
+    download_parser.add_argument('--force', action='store_true', help='Force download even if model exists')
+    
+    # Search command
+    search_parser = subparsers.add_parser('search', help='Search for models')
+    search_parser.add_argument('query', type=str, help='Search query')
+    search_parser.add_argument('--format', type=str, default='table', choices=['table', 'json'], help='Output format')
+    
+    # Misc
+    parser.add_argument('--log_level', type=str, default='info', help='Logging level')
+    
+    return parser.parse_args()
+
+
+def setup_logging(log_level: str) -> None:
+    """Set up logging with the specified log level."""
+    numeric_level = getattr(logging, log_level.upper(), None)
+    if not isinstance(numeric_level, int):
+        numeric_level = logging.INFO
+    logging.basicConfig(level=numeric_level, format='[%(levelname)s] %(message)s')
+
+
+def print_model_table(models):
+    """Print models in a formatted table."""
+    if not models:
+        print("No models found.")
+        return
+    
+    # Calculate column widths
+    name_width = max(len(model.get('name', '')) for model in models) + 2
+    type_width = max(len(model.get('type', '')) for model in models) + 2
+    size_width = max(len(str(model.get('size', ''))) for model in models) + 2
+    
+    # Print header
+    print(f"{'Name':<{name_width}} {'Type':<{type_width}} {'Size':<{size_width}} {'Description'}")
+    print("-" * (name_width + type_width + size_width + 50))
+    
+    # Print models
+    for model in models:
+        name = model.get('name', '')
+        model_type = model.get('type', '')
+        size = model.get('size', '')
+        description = model.get('description', '')
+        print(f"{name:<{name_width}} {model_type:<{type_width}} {size:<{size_width}} {description}")
+
+
+def print_model_info(model_info, format='table'):
+    """Print detailed model information."""
+    if format == 'json':
+        print(json.dumps(model_info, indent=2))
+        return
+    
+    print(f"\n🤖 Model: {model_info.get('name', 'Unknown')}")
+    print("=" * 50)
+    
+    for key, value in model_info.items():
+        if key == 'name':
+            continue
+        print(f"{key.replace('_', ' ').title()}: {value}")
+    
+    if 'config' in model_info:
+        print(f"\n📋 Configuration:")
+        config = model_info['config']
+        for key, value in config.items():
+            print(f"  {key}: {value}")
+
+
+def main():
+    """Main function for model zoo operations."""
+    args = parse_args()
+    setup_logging(args.log_level)
+    logger = logging.getLogger(__name__)
+    
+    try:
+        if args.command == 'list':
+            logger.info("📋 Fetching available models...")
+            models = get_available_models()
+            
+            if args.filter:
+                models = [m for m in models if args.filter.lower() in m.get('name', '').lower()]
+            
+            if args.format == 'json':
+                print(json.dumps(models, indent=2))
+            else:
+                print(f"\n🎯 Available Models ({len(models)} total):")
+                print_model_table(models)
+                
+        elif args.command == 'info':
+            logger.info(f"🔍 Getting information for model: {args.model_name}")
+            model_info = get_model_info(args.model_name)
+            print_model_info(model_info, args.format)
+            
+        elif args.command == 'download':
+            logger.info(f"⬇️ Downloading model: {args.model_name}")
+            output_path = download_model(args.model_name, args.output_dir, force=args.force)
+            logger.info(f"✅ Model downloaded to: {output_path}")
+            
+        elif args.command == 'search':
+            logger.info(f"🔍 Searching for: {args.query}")
+            models = get_available_models()
+            results = [m for m in models if args.query.lower() in m.get('name', '').lower() or 
+                      args.query.lower() in m.get('description', '').lower()]
+            
+            if args.format == 'json':
+                print(json.dumps(results, indent=2))
+            else:
+                print(f"\n🔍 Search Results for '{args.query}' ({len(results)} found):")
+                print_model_table(results)
+                
+    except Exception as e:
+        logger.error(f"❌ Operation failed: {e}")
+        return 1
+    
+    return 0
+
+
+if __name__ == '__main__':
+    sys.exit(main())