Update app to properly utilize ZeroGPU with GPU optimizations

app.py

@@ -28,7 +28,8 @@ class CodeReviewer:
     def __init__(self):
         self.model = None
         self.tokenizer = None
-        self.device = "cpu"
+        # Let ZeroGPU handle GPU allocation
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.review_history: List[Review] = []
         self.metrics = {
             'total_reviews': 0,
@@ -41,19 +42,25 @@ class CodeReviewer:
         """Initialize the model and tokenizer."""
         try:
             if HF_TOKEN:
-                login(token=HF_TOKEN)
+                login(token=HF_TOKEN, add_to_git_credential=False)
             
             logger.info("Loading tokenizer...")
-            self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                MODEL_NAME,
+                token=HF_TOKEN,
+                trust_remote_code=True
+            )
             
             logger.info("Loading model...")
+            # Let ZeroGPU handle device mapping
             self.model = AutoModelForCausalLM.from_pretrained(
                 MODEL_NAME,
-                device_map={"": self.device},
-                torch_dtype=torch.float32,
-                low_cpu_mem_usage=True
+                token=HF_TOKEN,
+                device_map="auto",
+                torch_dtype=torch.float16,  # Use fp16 for GPU
+                trust_remote_code=True
             )
-            logger.info("Model loaded successfully")
+            logger.info(f"Model loaded successfully on {self.device}")
         except Exception as e:
             logger.error(f"Error initializing model: {e}")
             raise
@@ -83,7 +90,7 @@ Code:
                 truncation=True,
                 max_length=512,
                 padding=True
-            )
+            ).to(self.device)  # Move inputs to GPU
             
             with torch.no_grad():
                 outputs = self.model.generate(
@@ -108,6 +115,11 @@ Code:
             # Update metrics
             self.update_metrics(review)
             
+            # Clear GPU memory
+            if torch.cuda.is_available():
+                del inputs, outputs
+                torch.cuda.empty_cache()
+            
             return suggestions
             
         except Exception as e:
@@ -148,43 +160,10 @@ Code:
         return {
             'Total Reviews': self.metrics['total_reviews'],
             'Average Response Time': f"{self.metrics['avg_response_time']:.2f}s",
-            'Reviews Today': self.metrics['reviews_today']
+            'Reviews Today': self.metrics['reviews_today'],
+            'Device': str(self.device)
         }
 
-# Initialize the reviewer
-reviewer = CodeReviewer()
-
-def review_code_interface(code: str, language: str) -> str:
-    """Gradio interface function for code review."""
-    if not code.strip():
-        return "Please enter some code to review."
-    
-    try:
-        result = reviewer.review_code(code, language)
-        return result
-    except Exception as e:
-        return f"Error: {str(e)}"
-
-def get_history_interface() -> str:
-    """Format history for display."""
-    history = reviewer.get_history()
-    if not history:
-        return "No reviews yet."
-    
-    result = ""
-    for review in history:
-        result += f"Time: {review['timestamp']}\n"
-        result += f"Language: {review['language']}\n"
-        result += f"Response Time: {review['response_time']}\n"
-        result += "Code:\n```\n" + review['code'] + "\n```\n"
-        result += "Suggestions:\n" + review['suggestions'] + "\n"
-        result += "-" * 80 + "\n\n"
-    return result
-
-def get_metrics_interface() -> Dict:
-    """Get metrics for display."""
-    return reviewer.get_metrics()
-
 # Create Gradio interface
 with gr.Blocks(theme=gr.themes.Soft()) as iface:
     gr.Markdown("# Code Review Assistant")
@@ -215,18 +194,44 @@ with gr.Blocks(theme=gr.themes.Soft()) as iface:
             refresh_history = gr.Button("Refresh History")
             history_output = gr.Textbox(
                 label="Review History",
-                lines=20,
-                value=get_history_interface()
+                lines=20
             )
         
         with gr.Tab("Metrics"):
             refresh_metrics = gr.Button("Refresh Metrics")
             metrics_output = gr.JSON(
-                label="Performance Metrics",
-                value=get_metrics_interface()
+                label="Performance Metrics"
             )
     
+    # Initialize reviewer
+    reviewer = CodeReviewer()
+    
     # Set up event handlers
+    def review_code_interface(code: str, language: str) -> str:
+        if not code.strip():
+            return "Please enter some code to review."
+        try:
+            return reviewer.review_code(code, language)
+        except Exception as e:
+            return f"Error: {str(e)}"
+    
+    def get_history_interface() -> str:
+        history = reviewer.get_history()
+        if not history:
+            return "No reviews yet."
+        result = ""
+        for review in history:
+            result += f"Time: {review['timestamp']}\n"
+            result += f"Language: {review['language']}\n"
+            result += f"Response Time: {review['response_time']}\n"
+            result += "Code:\n```\n" + review['code'] + "\n```\n"
+            result += "Suggestions:\n" + review['suggestions'] + "\n"
+            result += "-" * 80 + "\n\n"
+        return result
+    
+    def get_metrics_interface() -> Dict:
+        return reviewer.get_metrics()
+    
     submit_btn.click(
         review_code_interface,
         inputs=[code_input, language_input],