add device detection utils

src/utils/__init__.py

@@ -0,0 +1,18 @@
+"""Utility functions for the audit assistant."""
+
+from .device import (
+    get_device,
+    get_device_for_sentence_transformers,
+    get_device_for_colpali,
+    log_device_info,
+    get_cached_device
+)
+
+__all__ = [
+    "get_device",
+    "get_device_for_sentence_transformers",
+    "get_device_for_colpali",
+    "log_device_info",
+    "get_cached_device"
+]
+

src/utils/device.py

@@ -0,0 +1,109 @@
+"""
+Device detection utility for automatic hardware acceleration.
+
+Automatically detects and returns the best available device:
+- MPS (Apple Silicon)
+- CUDA (NVIDIA GPU)
+- CPU (fallback)
+"""
+import torch
+import logging
+
+logger = logging.getLogger(__name__)
+
+def get_device() -> str:
+    """
+    Automatically detect the best available device.
+    
+    Priority:
+    1. MPS (Apple Silicon) - if available
+    2. CUDA (NVIDIA GPU) - if available
+    3. CPU - fallback
+    
+    Returns:
+        str: Device name ('mps', 'cuda', or 'cpu')
+    """
+    # Check for MPS (Apple Silicon)
+    mps_available = torch.backends.mps.is_available()
+    mps_built = torch.backends.mps.is_built()
+    
+    if mps_available and mps_built:
+        try:
+            # Test if MPS is actually usable
+            torch.zeros(1, device='mps')
+            logger.info("🍎 Using MPS (Apple Silicon) device")
+            return 'mps'
+        except Exception as e:
+            logger.warning(f"⚠️ MPS available but not usable: {e}")
+    
+    # Check for CUDA (NVIDIA GPU)
+    if torch.cuda.is_available():
+        logger.info(f"🎮 Using CUDA device (GPU: {torch.cuda.get_device_name(0)})")
+        return 'cuda'
+    
+    # Fallback to CPU
+    logger.info(f"💻 Using CPU device (no GPU acceleration) [MPS available={mps_available}, MPS built={mps_built}]")
+    return 'cpu'
+
+def get_device_for_sentence_transformers() -> str:
+    """
+    Get device string for sentence-transformers library.
+    
+    Returns:
+        str: Device name compatible with sentence-transformers
+    """
+    device = get_device()
+    # sentence-transformers uses same naming
+    return device
+
+def get_device_for_colpali() -> str:
+    """
+    Get device string for ColPali models.
+    
+    Returns:
+        str: Device name compatible with ColPali/transformers
+    """
+    device = get_device()
+    # ColPali/transformers uses same naming
+    return device
+
+def log_device_info():
+    """Log detailed device information"""
+    device = get_device()
+    
+    logger.info("=" * 70)
+    logger.info("🖥️  DEVICE CONFIGURATION")
+    logger.info("=" * 70)
+    logger.info(f"Selected device: {device.upper()}")
+    
+    if device == 'mps':
+        logger.info("   Type: Apple Silicon (M1/M2/M3)")
+        logger.info("   Acceleration: Metal Performance Shaders")
+    elif device == 'cuda':
+        logger.info(f"   Type: NVIDIA GPU")
+        logger.info(f"   GPU: {torch.cuda.get_device_name(0)}")
+        logger.info(f"   CUDA Version: {torch.version.cuda}")
+        logger.info(f"   Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
+    else:
+        logger.info("   Type: CPU (no GPU acceleration)")
+        logger.info("   Note: Processing will be slower")
+    
+    logger.info("=" * 70)
+    
+    return device
+
+# Convenience function to get device once at module level
+_cached_device = None
+
+def get_cached_device() -> str:
+    """
+    Get device with caching to avoid repeated checks.
+    
+    Returns:
+        str: Device name ('mps', 'cuda', or 'cpu')
+    """
+    global _cached_device
+    if _cached_device is None:
+        _cached_device = get_device()
+    return _cached_device
+