Update handler.py

handler.py

@@ -2,84 +2,13 @@ import torch
 from typing import  Dict, List, Any
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# MAX_TOKENS_IN_BATCH = 4_000
-# DEFAULT_MAX_NEW_TOKENS = 10
-
-
-# class EndpointHandler():
-#     def __init__(self, path: str = ""):
-
-#         self.tokenizer = AutoTokenizer.from_pretrained(path)
-#         self.model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch.float16)
-#         self.model = self.model.to('cuda:0')
-
-#     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-#         """
-#         Args:
-#             data (:obj:):
-#                 includes the input data and the parameters for the inference.
-#         Return:
-#             A :obj:`list`:. The list contains the answer and scores of the inference inputs
-#         """
-        
-#         prompts = [f"<human>: {prompt}\n<bot>:" for prompt in data["inputs"]]
-
-#         self.tokenizer.pad_token = self.tokenizer.eos_token
-#         inputs = self.tokenizer(prompts, padding=True, return_tensors='pt').to(self.model.device)
-#         input_length = inputs.input_ids.shape[1]
-
-#         outputs = self.model.generate(
-#             **inputs, **data["parameters"]
-#          )
-        
-#         output_strs = self.tokenizer.batch_decode(outputs[:, input_length:], skip_special_tokens=True)
-        
-#         return [{"generated_text": output_strs}]
-
-# import torch
-# from typing import Dict, List, Any
-# from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-
-# # get dtype
-# dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
-
-
-# class EndpointHandler:
-#     def __init__(self, path=""):
-#         # load the model
-#         tokenizer = AutoTokenizer.from_pretrained(path)
-#         model = AutoModelForCausalLM.from_pretrained(path, device_map="auto", torch_dtype=torch.float16)
-
-#     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
-#         """
-#         Args:
-#             data (:obj:):
-#                 includes the input data and the parameters for the inference.
-#         Return:
-#             A :obj:`list`:. The list contains the answer and scores of the inference inputs
-#         """
-        
-#         inputs = data.pop("inputs", data)
-#         parameters = data.pop("parameters", None)
-
-#         prompts = [f"<human>: {prompt}\n<bot>:" for prompt in inputs]
-
-#         # pass inputs with all kwargs in data
-#         if parameters is not None:
-#             prediction = self.pipeline(inputs, **parameters)
-#         else:
-#             prediction = self.pipeline(inputs)
-#         # postprocess the prediction
-#         return prediction
-
-
-class EndpointHandler():
+class EndpointHandler:
     def __init__(self, path: str = ""):
 
-        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.tokenizer = AutoTokenizer.from_pretrained(path, padding_side = "left")
         self.model = AutoModelForCausalLM.from_pretrained(path, device_map = "auto", torch_dtype=torch.float16)
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         """
         Args:
             data (:obj:):
@@ -87,13 +16,22 @@ class EndpointHandler():
         Return:
             A :obj:`list`:. The list contains the answer and scores of the inference inputs
         """
+
+        # process input
+        inputs_dict = data.pop("inputs", data)
+        parameters = data.pop("parameters", {})
         
-        prompts = [f"<human>: {prompt}\n<bot>:" for prompt in data["inputs"]]
+        prompts = [f"<human>: {prompt}\n<bot>:" for prompt in inputs_dict]
 
         self.tokenizer.pad_token = self.tokenizer.eos_token
-        inputs = self.tokenizer(prompts, padding=True, return_tensors='pt').to(self.model.device)
+
+        inputs = self.tokenizer(prompts, truncation=True, max_length=2048-512,
+                                return_tensors='pt', padding=True).to(self.model.device)
         input_length = inputs.input_ids.shape[1]
 
+        if parameters.get("deterministic", False):
+            torch.manual_seed(42)
+
         outputs = self.model.generate(
             **inputs, max_new_tokens=512, do_sample=True, temperature=0.7, top_p=0.7, top_k=50
          )