Add proper ZeroGPU configuration with spaces decorator

README.md

@@ -10,6 +10,7 @@ pinned: false
 hf_oauth: false
 hardware: a10g-small
 inference: true
+zerogpu: true
 ---
 
 # Code Review Assistant
@@ -28,7 +29,7 @@ An automated code review system powered by Gemma-2b that provides intelligent co
 - Uses Gemma-2b for intelligent code analysis
 - Tracks model performance and accuracy
 - Monitors response times and token usage
-- Optimized for both CPU and GPU execution
+- GPU-accelerated with ZeroGPU
 
 ### Performance Monitoring
 - Real-time metrics dashboard
@@ -53,8 +54,9 @@ The following environment variables need to be set in your Hugging Face Space:
 
 This Space uses:
 - Hardware: A10G Small
-- Memory: Optimized for both CPU and GPU
-- Fallback: Graceful degradation to CPU if needed
+- Runtime: ZeroGPU enabled
+- Memory: Dynamic GPU allocation
+- Optimization: @spaces.GPU decorator
 
 ## License
 

app.py

@@ -8,6 +8,7 @@ from datetime import datetime
 import json
 from typing import List, Dict
 import warnings
+import spaces
 
 # Filter out warnings
 warnings.filterwarnings('ignore')
@@ -45,6 +46,7 @@ class CodeReviewer:
         }
         self.initialize_model()
         
+    @spaces.GPU
     def initialize_model(self):
         """Initialize the model and tokenizer."""
         try:
@@ -60,37 +62,16 @@ class CodeReviewer:
             )
             
             logger.info("Loading model...")
-            # Initialize model with specific configuration
-            model_kwargs = {
-                "trust_remote_code": True,
-                "low_cpu_mem_usage": True,
-                "cache_dir": CACHE_DIR,
-                "token": HF_TOKEN
-            }
-            
-            # First try loading with CPU
-            try:
-                self.model = AutoModelForCausalLM.from_pretrained(
-                    MODEL_NAME,
-                    device_map=None,  # Start with no device map
-                    **model_kwargs
-                )
-                self.device = torch.device("cpu")
-                logger.info("Model loaded on CPU, will attempt GPU transfer")
-            except Exception as e1:
-                logger.error(f"Failed to load model on CPU: {e1}")
-                raise
-            
-            # Try moving to GPU if available
-            try:
-                if torch.cuda.is_available():
-                    logger.info("Moving model to GPU")
-                    self.model = self.model.to("cuda")
-                    self.device = torch.device("cuda")
-            except Exception as e2:
-                logger.warning(f"Could not move model to GPU: {e2}")
-                logger.info("Continuing with CPU")
-            
+            self.model = AutoModelForCausalLM.from_pretrained(
+                MODEL_NAME,
+                device_map="auto",
+                torch_dtype=torch.float16,
+                trust_remote_code=True,
+                low_cpu_mem_usage=True,
+                cache_dir=CACHE_DIR,
+                token=HF_TOKEN
+            )
+            self.device = next(self.model.parameters()).device
             logger.info(f"Model loaded successfully on {self.device}")
         except Exception as e:
             logger.error(f"Error initializing model: {e}")
@@ -109,6 +90,7 @@ Code:
 {code}
 ```"""
 
+    @spaces.GPU
     def review_code(self, code: str, language: str) -> str:
         """Perform code review using the model."""
         try:
@@ -161,7 +143,7 @@ Code:
             # Update metrics
             self.update_metrics(review)
             
-            # Clear GPU memory if using CUDA
+            # Clear GPU memory
             if self.device.type == "cuda":
                 del inputs, outputs
                 torch.cuda.empty_cache()
@@ -253,6 +235,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as iface:
             )
     
     # Set up event handlers
+    @spaces.GPU
     def review_code_interface(code: str, language: str) -> str:
         if not code.strip():
             return "Please enter some code to review."

requirements.txt

@@ -12,6 +12,7 @@ scipy>=1.11.0
 
 # Hugging Face
 huggingface-hub>=0.20.3
+spaces>=0.19.4
 
 # Utilities
 python-dotenv>=1.0.0