Update app.py

app.py

@@ -8,12 +8,24 @@ import spaces
 peft_model_id = "rootxhacker/CodeAstra-7B"
 config = PeftConfig.from_pretrained(peft_model_id)
 
-# Load the model without explicit device mapping
+# Function to move tensors to CPU
+def to_cpu(obj):
+    if isinstance(obj, torch.Tensor):
+        return obj.cpu()
+    elif isinstance(obj, list):
+        return [to_cpu(item) for item in obj]
+    elif isinstance(obj, tuple):
+        return tuple(to_cpu(item) for item in obj)
+    elif isinstance(obj, dict):
+        return {key: to_cpu(value) for key, value in obj.items()}
+    return obj
+
+# Load the model
 model = AutoModelForCausalLM.from_pretrained(
     config.base_model_name_or_path,
     return_dict=True,
     load_in_4bit=True,
-    device_map=None  # Let the Spaces environment handle device mapping
+    device_map='auto'
 )
 
 tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
@@ -24,12 +36,25 @@ model = PeftModel.from_pretrained(model, peft_model_id)
 @spaces.GPU(duration=200)
 def get_completion(query, model, tokenizer):
     try:
-        inputs = tokenizer(query, return_tensors="pt")
+        # Move model to CUDA
+        model = model.cuda()
+        
+        # Ensure input is on CUDA
+        inputs = tokenizer(query, return_tensors="pt").to('cuda')
+        
         with torch.no_grad():
             outputs = model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.7)
+        
+        # Move outputs to CPU before decoding
+        outputs = to_cpu(outputs)
+        
         return tokenizer.decode(outputs[0], skip_special_tokens=True)
     except Exception as e:
         return f"An error occurred: {str(e)}"
+    finally:
+        # Move model back to CPU to free up GPU memory
+        model = model.cpu()
+        torch.cuda.empty_cache()
 
 @spaces.GPU(duration=200)
 def code_review(code_to_analyze):
@@ -47,4 +72,4 @@ iface = gr.Interface(
 )
 
 # Launch the Gradio app
-iface.launch()
+iface.launch()