adding additional bad words ids

handler.py

@@ -20,6 +20,7 @@ class EndpointHandler():
             A :obj:`list` | `dict`: will be serialized and returned
         """
         inputs = data.pop("inputs", data)
+        additional_bad_words_ids = data.pop("bad_words_ids", [])
 
 
         # 3070, 10456, [313, 334] corresponds to "(*", and we do not want to output a comment
@@ -27,6 +28,7 @@ class EndpointHandler():
         # [1976, 441, 29889] is "Abort."
         # [2087, 29885, 4430, 29889] is "Admitted."
         bad_words_ids = [[3070], [313, 334], [10456], [13], [1976, 441, 29889], [2087, 29885, 4430, 29889]]
+        bad_words_ids.extend(additional_bad_words_ids)
 
         input_ids = self.tokenizer.encode(inputs, return_tensors="pt")