Add ZeroGPU support and update dependencies

app.py

@@ -7,6 +7,9 @@ import logging
 from datetime import datetime
 import json
 from typing import List, Dict
+from huggingface_hub import HfApi
+from huggingface_hub.spaces import SpaceHardware, SpaceStage
+from huggingface_hub.spaces.space_sdk import SpaceRuntime
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -16,6 +19,16 @@ logger = logging.getLogger(__name__)
 HF_TOKEN = os.getenv("HUGGING_FACE_TOKEN")
 MODEL_NAME = os.getenv("MODEL_NAME", "google/gemma-2b-it")
 
+# Initialize Hugging Face API
+api = HfApi()
+
+# Space hardware configuration
+space_config = {
+    "hardware": SpaceHardware.A10G_SMALL,
+    "stage": SpaceStage.RUNTIME,
+    "runtime": SpaceRuntime.ZEROGPU
+}
+
 class Review:
     def __init__(self, code: str, language: str, suggestions: str):
         self.code = code
@@ -28,8 +41,7 @@ class CodeReviewer:
     def __init__(self):
         self.model = None
         self.tokenizer = None
-        # Let ZeroGPU handle GPU allocation
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = None
         self.review_history: List[Review] = []
         self.metrics = {
             'total_reviews': 0,
@@ -52,14 +64,17 @@ class CodeReviewer:
             )
             
             logger.info("Loading model...")
-            # Let ZeroGPU handle device mapping
+            # Initialize with ZeroGPU configuration
             self.model = AutoModelForCausalLM.from_pretrained(
                 MODEL_NAME,
                 token=HF_TOKEN,
                 device_map="auto",
-                torch_dtype=torch.float16,  # Use fp16 for GPU
-                trust_remote_code=True
+                torch_dtype=torch.float16,
+                trust_remote_code=True,
+                low_cpu_mem_usage=True,
+                use_zerogpu=True  # Enable ZeroGPU
             )
+            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}")
@@ -164,6 +179,13 @@ Code:
             'Device': str(self.device)
         }
 
+# Initialize reviewer with ZeroGPU
+@space_config
+def create_reviewer():
+    return CodeReviewer()
+
+reviewer = create_reviewer()
+
 # Create Gradio interface
 with gr.Blocks(theme=gr.themes.Soft()) as iface:
     gr.Markdown("# Code Review Assistant")
@@ -203,9 +225,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as iface:
                 label="Performance Metrics"
             )
     
-    # Initialize reviewer
-    reviewer = CodeReviewer()
-    
     # Set up event handlers
     def review_code_interface(code: str, language: str) -> str:
         if not code.strip():
@@ -264,4 +283,9 @@ with gr.Blocks(theme=gr.themes.Soft()) as iface:
 
 # Launch the app
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch(
+        share=False,
+        server_name="0.0.0.0",
+        server_port=7860,
+        enable_queue=True
+    )

requirements.txt

@@ -6,6 +6,11 @@ accelerate>=0.27.2
 safetensors>=0.4.2
 sentencepiece>=0.1.99
 
+# Hugging Face
+huggingface-hub>=0.20.3
+huggingface_hub[cli]>=0.20.3
+huggingface_hub[space]>=0.20.3
+
 # Utilities
 python-dotenv>=1.0.0
 pydantic>=2.4.2