Update app.py

app.py

@@ -2,46 +2,49 @@ import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel, PeftConfig
-import spaces
+from huggingface_hub import InferenceClient
 
-# Check if CUDA is available and set the device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {device}")
-
-# Load model and tokenizer#
+# Load configuration
 MODEL_PATH = "sagar007/phi2_25k"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
-tokenizer.pad_token = tokenizer.eos_token
+peft_config = PeftConfig.from_pretrained(MODEL_PATH)
 
-base_model = AutoModelForCausalLM.from_pretrained(
-    "microsoft/phi-2",
-    torch_dtype=torch.float16 if device.type == "cuda" else torch.float32,
-    device_map="auto",
-    trust_remote_code=True
-)
+# Initialize client for Zero-GPU environment
+client = InferenceClient()
 
-peft_config = PeftConfig.from_pretrained(MODEL_PATH)
-model = PeftModel.from_pretrained(base_model, MODEL_PATH)
-model.to(device)
-model.eval()
+def load_model():
+    # Load base model
+    base_model = AutoModelForCausalLM.from_pretrained(
+        "microsoft/phi-2",
+        torch_dtype=torch.float16,
+        device_map="auto",
+        trust_remote_code=True
+    )
+    
+    # Load PEFT model
+    model = PeftModel.from_pretrained(base_model, MODEL_PATH)
+    return model, AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
 
-@spaces.GPU(duration=60)
+@client.gpu(timeout=120)
 def generate_response(instruction, max_length=512):
-    prompt = f"Instruction: {instruction}\nResponse:"
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_length=max_length,
-            num_return_sequences=1,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True
-        )
-    
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response.split("Response:")[1].strip()
+    try:
+        model, tokenizer = load_model()
+        prompt = f"Instruction: {instruction}\nResponse:"
+        
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_length=max_length,
+                temperature=0.7,
+                top_p=0.9,
+                do_sample=True
+            )
+            
+        return tokenizer.decode(outputs[0], skip_special_tokens=True).split("Response:")[-1].strip()
+    except Exception as e:
+        print(f"Error: {str(e)}")
+        return "Sorry, I encountered an error. Please try again."
 
 def chatbot(message, history):
     response = generate_response(message)
@@ -49,15 +52,14 @@ def chatbot(message, history):
 
 demo = gr.ChatInterface(
     chatbot,
-    title="Fine-tuned Phi-2 Chatbot",
-    description="This is a chatbot using a fine-tuned version of the Phi-2 model.",
-    theme="default",
+    title="Phi-2 Zero-GPU Chat",
+    description="Fine-tuned Phi-2 model running on Hugging Face Zero-GPU Spaces",
     examples=[
-        "Explain the concept of machine learning.",
-        "Write a short story about a robot learning to paint.",
-        "What are some effective ways to reduce stress?",
+        ["Explain quantum computing in simple terms"],
+        ["Write a poem about artificial intelligence"],
+        ["How do I make a perfect omelette?"]
     ],
-    cache_examples=True,
+    cache_examples=False
 )
 
 if __name__ == "__main__":