update handler py in order to deal with batch inputs (#2)

- update handler py in order to deal with batch inputs (b8b060e69b5d03b9a2e58faf8b1f29460e68454a)


Co-authored-by: parpaillon <clemparpa@users.noreply.huggingface.co>

handler.py

@@ -32,34 +32,50 @@ class EndpointHandler():
         )
         FastLanguageModel.for_inference(self.model) # Enable native 2x faster inference
 
+        self.tokenizer.padding_side="left"
+        self.tokenizer.pad_token=tokenizer.eos_token
+          
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+    def _secure_inputs(self, data: Dict[str, Any]):
         inputs = data.get("inputs", None)
         if inputs is None:
-          return [{"err": "no inputs"}]
-
-        if not isinstance(inputs, str):
-          return [{"err": "inputs must bet of type {'inputs': str}"}]
-
-
-        tokenized_inputs = (
-              self.tokenizer(
-                self.summary_prompt.format(inputs, ""), 
-                return_tensors="pt"
-              )
-              .to("cuda")
-        )
-
-        outputs = self.model.generate(
-            **tokenized_inputs, 
-            max_new_tokens=self.max_new_tokens,
-            use_cache=True
-        )
-        outputs = outputs[:, tokenized_inputs["input_ids"].shape[1]:]
+            return [{"error": "inputs should be shaped like {'inputs': <string or List of strings (abstracts)>}"}], False
+
+        if isinstance(inputs, str):
+            inputs = [inputs]
+
+        return inputs, True
+
+    def _format_inputs(self, inputs: list[str]):
+        prompts = [summary_prompt.format(abstract, "") for abstract in abstracts]
+        prompts_lengths = [len(prompt) for prompt in prompts]
+        return prompts, prompts_lengths
+
+    def _generate_outputs(self, inputs):
+        tokenized = tokenizer(inputs, return_tensors="pt", padding=True).to("cuda")
+        outputs = model.generate(**tokenized, max_new_tokens=500, use_cache=True)
+        decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True) 
+        return decoded
+
+    def _format_outputs(self, outputs: list[str], inputs_lengths: list[int]):
+        decoded_without_input = [
+            output_str[input_len:].strip()
+            for output_str, input_len
+            in zip(outputs, inputs_lengths)
+        ]
+        return decoded_without_input
+        
+        
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        inputs, is_secure = self._secure_inputs(data)
+        
+        if not is_secure:
+            return inputs
         
-        outputs = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
-        outputs = outputs.strip(" ").strip("\n")
+        inputs, inputs_length = self._format_inputs(inputs)
+        outputs = self._generate_outputs(inputs)
+        outputs = self._format_outputs(outputs, inputs_length)
 
-        return [{
-            "result": outputs
-        }]
+        outputs = [{"summary": output_} for output_ in outputs]
+        
+        return outputs