Update handler.py

handler.py

@@ -2,22 +2,32 @@
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
+from huggingface_hub import snapshot_download
 
-# Model path in the repo
 BASE_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
-ADAPTER_PATH = "."
+ADAPTER_PATH = "GilbertAkham/deepseek-R1-multitask-lora"
 
 class EndpointHandler:
     def __init__(self, path=""):
-        print("Loading tokenizer and model...")
+        print("🚀 Loading base model...")
         self.tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
+
         base_model = AutoModelForCausalLM.from_pretrained(
-            BASE_MODEL, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True
+            BASE_MODEL,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            trust_remote_code=True
         )
-        self.model = PeftModel.from_pretrained(base_model, ADAPTER_PATH)
-        self.model = self.model.merge_and_unload()
+
+        print(f"🔗 Downloading LoRA adapter from {ADAPTER_PATH}...")
+        adapter_local_path = snapshot_download(repo_id=ADAPTER_PATH, allow_patterns=["*adapter*"])
+        print(f"📁 Adapter files cached at {adapter_local_path}")
+
+        print("🧩 Attaching LoRA adapter...")
+        self.model = PeftModel.from_pretrained(base_model, adapter_local_path)
         self.model.eval()
-        print("Model loaded successfully.")
+
+        print("✅ Model + LoRA adapter loaded successfully.")
 
     def __call__(self, data):
         prompt = data.get("inputs", "")