Update handler.py

handler.py

@@ -1,21 +1,16 @@
 # handler.py
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from peft import PeftModel
 
-# Model path in the repo
-BASE_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
-ADAPTER_PATH = "."
+MODEL_PATH = "."
 
 class EndpointHandler:
     def __init__(self, path=""):
-        print("Loading tokenizer and 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
+        print("Loading merged model...")
+        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            MODEL_PATH, 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()
         self.model.eval()
         print("Model loaded successfully.")
 
@@ -32,5 +27,4 @@ class EndpointHandler:
                 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}
+        return {"generated_text": self.tokenizer.decode(outputs[0], skip_special_tokens=True)}