Update app.py

app.py

@@ -11,7 +11,7 @@ from peft import PeftModel
 
 
 # ============================================================
-# MODEL
+# CONFIGURATION
 # ============================================================
 
 BASE_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
@@ -20,7 +20,7 @@ ADAPTER_MODEL = "dd253B/DhanushAI-0.5B"
 
 
 # ============================================================
-# GLOBAL VARIABLES
+# GLOBAL MODEL
 # ============================================================
 
 tokenizer = None
@@ -39,6 +39,14 @@ def load_model():
     if model is not None:
         return
 
+    print("====================================")
+    print("Loading DhanushAI...")
+    print("====================================")
+
+    # -----------------------------
+    # Tokenizer
+    # -----------------------------
+
     print("Loading tokenizer...")
 
     tokenizer = AutoTokenizer.from_pretrained(
@@ -48,6 +56,11 @@ def load_model():
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
+
+    # -----------------------------
+    # Base model
+    # -----------------------------
+
     print("Loading base model...")
 
     base_model = AutoModelForCausalLM.from_pretrained(
@@ -55,6 +68,11 @@ def load_model():
         torch_dtype=torch.float16
     )
 
+
+    # -----------------------------
+    # LoRA adapter
+    # -----------------------------
+
     print("Loading DhanushAI adapter...")
 
     model = PeftModel.from_pretrained(
@@ -62,44 +80,74 @@ def load_model():
         ADAPTER_MODEL
     )
 
+
+    # -----------------------------
+    # Move to GPU
+    # -----------------------------
+
     model = model.to("cuda")
 
     model.eval()
 
+    print("====================================")
     print("DhanushAI loaded successfully!")
+    print("====================================")
 
 
 # ============================================================
-# CHAT
+# CHAT FUNCTION
 # ============================================================
 
 @spaces.GPU
 def chat(message):
 
+    # Load model after ZeroGPU allocation
     load_model()
 
-    if not message or not message.strip():
+    if message is None:
+        return "Please enter a message."
+
+    message = message.strip()
+
+    if not message:
         return "Please enter a message."
 
+
+    # -----------------------------
+    # Prompt
+    # -----------------------------
+
     prompt = f"""You are DhanushAI, a helpful AI assistant.
 
 User: {message}
 
 Assistant:"""
 
+
+    # -----------------------------
+    # Tokenize
+    # -----------------------------
+
     inputs = tokenizer(
         prompt,
         return_tensors="pt"
     )
 
+
     inputs = {
         key: value.to("cuda")
         for key, value in inputs.items()
     }
 
+
+    # -----------------------------
+    # Generate
+    # -----------------------------
+
     with torch.no_grad():
 
         outputs = model.generate(
+
             **inputs,
 
             max_new_tokens=200,
@@ -113,11 +161,21 @@ Assistant:"""
             repetition_penalty=1.1
         )
 
+
+    # -----------------------------
+    # Decode
+    # -----------------------------
+
     generated = tokenizer.decode(
         outputs[0],
         skip_special_tokens=True
     )
 
+
+    # -----------------------------
+    # Remove prompt
+    # -----------------------------
+
     if "Assistant:" in generated:
 
         answer = generated.split(
@@ -129,11 +187,12 @@ Assistant:"""
 
         answer = generated.strip()
 
+
     return answer
 
 
 # ============================================================
-# GRADIO
+# GRADIO UI + API
 # ============================================================
 
 demo = gr.Interface(
@@ -142,7 +201,7 @@ demo = gr.Interface(
 
     inputs=gr.Textbox(
         label="Message",
-        placeholder="Ask DhanushAI..."
+        placeholder="Ask DhanushAI something..."
     ),
 
     outputs=gr.Textbox(
@@ -151,8 +210,14 @@ demo = gr.Interface(
 
     title="DhanushAI",
 
-    description="My custom AI model"
+    description="My custom AI model",
+
+    api_name="chat"
 )
 
 
+# ============================================================
+# START
+# ============================================================
+
 demo.launch()