Joash commited on
Commit
cdd6dbd
·
1 Parent(s): 065d6d7

Improve model initialization with CPU-first approach

Browse files
Files changed (1) hide show
  1. app.py +23 -19
app.py CHANGED
@@ -24,10 +24,6 @@ MODEL_NAME = os.getenv("MODEL_NAME", "google/gemma-2b-it")
24
  CACHE_DIR = "/home/user/.cache/huggingface"
25
  os.makedirs(CACHE_DIR, exist_ok=True)
26
 
27
- # Enable ZeroGPU features
28
- os.environ["ZERO_GPU"] = "1"
29
- os.environ["CUDA_VISIBLE_DEVICES"] = "0"
30
-
31
  class Review:
32
  def __init__(self, code: str, language: str, suggestions: str):
33
  self.code = code
@@ -64,28 +60,36 @@ class CodeReviewer:
64
  )
65
 
66
  logger.info("Loading model...")
67
- # Initialize model with ZeroGPU configuration
68
  model_kwargs = {
69
- "torch_dtype": torch.float16,
70
  "trust_remote_code": True,
71
  "low_cpu_mem_usage": True,
72
  "cache_dir": CACHE_DIR,
73
- "token": HF_TOKEN,
74
- "device_map": {"": 0} # Use first GPU
75
  }
76
 
77
- # Load model with ZeroGPU support
78
- self.model = AutoModelForCausalLM.from_pretrained(
79
- MODEL_NAME,
80
- **model_kwargs
81
- )
 
 
 
 
 
 
 
82
 
83
- # Set device based on ZeroGPU availability
84
- self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
85
- if self.device.type == "cpu":
86
- logger.warning("Running on CPU - ZeroGPU not initialized")
87
- else:
88
- logger.info("ZeroGPU initialized successfully")
 
 
 
89
 
90
  logger.info(f"Model loaded successfully on {self.device}")
91
  except Exception as e:
 
24
  CACHE_DIR = "/home/user/.cache/huggingface"
25
  os.makedirs(CACHE_DIR, exist_ok=True)
26
 
 
 
 
 
27
  class Review:
28
  def __init__(self, code: str, language: str, suggestions: str):
29
  self.code = code
 
60
  )
61
 
62
  logger.info("Loading model...")
63
+ # Initialize model with specific configuration
64
  model_kwargs = {
 
65
  "trust_remote_code": True,
66
  "low_cpu_mem_usage": True,
67
  "cache_dir": CACHE_DIR,
68
+ "token": HF_TOKEN
 
69
  }
70
 
71
+ # First try loading with CPU
72
+ try:
73
+ self.model = AutoModelForCausalLM.from_pretrained(
74
+ MODEL_NAME,
75
+ device_map=None, # Start with no device map
76
+ **model_kwargs
77
+ )
78
+ self.device = torch.device("cpu")
79
+ logger.info("Model loaded on CPU, will attempt GPU transfer")
80
+ except Exception as e1:
81
+ logger.error(f"Failed to load model on CPU: {e1}")
82
+ raise
83
 
84
+ # Try moving to GPU if available
85
+ try:
86
+ if torch.cuda.is_available():
87
+ logger.info("Moving model to GPU")
88
+ self.model = self.model.to("cuda")
89
+ self.device = torch.device("cuda")
90
+ except Exception as e2:
91
+ logger.warning(f"Could not move model to GPU: {e2}")
92
+ logger.info("Continuing with CPU")
93
 
94
  logger.info(f"Model loaded successfully on {self.device}")
95
  except Exception as e: