Improve model loading with fallback options and memory settings

app.py

@@ -9,8 +9,8 @@ import json
 from typing import List, Dict
 import warnings
 
-# Filter CUDA warnings
-warnings.filterwarnings('ignore', category=UserWarning, message='Can\'t initialize NVML')
+# Filter out CUDA/NVML warnings
+warnings.filterwarnings('ignore', category=UserWarning)
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -24,6 +24,11 @@ MODEL_NAME = os.getenv("MODEL_NAME", "google/gemma-2b-it")
 CACHE_DIR = "/home/user/.cache/huggingface"
 os.makedirs(CACHE_DIR, exist_ok=True)
 
+# Set environment variables for GPU
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512"
+
 class Review:
     def __init__(self, code: str, language: str, suggestions: str):
         self.code = code
@@ -36,7 +41,7 @@ class CodeReviewer:
     def __init__(self):
         self.model = None
         self.tokenizer = None
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = None
         self.review_history: List[Review] = []
         self.metrics = {
             'total_reviews': 0,
@@ -62,7 +67,6 @@ class CodeReviewer:
             logger.info("Loading model...")
             # Initialize model with specific configuration
             model_kwargs = {
-                "device_map": "auto",
                 "torch_dtype": torch.float16,
                 "trust_remote_code": True,
                 "low_cpu_mem_usage": True,
@@ -70,20 +74,31 @@ class CodeReviewer:
                 "token": HF_TOKEN
             }
             
-            # Load model with error handling
+            # Try loading with different configurations
             try:
+                # First try with device_map="auto"
                 self.model = AutoModelForCausalLM.from_pretrained(
                     MODEL_NAME,
+                    device_map="auto",
                     **model_kwargs
                 )
-            except Exception as model_error:
-                logger.error(f"Error loading model: {model_error}")
-                # Try loading with safetensors
-                model_kwargs["use_safetensors"] = True
-                self.model = AutoModelForCausalLM.from_pretrained(
-                    MODEL_NAME,
-                    **model_kwargs
-                )
+                self.device = next(self.model.parameters()).device
+            except Exception as e1:
+                logger.warning(f"Failed to load with device_map='auto': {e1}")
+                try:
+                    # Try with specific device
+                    if torch.cuda.is_available():
+                        self.device = torch.device("cuda:0")
+                    else:
+                        self.device = torch.device("cpu")
+                    model_kwargs["device_map"] = None
+                    self.model = AutoModelForCausalLM.from_pretrained(
+                        MODEL_NAME,
+                        **model_kwargs
+                    ).to(self.device)
+                except Exception as e2:
+                    logger.error(f"Failed to load model on specific device: {e2}")
+                    raise
             
             logger.info(f"Model loaded successfully on {self.device}")
         except Exception as e: