Update handler.py

handler.py

@@ -1,12 +1,18 @@
-# handler.py
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 from huggingface_hub import snapshot_download
 
+# === Base & adapter config ===
 BASE_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
 ADAPTER_PATH = "GilbertAkham/deepseek-R1-multitask-lora"
 
+# === System message (always prepended to input) ===
+SYSTEM_PROMPT = (
+    "You are Chat-Bot, a helpful and logical assistant trained for reasoning, "
+    "email, chatting, summarization, story continuation, and report writing.\n\n"
+)
+
 class EndpointHandler:
     def __init__(self, path=""):
         print("🚀 Loading base model...")
@@ -30,17 +36,27 @@ class EndpointHandler:
         print("✅ Model + LoRA adapter loaded successfully.")
 
     def __call__(self, data):
+        # === Combine system + user prompt ===
         prompt = data.get("inputs", "")
-        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+        full_prompt = SYSTEM_PROMPT + prompt
+
+        params = data.get("parameters", {})
+        max_new_tokens = params.get("max_new_tokens", 512)
+        temperature = params.get("temperature", 0.7)
+        top_p = params.get("top_p", 0.9)
+
+        # === Tokenize and run generation ===
+        inputs = self.tokenizer(full_prompt, return_tensors="pt").to(self.model.device)
         with torch.no_grad():
             outputs = self.model.generate(
                 **inputs,
-                max_new_tokens=512,
-                temperature=0.7,
-                top_p=0.9,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_p=top_p,
                 do_sample=True,
                 pad_token_id=self.tokenizer.eos_token_id,
                 eos_token_id=self.tokenizer.eos_token_id,
             )
+
         text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
         return {"generated_text": text}