scripts/rescue/recover_model.py

#!/usr/bin/env python3
"""
Model Recovery and Deployment Script
Recovers trained model from cloud instance, quantizes it, and pushes to Hugging Face Hub
"""

import os
import sys
import json
import argparse
import logging
import subprocess
from pathlib import Path
from typing import Dict, Any, Optional
from datetime import datetime

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
sys.path.append(os.path.join(os.path.dirname(__file__), 'src'))


class ModelRecoveryPipeline:
    """Complete model recovery and deployment pipeline"""
    
    def __init__(
        self,
        model_path: str,
        repo_name: str,
        hf_token: Optional[str] = None,
        private: bool = False,
        quantize: bool = True,
        quant_types: Optional[list] = None,
        trackio_url: Optional[str] = None,
        experiment_name: Optional[str] = None,
        dataset_repo: Optional[str] = None,
        author_name: Optional[str] = None,
        model_description: Optional[str] = None
    ):
        self.model_path = Path(model_path)
        self.repo_name = repo_name
        self.hf_token = hf_token or os.getenv('HF_TOKEN')
        self.private = private
        self.quantize = quantize
        self.quant_types = quant_types or ["int8_weight_only", "int4_weight_only"]
        self.trackio_url = trackio_url
        self.experiment_name = experiment_name
        self.dataset_repo = dataset_repo
        self.author_name = author_name
        self.model_description = model_description
        
        # Validate HF token
        if not self.hf_token:
            raise ValueError("HF_TOKEN environment variable or --hf-token argument is required")
        
        logger.info(f"Initialized ModelRecoveryPipeline for {repo_name}")
        logger.info(f"Model path: {self.model_path}")
        logger.info(f"Quantization enabled: {self.quantize}")
        if self.quantize:
            logger.info(f"Quantization types: {self.quant_types}")
    
    def validate_model_path(self) -> bool:
        """Validate that the model path contains required files"""
        if not self.model_path.exists():
            logger.error(f"❌ Model path does not exist: {self.model_path}")
            return False
        
        # Check for essential model files
        required_files = ['config.json']
        
        # Check for model files (either safetensors or pytorch)
        model_files = [
            "model.safetensors.index.json",  # Safetensors format
            "pytorch_model.bin"  # PyTorch format
        ]
        
        missing_files = []
        for file in required_files:
            if not (self.model_path / file).exists():
                missing_files.append(file)
        
        # Check if at least one model file exists
        model_file_exists = any((self.model_path / file).exists() for file in model_files)
        if not model_file_exists:
            missing_files.extend(model_files)
        
        if missing_files:
            logger.error(f"❌ Missing required model files: {missing_files}")
            return False
        
        logger.info("✅ Model files validated")
        return True
    
    def load_training_config(self) -> Dict[str, Any]:
        """Load training configuration from model directory"""
        config_files = [
            "training_config.json",
            "config_petite_llm_3_fr_1_20250727_152504.json",
            "config_petite_llm_3_fr_1_20250727_152524.json"
        ]
        
        for config_file in config_files:
            config_path = self.model_path / config_file
            if config_path.exists():
                with open(config_path, 'r') as f:
                    config = json.load(f)
                logger.info(f"✅ Loaded training config from: {config_file}")
                return config
        
        # Fallback to basic config
        logger.warning("⚠️ No training config found, using default")
        return {
            "model_name": "HuggingFaceTB/SmolLM3-3B",
            "dataset_name": "OpenHermes-FR",
            "training_config_type": "Custom Configuration",
            "trainer_type": "SFTTrainer",
            "per_device_train_batch_size": 8,
            "gradient_accumulation_steps": 16,
            "learning_rate": "5e-6",
            "num_train_epochs": 3,
            "max_seq_length": 2048,
            "dataset_size": "~80K samples",
            "dataset_format": "Chat format"
        }
    
    def load_training_results(self) -> Dict[str, Any]:
        """Load training results from model directory"""
        results_files = [
            "train_results.json",
            "training_summary_petite_llm_3_fr_1_20250727_152504.json",
            "training_summary_petite_llm_3_fr_1_20250727_152524.json"
        ]
        
        for results_file in results_files:
            results_path = self.model_path / results_file
            if results_path.exists():
                with open(results_path, 'r') as f:
                    results = json.load(f)
                logger.info(f"✅ Loaded training results from: {results_file}")
                return results
        
        # Fallback to basic results
        logger.warning("⚠️ No training results found, using default")
        return {
            "final_loss": "Unknown",
            "total_steps": "Unknown",
            "train_loss": "Unknown",
            "eval_loss": "Unknown"
        }
    
    def push_main_model(self) -> bool:
        """Push the main model to Hugging Face Hub"""
        try:
            logger.info("🚀 Pushing main model to Hugging Face Hub...")
            
            # Import push script
            from scripts.model_tonic.push_to_huggingface import HuggingFacePusher
            
            # Load training data
            training_config = self.load_training_config()
            training_results = self.load_training_results()
            
            # Initialize pusher
            pusher = HuggingFacePusher(
                model_path=str(self.model_path),
                repo_name=self.repo_name,
                token=self.hf_token,
                private=self.private,
                trackio_url=self.trackio_url,
                experiment_name=self.experiment_name,
                dataset_repo=self.dataset_repo,
                hf_token=self.hf_token,
                author_name=self.author_name,
                model_description=self.model_description
            )
            
            # Push model
            success = pusher.push_model(training_config, training_results)
            
            if success:
                logger.info(f"✅ Main model pushed successfully to: https://huggingface.co/{self.repo_name}")
                return True
            else:
                logger.error("❌ Failed to push main model")
                return False
                
        except Exception as e:
            logger.error(f"❌ Error pushing main model: {e}")
            return False
    
    def quantize_and_push_models(self) -> bool:
        """Quantize and push models to Hugging Face Hub"""
        if not self.quantize:
            logger.info("⏭️ Skipping quantization (disabled)")
            return True
        
        try:
            logger.info("🔄 Starting quantization and push process...")
            
            # Import quantization script
            from scripts.model_tonic.quantize_model import ModelQuantizer
            
            success_count = 0
            total_count = len(self.quant_types)
            
            for quant_type in self.quant_types:
                logger.info(f"🔄 Processing quantization type: {quant_type}")
                
                # Initialize quantizer
                quantizer = ModelQuantizer(
                    model_path=str(self.model_path),
                    repo_name=self.repo_name,
                    token=self.hf_token,
                    private=self.private,
                    trackio_url=self.trackio_url,
                    experiment_name=self.experiment_name,
                    dataset_repo=self.dataset_repo,
                    hf_token=self.hf_token
                )
                
                # Perform quantization and push
                success = quantizer.quantize_and_push(
                    quant_type=quant_type,
                    device="auto",
                    group_size=128
                )
                
                if success:
                    logger.info(f"✅ {quant_type} quantization and push completed")
                    success_count += 1
                else:
                    logger.error(f"❌ {quant_type} quantization and push failed")
            
            logger.info(f"📊 Quantization summary: {success_count}/{total_count} successful")
            return success_count > 0
                
        except Exception as e:
            logger.error(f"❌ Error during quantization: {e}")
            return False
    
    def run_complete_pipeline(self) -> bool:
        """Run the complete model recovery and deployment pipeline"""
        logger.info("🚀 Starting complete model recovery and deployment pipeline")
        
        # Step 1: Validate model path
        if not self.validate_model_path():
            logger.error("❌ Model validation failed")
            return False
        
        # Step 2: Push main model
        if not self.push_main_model():
            logger.error("❌ Main model push failed")
            return False
        
        # Step 3: Quantize and push models
        if not self.quantize_and_push_models():
            logger.warning("⚠️ Quantization failed, but main model was pushed successfully")
        
        logger.info("🎉 Model recovery and deployment pipeline completed!")
        logger.info(f"🌐 View your model at: https://huggingface.co/{self.repo_name}")
        
        return True

def parse_args():
    """Parse command line arguments"""
    parser = argparse.ArgumentParser(description='Recover and deploy trained model to Hugging Face Hub')
    
    # Required arguments
    parser.add_argument('model_path', type=str, help='Path to trained model directory')
    parser.add_argument('repo_name', type=str, help='Hugging Face repository name (username/repo-name)')
    
    # Optional arguments
    parser.add_argument('--hf-token', type=str, default=None, help='Hugging Face token')
    parser.add_argument('--private', action='store_true', help='Make repository private')
    parser.add_argument('--no-quantize', action='store_true', help='Skip quantization')
    parser.add_argument('--quant-types', nargs='+', 
                       choices=['int8_weight_only', 'int4_weight_only', 'int8_dynamic'],
                       default=['int8_weight_only', 'int4_weight_only'],
                       help='Quantization types to apply')
    parser.add_argument('--trackio-url', type=str, default=None, help='Trackio Space URL for logging')
    parser.add_argument('--experiment-name', type=str, default=None, help='Experiment name for Trackio')
    parser.add_argument('--dataset-repo', type=str, default=None, help='HF Dataset repository for experiment storage')
    parser.add_argument('--author-name', type=str, default=None, help='Author name for model card')
    parser.add_argument('--model-description', type=str, default=None, help='Model description for model card')
    
    return parser.parse_args()

def main():
    """Main function"""
    args = parse_args()
    
    # Setup logging
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    logger.info("Starting model recovery and deployment pipeline")
    
    # Initialize pipeline
    try:
        pipeline = ModelRecoveryPipeline(
            model_path=args.model_path,
            repo_name=args.repo_name,
            hf_token=args.hf_token,
            private=args.private,
            quantize=not args.no_quantize,
            quant_types=args.quant_types,
            trackio_url=args.trackio_url,
            experiment_name=args.experiment_name,
            dataset_repo=args.dataset_repo,
            author_name=args.author_name,
            model_description=args.model_description
        )
        
        # Run complete pipeline
        success = pipeline.run_complete_pipeline()
        
        if success:
            logger.info("✅ Model recovery and deployment completed successfully!")
            return 0
        else:
            logger.error("❌ Model recovery and deployment failed!")
        return 1
    
    except Exception as e:
        logger.error(f"❌ Error during model recovery: {e}")
        return 1

if __name__ == "__main__":
    exit(main())