Update app.py

app.py

@@ -1,24 +1,34 @@
+import os
 import gradio as gr
-from transformers import pipeline
 import torch
+from peft import AutoPeftModelForCausalLM
+from transformers import AutoTokenizer, pipeline
 
 MODEL_ID = "SatyamSinghal/taskmind-1.1b-chat-lora"
 
-# Load once when Space starts
-pipe = pipeline(
-    "text-generation",
-    model=MODEL_ID,
-    device_map="auto",
+# Optional: use HF token from Space secrets if you add one
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID,
+    token=HF_TOKEN,
 )
 
-SYSTEM_PROMPT = (
-    "You are TaskMind, a helpful assistant for understanding team messages "
-    "and converting them into structured task information when appropriate."
+model = AutoPeftModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    token=HF_TOKEN,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    low_cpu_mem_usage=True,
 )
 
-def chat_fn(message, history):
-    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+)
 
+def respond(message, history):
+    messages = []
     for user_msg, assistant_msg in history:
         if user_msg:
             messages.append({"role": "user", "content": user_msg})
@@ -27,7 +37,7 @@ def chat_fn(message, history):
 
     messages.append({"role": "user", "content": message})
 
-    output = pipe(
+    out = pipe(
         messages,
         max_new_tokens=256,
         do_sample=True,
@@ -35,27 +45,23 @@ def chat_fn(message, history):
         top_p=0.9,
     )
 
-    # Pipeline output for chat models usually returns generated conversation
-    generated = output[0]["generated_text"]
-
+    generated = out[0]["generated_text"]
     if isinstance(generated, list):
         return generated[-1]["content"]
-
     return str(generated)
 
 demo = gr.ChatInterface(
-    fn=chat_fn,
-    title="TaskMind Chat Demo",
-    description="Try the TaskMind LoRA model hosted on Hugging Face Spaces.",
+    fn=respond,
+    title="TaskMind Demo",
+    description="Try the TaskMind LoRA model.",
     examples=[
         "Who are you?",
-        "@Agrim fix the growstreams deck ASAP no delay",
-        "login page 60% ho gaya",
+        "@Agrim fix the growstreams deck ASAP NO Delay",
         "done bhai, merged the PR",
+        "login page 60% ho gaya",
         "getting 500 error on registration",
     ],
-    theme="soft",
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()