Update handler.py

handler.py

@@ -15,9 +15,13 @@ class EndpointHandler:
 
     def __call__(self, data: Dict[str, Any]) -> List[str]:
         prompt = data["inputs"]
+        if "config" in data:
+          config = data.pop("config", None)
+        else:
+          config = {'max_new_tokens':100}
         input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
-        generated_ids = self.model.generate(input_ids)
-        return [self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)]
+        generated_ids = self.model.generate(input_ids, **config)
+        return self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)