Upload handler.py

handler.py

@@ -1,9 +1,9 @@
 from typing import Any, Dict
 
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# from peft import PeftConfig, PeftModel
+from peft import PeftConfig, PeftModel
 
 
 class EndpointHandler:
@@ -11,17 +11,17 @@ class EndpointHandler:
         # load model and processor from path
         self.tokenizer = AutoTokenizer.from_pretrained(path)
         # try:
-        # config = AutoConfig.from_pretrained(path)
+        config = PeftConfig.from_pretrained(path)
         model = AutoModelForCausalLM.from_pretrained(
-            path,
+            config.base_model_name_or_path,
             # return_dict=True,
             # load_in_8bit=True,
             device_map="auto",
             torch_dtype=torch.float16,
-            # trust_remote_code=True,
+            trust_remote_code=True,
         )
         # model.resize_token_embeddings(len(self.tokenizer))
-        # model = PeftModel.from_pretrained(model, path)
+        model = PeftModel.from_pretrained(model, path)
         # except Exception:
         #     model = AutoModelForCausalLM.from_pretrained(
         #         path, device_map="auto", load_in_8bit=True, torch_dtype=torch.float16, trust_remote_code=True
@@ -35,13 +35,13 @@ class EndpointHandler:
         parameters = data.pop("parameters", None)
 
         # preprocess
-        inputs = self.tokenizer(f"User: {inputs}\n\n", return_tensors="pt")
+        inputs = self.tokenizer(inputs, return_tensors="pt").to(self.device)
 
         # pass inputs with all kwargs in data
         if parameters is not None:
-            outputs = self.model.generate(**inputs.to(self.device), max_new_tokens=880, **parameters)
+            outputs = self.model.generate(**inputs, max_new_tokens=880, **parameters)
         else:
-            outputs = self.model.generate(**inputs.to(self.device), max_new_tokens=880)
+            outputs = self.model.generate(**inputs, max_new_tokens=880)
 
         # postprocess the prediction
         prediction = self.tokenizer.decode(outputs[0], skip_special_tokens=True)