Spaces:
Sleeping
Sleeping
Joash
commited on
Commit
·
b1ededf
1
Parent(s):
8aef6ee
Configure app for ZeroGPU dynamic GPU allocation
Browse files
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 =
|
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
|
64 |
model_kwargs = {
|
65 |
-
"torch_dtype": torch.
|
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
|
73 |
self.model = AutoModelForCausalLM.from_pretrained(
|
74 |
MODEL_NAME,
|
75 |
-
device_map=None, # Don't use device_map
|
76 |
**model_kwargs
|
77 |
-
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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:
|