Joash commited on
Commit
b1ededf
·
1 Parent(s): 8aef6ee

Configure app for ZeroGPU dynamic GPU allocation

Browse files
Files changed (1) hide show
  1. app.py +23 -7
app.py CHANGED
@@ -24,6 +24,10 @@ 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
  class Review:
28
  def __init__(self, code: str, language: str, suggestions: str):
29
  self.code = code
@@ -36,7 +40,7 @@ class CodeReviewer:
36
  def __init__(self):
37
  self.model = None
38
  self.tokenizer = None
39
- self.device = "cpu" # Default to CPU
40
  self.review_history: List[Review] = []
41
  self.metrics = {
42
  'total_reviews': 0,
@@ -60,21 +64,28 @@ class CodeReviewer:
60
  )
61
 
62
  logger.info("Loading model...")
63
- # Initialize model with specific configuration
64
  model_kwargs = {
65
- "torch_dtype": torch.float32, # Use float32 for CPU
66
  "trust_remote_code": True,
67
  "low_cpu_mem_usage": True,
68
  "cache_dir": CACHE_DIR,
69
- "token": HF_TOKEN
 
70
  }
71
 
72
- # Load model directly to CPU
73
  self.model = AutoModelForCausalLM.from_pretrained(
74
  MODEL_NAME,
75
- device_map=None, # Don't use device_map
76
  **model_kwargs
77
- ).to(self.device)
 
 
 
 
 
 
 
78
 
79
  logger.info(f"Model loaded successfully on {self.device}")
80
  except Exception as e:
@@ -146,6 +157,11 @@ Code:
146
  # Update metrics
147
  self.update_metrics(review)
148
 
 
 
 
 
 
149
  return suggestions
150
 
151
  except Exception as e:
 
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
 
40
  def __init__(self):
41
  self.model = None
42
  self.tokenizer = None
43
+ self.device = None
44
  self.review_history: List[Review] = []
45
  self.metrics = {
46
  'total_reviews': 0,
 
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:
 
157
  # Update metrics
158
  self.update_metrics(review)
159
 
160
+ # Clear GPU memory if using CUDA
161
+ if self.device.type == "cuda":
162
+ del inputs, outputs
163
+ torch.cuda.empty_cache()
164
+
165
  return suggestions
166
 
167
  except Exception as e: