Update app.py

app.py

@@ -1,7 +1,8 @@
 import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel, PeftConfig
+import spaces
 
 # Load model and tokenizer
 MODEL_PATH = "sagar007/phi2_finetune"
@@ -9,16 +10,9 @@ MODEL_PATH = "sagar007/phi2_finetune"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token
 
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.float16,
-    bnb_4bit_use_double_quant=False
-)
-
 base_model = AutoModelForCausalLM.from_pretrained(
     "microsoft/phi-2",
-    quantization_config=bnb_config,
+    torch_dtype=torch.float32,  # Use float32 for CPU
     device_map="auto",
     trust_remote_code=True
 )
@@ -27,9 +21,10 @@ peft_config = PeftConfig.from_pretrained(MODEL_PATH)
 model = PeftModel.from_pretrained(base_model, MODEL_PATH)
 model.eval()
 
+@spaces.GPU(duration=60)
 def generate_response(instruction, max_length=512):
     prompt = f"Instruction: {instruction}\nResponse:"
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    inputs = tokenizer(prompt, return_tensors="pt")
     
     with torch.no_grad():
         outputs = model.generate(