Updated handler.py

handler.py

@@ -7,30 +7,32 @@ class EndpointHandler:
     def __init__(self, path=""):
         # load model and tokenizer from path
         self.tokenizer = AutoTokenizer.from_pretrained(path)
-        self.model = AutoModelForCausalLM.from_pretrained(
-            path, device_map="auto", torch_dtype=torch.float16, trust_remote_code=True
-        )
+        self.model = AutoModelForCausalLM.from_pretrained(path,
+                                                          device_map="auto",
+                                                          torch_dtype=torch.float16,
+                                                          trust_remote_code=True)
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
         # process input
         inputs = data.pop("inputs", data)
-        parameters = data.pop("parameters", None)
+        parameters = data.pop("parameters", {})
+        return_full_text = parameters.pop("return_full_text", True)
 
         # preprocess
-        inputs = self.tokenizer(inputs, return_tensors="pt")
-        inputs.pop("token_type_ids")
-        inputs.pop("attention_mask")
+        inputs = self.tokenizer(inputs,
+                                return_tensors="pt",
+                                return_token_type_ids=False)
         inputs = inputs.to(self.device)
+        input_len = len(inputs[0])
 
-        # pass inputs with all kwargs in data
-        if parameters is not None:
-            outputs = self.model.generate(**inputs, **parameters)
-        else:
-            outputs = self.model.generate(**inputs)
+        outputs = self.model.generate(**inputs, **parameters)[0]
+
+        if not return_full_text:
+            outputs = outputs[input_len:]
 
         # postprocess the prediction
-        prediction = self.tokenizer.decode(
-            outputs[0], skip_special_tokens=True)
+        prediction = self.tokenizer.decode(outputs,
+                                           skip_special_tokens=True)
 
         return [{"generated_text": prediction}]