chore: device optimization

app.py

@@ -24,13 +24,21 @@ class MedGemmaSymptomAnalyzer:
         logger.info("Initializing MedGemma Symptom Analyzer...")
     
     def load_model(self):
-        """Load MedGemma model with optimizations for deployment"""
+        """Load MedGemma model with optimizations for deployment and CPU compatibility"""
         if self.model_loaded:
             return True
             
         model_name = "google/medgemma-4b-it"
         logger.info(f"Loading model: {model_name}")
         
+        # Detect available device and log system info
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        logger.info(f"Device detected: {device}")
+        if device == "cpu":
+            logger.info(f"CPU threads available: {torch.get_num_threads()}")
+        else:
+            logger.info(f"CUDA device: {torch.cuda.get_device_name()}")
+        
         try:
             # Get HF token from environment (set in Hugging Face Spaces secrets)
             hf_token = os.getenv("HF_TOKEN")
@@ -39,33 +47,124 @@ class MedGemmaSymptomAnalyzer:
             else:
                 logger.warning("HF_TOKEN not found in environment variables")
             
-            # First try without quantization for CPU compatibility
+            # Optimize settings based on device
+            if device == "cpu":
+                logger.info("Configuring for CPU-optimized loading...")
+                torch_dtype = torch.float32
+                device_map = "cpu"
+                # Set optimal number of threads for CPU inference
+                torch.set_num_threads(max(1, torch.get_num_threads() // 2))
+                
+                # Additional CPU optimizations
+                import psutil
+                available_memory_gb = psutil.virtual_memory().available / (1024**3)
+                logger.info(f"Available memory: {available_memory_gb:.1f} GB")
+                
+                # Enable memory-efficient loading for low-memory systems
+                cpu_loading_kwargs = {
+                    "low_cpu_mem_usage": True,
+                    "torch_dtype": torch_dtype,
+                    "device_map": device_map
+                }
+                
+                # Use offloading for very low memory systems (< 8GB available)
+                if available_memory_gb < 8:
+                    logger.warning("Low memory detected, enabling aggressive memory optimizations")
+                    cpu_loading_kwargs.update({
+                        "offload_folder": "/tmp/model_offload",
+                        "offload_state_dict": True
+                    })
+            else:
+                logger.info("Configuring for GPU loading...")
+                torch_dtype = torch.float16
+                device_map = "auto"
+                cpu_loading_kwargs = {
+                    "torch_dtype": torch_dtype,
+                    "device_map": device_map,
+                    "low_cpu_mem_usage": True
+                }
+            
             logger.info("Loading tokenizer...")
             self.tokenizer = AutoTokenizer.from_pretrained(
                 model_name,
-                token=hf_token
+                token=hf_token,
+                use_fast=True  # Use fast tokenizer for better performance
             )
             
-            logger.info("Loading model...")
-            # Simplified loading for CPU/compatibility
+            logger.info(f"Loading model with dtype={torch_dtype}, device_map={device_map}...")
             self.model = AutoModelForCausalLM.from_pretrained(
                 model_name,
-                torch_dtype=torch.float32,  # Use float32 for CPU
-                device_map="cpu",  # Force CPU for compatibility
-                low_cpu_mem_usage=True,
-                token=hf_token
+                token=hf_token,
+                trust_remote_code=False,  # Security best practice
+                **cpu_loading_kwargs
             )
             
+            # Ensure pad token is set
             if self.tokenizer.pad_token is None:
                 self.tokenizer.pad_token = self.tokenizer.eos_token
+            
+            # Move model to appropriate device if needed
+            if device == "cpu" and hasattr(self.model, 'to'):
+                self.model = self.model.to('cpu')
+                logger.info("Model moved to CPU")
                 
             self.model_loaded = True
-            logger.info("Model loaded successfully!")
+            logger.info(f"Model loaded successfully on {device}!")
             return True
                 
+        except torch.cuda.OutOfMemoryError as e:
+            logger.error(f"GPU out of memory: {str(e)}")
+            logger.info("Attempting CPU fallback due to GPU memory constraints...")
+            try:
+                # Force CPU loading if GPU fails
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    model_name,
+                    token=hf_token,
+                    trust_remote_code=False,
+                    torch_dtype=torch.float32,
+                    device_map="cpu",
+                    low_cpu_mem_usage=True
+                )
+                self.model_loaded = True
+                logger.info("Model loaded successfully on CPU after GPU failure!")
+                return True
+            except Exception as fallback_e:
+                logger.error(f"CPU fallback also failed: {str(fallback_e)}")
+                self.model = None
+                self.tokenizer = None
+                self.model_loaded = False
+                return False
+        except ImportError as e:
+            logger.error(f"Missing dependency for model loading: {str(e)}")
+            logger.info("Please ensure all required packages are installed: pip install -r requirements.txt")
+            self.model = None
+            self.tokenizer = None
+            self.model_loaded = False
+            return False
+        except OSError as e:
+            if "disk quota exceeded" in str(e).lower() or "no space left" in str(e).lower():
+                logger.error("Insufficient disk space for model loading")
+                logger.info("Please free up disk space and try again")
+            elif "connection" in str(e).lower() or "timeout" in str(e).lower():
+                logger.error("Network connection issue during model download")
+                logger.info("Please check your internet connection and try again")
+            else:
+                logger.error(f"OS error during model loading: {str(e)}")
+            self.model = None
+            self.tokenizer = None
+            self.model_loaded = False
+            return False
         except Exception as e:
             logger.error(f"Failed to load model {model_name}: {str(e)}", exc_info=True)
             logger.warning("Falling back to demo mode due to model loading failure")
+            
+            # Provide helpful troubleshooting info
+            if device == "cpu":
+                logger.info("CPU loading troubleshooting tips:")
+                logger.info("- Ensure sufficient RAM (minimum 8GB recommended)")
+                logger.info("- Check that PyTorch CPU version is installed")
+                logger.info("- Verify HuggingFace token is valid")
+            
             self.model = None
             self.tokenizer = None
             self.model_loaded = False