Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import torch
+import spaces
 
 from transformers import (
     AutoTokenizer,
@@ -10,7 +11,7 @@ from peft import PeftModel
 
 
 # ==========================================
-# MODEL CONFIG
+# MODEL
 # ==========================================
 
 BASE_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
@@ -19,22 +20,7 @@ ADAPTER_MODEL = "dd253B/DhanushAI-0.5B"
 
 
 # ==========================================
-# DEVICE
-# ==========================================
-
-if torch.cuda.is_available():
-    DEVICE = "cuda"
-    DTYPE = torch.float16
-else:
-    DEVICE = "cpu"
-    DTYPE = torch.float32
-
-
-print("Device:", DEVICE)
-
-
-# ==========================================
-# TOKENIZER
+# LOAD TOKENIZER
 # ==========================================
 
 print("Loading tokenizer...")
@@ -48,41 +34,34 @@ if tokenizer.pad_token is None:
 
 
 # ==========================================
-# BASE MODEL
+# LOAD MODEL
 # ==========================================
 
 print("Loading base model...")
 
-base_model = AutoModelForCausalLM.from_pretrained(
+model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
-    torch_dtype=DTYPE
+    torch_dtype=torch.float16,
+    device_map="auto"
 )
 
-base_model.to(DEVICE)
-
-
-# ==========================================
-# YOUR LORA MODEL
-# ==========================================
-
-print("Loading DhanushAI...")
+print("Loading DhanushAI adapter...")
 
 model = PeftModel.from_pretrained(
-    base_model,
+    model,
     ADAPTER_MODEL
 )
 
-model.to(DEVICE)
-
 model.eval()
 
 print("DhanushAI loaded!")
 
 
 # ==========================================
-# CHAT FUNCTION
+# GPU FUNCTION
 # ==========================================
 
+@spaces.GPU
 def chat(message):
 
     if not message or not message.strip():
@@ -99,8 +78,9 @@ Assistant:"""
         return_tensors="pt"
     )
 
+    # Move inputs to the same device as model
     inputs = {
-        key: value.to(DEVICE)
+        key: value.to(model.device)
         for key, value in inputs.items()
     }
 
@@ -140,16 +120,15 @@ Assistant:"""
 
 
 # ==========================================
-# GRADIO API
+# GRADIO
 # ==========================================
 
 demo = gr.Interface(
-
     fn=chat,
 
     inputs=gr.Textbox(
         label="Message",
-        placeholder="Ask DhanushAI something..."
+        placeholder="Ask DhanushAI..."
     ),
 
     outputs=gr.Textbox(