model_wrapper.py

"""
Model wrapper and configuration validation for the planning system.
Provides robust model loading with proper error handling and file validation.
"""

import os
import json
import logging
from dataclasses import dataclass
from typing import Optional, List, Dict, Any
from pathlib import Path

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@dataclass
class ModelWrapper:
    """
    Encapsulates model state and provides access to loaded model components.
    
    Attributes:
        model: The loaded PeftModel instance (None if not loaded)
        tokenizer: The loaded AutoTokenizer instance (None if not loaded)
        is_loaded: Whether the model was successfully loaded
        error_message: Error message if loading failed
        fallback_mode: Whether the system is operating in fallback mode
        model_path: Path to the model directory
        base_model_path: Path to the base model
    """
    model: Optional[Any] = None
    tokenizer: Optional[Any] = None
    is_loaded: bool = False
    error_message: Optional[str] = None
    fallback_mode: bool = False
    model_path: Optional[str] = None
    base_model_path: Optional[str] = None


@dataclass
class ValidationResult:
    """
    Result of configuration file validation.
    
    Attributes:
        is_valid: Whether all required files are present and valid
        missing_files: List of missing file paths
        corrupted_files: List of corrupted file paths
        repair_suggestions: List of suggestions for fixing issues
    """
    is_valid: bool
    missing_files: List[str]
    corrupted_files: List[str]
    repair_suggestions: List[str]


class ConfigurationValidator:
    """
    Validates model configuration files and provides detailed error reporting.
    """
    
    REQUIRED_FILES = ["adapter_config.json"]

    # PEFT may save either format depending on save settings/version.
    ADAPTER_WEIGHT_ALTERNATIVES = [
        "adapter_model.safetensors",
        "adapter_model.bin",
    ]
    
    OPTIONAL_FILES = [
        "README.md",
        "training_args.bin"
    ]
    
    @staticmethod
    def validate_config(config_path: str) -> ValidationResult:
        """
        Validate model configuration directory and files.
        
        Parameters
        ----------
        config_path : str
            Path to the model configuration directory
            
        Returns
        -------
        ValidationResult
            Detailed validation results with missing files and repair suggestions
        """
        config_path = Path(config_path)
        missing_files = []
        corrupted_files = []
        repair_suggestions = []
        
        logger.info(f"Validating configuration at: {config_path}")
        
        # Check if directory exists
        if not config_path.exists():
            missing_files.append(str(config_path))
            repair_suggestions.append(
                f"Model directory '{config_path}' does not exist. "
                f"Please ensure the model is properly installed or check the path."
            )
            return ValidationResult(
                is_valid=False,
                missing_files=missing_files,
                corrupted_files=corrupted_files,
                repair_suggestions=repair_suggestions
            )
        
        if not config_path.is_dir():
            corrupted_files.append(str(config_path))
            repair_suggestions.append(
                f"Path '{config_path}' exists but is not a directory. "
                f"Please ensure it points to a valid model directory."
            )
            return ValidationResult(
                is_valid=False,
                missing_files=missing_files,
                corrupted_files=corrupted_files,
                repair_suggestions=repair_suggestions
            )
        
        # Check required files
        for required_file in ConfigurationValidator.REQUIRED_FILES:
            file_path = config_path / required_file
            
            if not file_path.exists():
                missing_files.append(str(file_path))
                repair_suggestions.append(
                    f"Required file '{required_file}' is missing from '{config_path}'. "
                    f"Please reinstall the model or restore this file."
                )
            else:
                # Check file integrity
                try:
                    if required_file.endswith('.json'):
                        ConfigurationValidator._validate_json_file(file_path)
                    elif required_file.endswith('.safetensors'):
                        ConfigurationValidator._validate_safetensors_file(file_path)
                except Exception as e:
                    corrupted_files.append(str(file_path))
                    repair_suggestions.append(
                        f"File '{required_file}' appears to be corrupted: {str(e)}. "
                        f"Please reinstall the model or restore this file."
                    )

        weight_files = [config_path / name for name in ConfigurationValidator.ADAPTER_WEIGHT_ALTERNATIVES]
        existing_weight_files = [path for path in weight_files if path.exists()]
        if not existing_weight_files:
            missing_files.append(
                "one of: "
                + ", ".join(str(path) for path in weight_files)
            )
            repair_suggestions.append(
                "The LoRA directory must contain adapter_model.safetensors or adapter_model.bin."
            )
        else:
            for weight_file in existing_weight_files:
                try:
                    if weight_file.suffix == ".safetensors":
                        ConfigurationValidator._validate_safetensors_file(weight_file)
                    elif weight_file.stat().st_size < 1024:
                        raise ValueError(
                            f"File size ({weight_file.stat().st_size} bytes) is too small"
                        )
                except Exception as exc:
                    corrupted_files.append(str(weight_file))
                    repair_suggestions.append(
                        f"Adapter weights '{weight_file.name}' appear invalid: {exc}"
                    )
        
        # Validate adapter_config.json content if it exists
        adapter_config_path = config_path / "adapter_config.json"
        if adapter_config_path.exists() and str(adapter_config_path) not in corrupted_files:
            try:
                ConfigurationValidator._validate_adapter_config_content(adapter_config_path)
            except Exception as e:
                corrupted_files.append(str(adapter_config_path))
                repair_suggestions.append(
                    f"adapter_config.json has invalid content: {str(e)}. "
                    f"Please check the file format and required fields."
                )
        
        # Check for common issues and provide specific guidance
        if missing_files or corrupted_files:
            repair_suggestions.append(
                "To fix model loading issues:\n"
                "1. Verify the model path is correct\n"
                "2. Ensure all required files are present\n"
                "3. Check file permissions\n"
                "4. Consider re-downloading the model if files are corrupted\n"
                "5. Evaluation fails closed unless rule-based fallback is explicitly enabled"
            )
        
        is_valid = len(missing_files) == 0 and len(corrupted_files) == 0
        
        if is_valid:
            logger.info("Configuration validation passed")
        else:
            logger.warning(f"Configuration validation failed: {len(missing_files)} missing, {len(corrupted_files)} corrupted")
        
        return ValidationResult(
            is_valid=is_valid,
            missing_files=missing_files,
            corrupted_files=corrupted_files,
            repair_suggestions=repair_suggestions
        )
    
    @staticmethod
    def _validate_json_file(file_path: Path) -> None:
        """
        Validate that a JSON file can be loaded and parsed.
        
        Parameters
        ----------
        file_path : Path
            Path to the JSON file
            
        Raises
        ------
        Exception
            If the file cannot be parsed as valid JSON
        """
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                json.load(f)
        except json.JSONDecodeError as e:
            raise Exception(f"Invalid JSON format: {str(e)}")
        except Exception as e:
            raise Exception(f"Cannot read file: {str(e)}")
    
    @staticmethod
    def _validate_safetensors_file(file_path: Path) -> None:
        """
        Validate that a safetensors file exists and has reasonable size.
        
        Parameters
        ----------
        file_path : Path
            Path to the safetensors file
            
        Raises
        ------
        Exception
            If the file appears to be corrupted or invalid
        """
        try:
            stat = file_path.stat()
            # Check if file is too small (likely corrupted)
            if stat.st_size < 1024:  # Less than 1KB is suspicious for a model file
                raise Exception(f"File size ({stat.st_size} bytes) is too small for a model file")
        except Exception as e:
            raise Exception(f"Cannot access file: {str(e)}")
    
    @staticmethod
    def _validate_adapter_config_content(file_path: Path) -> None:
        """
        Validate the content of adapter_config.json for required fields.
        
        Parameters
        ----------
        file_path : Path
            Path to the adapter_config.json file
            
        Raises
        ------
        Exception
            If required fields are missing or invalid
        """
        with open(file_path, 'r', encoding='utf-8') as f:
            config = json.load(f)
        
        required_fields = ['peft_type', 'task_type']
        missing_fields = []
        
        for field in required_fields:
            if field not in config:
                missing_fields.append(field)
        
        if missing_fields:
            raise Exception(f"Missing required fields: {missing_fields}")
        
        # Validate peft_type
        if config.get('peft_type') != 'LORA':
            logger.warning(f"Unexpected peft_type: {config.get('peft_type')}, expected 'LORA'")


# Module-level cache for validation results to avoid repeated disk I/O
_validation_cache: Dict[str, ValidationResult] = {}


def clear_model_validation_cache(model_path: str | None = None) -> None:
    """Clear cached adapter validation after files are created or replaced."""
    if model_path is None:
        _validation_cache.clear()
        return
    _validation_cache.pop(str(Path(model_path).resolve()), None)


def validate_model_path(model_path: str) -> ValidationResult:
    """
    Convenience function to validate a model path.
    Results are cached in memory so subsequent calls for the same path skip
    all disk I/O (file-stat and JSON parse).

    Parameters
    ----------
    model_path : str
        Path to the model directory

    Returns
    -------
    ValidationResult
        Validation results
    """
    resolved = str(Path(model_path).resolve())
    if resolved not in _validation_cache:
        _validation_cache[resolved] = ConfigurationValidator.validate_config(model_path)
    return _validation_cache[resolved]


def create_model_wrapper(model_path: str, base_model_path: str = "Qwen/Qwen3-1.7B") -> ModelWrapper:
    """
    Create a ModelWrapper with initial validation but without loading the model.
    
    Parameters
    ----------
    model_path : str
        Path to the LoRA model directory
    base_model_path : str
        Path to the base model
        
    Returns
    -------
    ModelWrapper
        ModelWrapper instance with validation results
    """
    wrapper = ModelWrapper(
        model_path=model_path,
        base_model_path=base_model_path
    )
    
    # Validate configuration
    validation_result = validate_model_path(model_path)
    
    if not validation_result.is_valid:
        wrapper.error_message = f"Configuration validation failed: {validation_result.repair_suggestions[0] if validation_result.repair_suggestions else 'Unknown error'}"
        wrapper.fallback_mode = True
        logger.error(f"Model validation failed: {wrapper.error_message}")
    else:
        logger.info("Model configuration validation passed")
    
    return wrapper