Spaces:
Sleeping
Sleeping
Joash
commited on
Commit
·
cdd6dbd
1
Parent(s):
065d6d7
Improve model initialization with CPU-first approach
Browse files
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
|
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 |
-
#
|
78 |
-
|
79 |
-
|
80 |
-
|
81 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
82 |
|
83 |
-
#
|
84 |
-
|
85 |
-
|
86 |
-
|
87 |
-
|
88 |
-
|
|
|
|
|
|
|
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:
|