src/model_loader.py

# src/model_loader.py
import logging
import os
import time

import requests
import torch
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from PIL import Image

from src.constants import CLASS_NAMES, HF_MODEL_FILENAME, HF_REPO_ID
from src.utils import preprocess_image

logger = logging.getLogger(__name__)

HF_REVISION = "217a9639ec46e2c5fd241973433c6ad69f984f54"


class ModelLoader:
    """A class to load and manage a trained PyTorch model for making predictions.

    This class handles model loading, device management, and inference operations
    for image classification tasks.

    Attributes
    ----------
        model_path (str): Path to the saved model file.
        device (torch.device): Device to run the model on (CPU or CUDA).
        model: The loaded PyTorch model.
        class_names (list): List of class names for classification.

    """

    def __init__(self):
        """Initialize the ModelLoader by downloading model from Hugging Face."""
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = None
        self.class_names = list(CLASS_NAMES.keys())

        # Get HF token from environment
        self.hf_token = os.getenv("HF_TOKEN")
        if self.hf_token:
            logger.info("HuggingFace token found, using authenticated access")
        else:
            logger.warning("No HF_TOKEN found, using anonymous access (rate limited)")

        logger.info(f"Using device: {self.device}")

        # Download and load model from HF with retry logic
        try:
            logger.info(f"Downloading model from Hugging Face: {HF_REPO_ID}")
            self.model_path = self._download_with_retry(HF_REPO_ID, HF_MODEL_FILENAME)
            logger.info(f"Model downloaded to: {self.model_path}")
            self._load_model()
        except Exception as e:
            logger.error(f"Failed to download model from Hugging Face: {e}")
            raise e

    def _download_with_retry(
        self, repo_id: str, filename: str, max_retries: int = 3, delay: float = 2.0
    ) -> str:
        """Download model with retry logic for network failures."""
        for attempt in range(max_retries):
            try:
                return hf_hub_download(
                    repo_id=repo_id,
                    filename=filename,
                    revision="217a9639ec46e2c5fd241973433c6ad69f984f54",
                    token=self.hf_token,
                )
            except requests.HTTPError as e:
                if e.response.status_code == 401:
                    raise ValueError(
                        "Authentication failed. Check HF_TOKEN environment variable"
                    ) from e
                elif e.response.status_code == 429:
                    # Rate limited - use longer delay
                    if attempt == max_retries - 1:
                        raise ValueError(
                            "Rate limited by HuggingFace. Try again later"
                        ) from e
                    logger.warning(f"Rate limited. Retrying in {delay * 5} seconds...")
                    time.sleep(delay * 5)
                    continue
                elif e.response.status_code >= 500:
                    # Server error - retry
                    if attempt == max_retries - 1:
                        raise e
                    logger.warning(
                        f"Server error {e.response.status_code}. Retrying..."
                    )
                    time.sleep(delay * (attempt + 1))
                else:
                    # Other HTTP errors - don't retry
                    raise e
            except (ConnectionError, TimeoutError) as e:
                if attempt == max_retries - 1:
                    raise e
                logger.warning(f"Network error: {e}. Retrying in {delay} seconds...")
                time.sleep(delay * (attempt + 1))
            except Exception as e:
                # Handle repository not found and other general errors
                if "not found" in str(e).lower() or "access denied" in str(e).lower():
                    raise ValueError(
                        f"Repository {repo_id} not found or access denied"
                    ) from e
                # For other unknown errors, retry
                if attempt == max_retries - 1:
                    raise e
                logger.warning(
                    f"Download attempt {attempt + 1} failed: {e}. Retrying..."
                )
                time.sleep(delay * (attempt + 1))

        raise Exception("All download attempts failed")

    def _load_model(self):
        """Load the trained model."""
        try:
            logger.info(f"Loading model from {self.model_path}")

            # Load model
            self.model = torch.load(  # nosec B614
                self.model_path, map_location=self.device, weights_only=False
            )
            self.model.to(self.device)
            self.model.eval()

            logger.info("Model loaded and set to evaluation mode")

            # Test model with dummy input
            self._validate_model()
            # Prevent slow first requests caused by model initialization.
            self._warmup_model()

        except Exception as e:
            logger.error(f"Failed to load model: {e}")
            raise e

    def _validate_model(self):
        """Validate model has expected architecture and outputs."""
        try:
            # Test with dummy input to check output shape
            dummy_input = torch.randn(1, 3, 224, 224).to(self.device)
            with torch.no_grad():
                output = self.model(dummy_input)

            # Validate output shape matches expected number of classes
            expected_classes = len(self.class_names)
            if output.shape[1] != expected_classes:
                raise ValueError(
                    f"Model output shape {output.shape} "
                    f"doesn't match expected classes {expected_classes}"
                )

            # Validate output is proper logits (not already softmaxed)
            if torch.any(output < 0):
                logger.info("Model outputs raw logits (expected)")
            else:
                logger.warning("Model might output probabilities instead of logits")

            logger.info(f"Model validation successful. Output shape: {output.shape}")

        except Exception as e:
            logger.error(f"Model validation failed: {e}")
            raise e

    def _warmup_model(self):
        """Warm up the model with multiple dummy predictions to optimize performance."""
        logger.info("Warming up model for optimal performance...")
        try:
            # Run multiple warmup predictions with different batch sizes
            dummy_inputs = [
                torch.randn(1, 3, 224, 224).to(self.device),
                torch.randn(2, 3, 224, 224).to(self.device),  # Batch of 2
            ]

            with torch.no_grad():
                for i, dummy_input in enumerate(dummy_inputs):
                    start_time = time.time()
                    _ = self.model(dummy_input)
                    warmup_time = time.time() - start_time
                    logger.info(f"Warmup {i + 1} completed in {warmup_time:.3f}s")

            logger.info("Model warmup completed successfully")

        except Exception as e:
            logger.warning(f"Model warmup failed: {e}. First prediction may be slower.")

    def predict(self, image: Image.Image) -> dict:
        """Make prediction on a single image."""
        if self.model is None:
            raise ValueError("Model not loaded")

        try:
            # Preprocess image
            input_tensor = preprocess_image(image).unsqueeze(0).to(self.device)

            # Make prediction
            with torch.no_grad():
                outputs = self.model(input_tensor)
                probabilities = F.softmax(outputs, dim=1)
                confidence, predicted_idx = torch.max(probabilities, 1)

                # Convert to CPU and numpy for JSON serialization
                probabilities = probabilities.cpu().numpy()[0]
                confidence = confidence.cpu().item()
                predicted_idx = predicted_idx.cpu().item()

                # Create results dictionary
                predicted_class = self.class_names[predicted_idx]

                all_probabilities = {
                    self.class_names[i]: float(prob)
                    for i, prob in enumerate(probabilities)
                }

                return {
                    "predicted_class": predicted_class,
                    "confidence": confidence,
                    "probabilities": all_probabilities,
                    "predicted_index": predicted_idx,
                }

        except Exception as e:
            logger.error(f"Prediction error: {e}")
            raise e