testing custom logit processor for bad words

handler.py

@@ -1,5 +1,5 @@
 from typing import Dict, List, Any
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, StoppingCriteria, StoppingCriteriaList
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, StoppingCriteria, StoppingCriteriaList, LogitsProcessor, LogitsProcessorList
 
 
 
@@ -11,6 +11,7 @@ class EndpointHandler():
         tokenizer.pad_token = tokenizer.eos_token
         self.pipeline = pipeline('text-generation', model=model, tokenizer=tokenizer)
         self.stopping_criteria = StoppingCriteriaList([StopAtPeriodCriteria(tokenizer)])
+        self.logits_processor = LogitsProcessorList([BanSpecificTokensLogitsProcessor(tokenizer, [3070])])
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
@@ -28,8 +29,9 @@ class EndpointHandler():
             stopping_criteria=self.stopping_criteria, 
             max_new_tokens=50, 
             return_full_text=False, 
-            bad_words_ids=[[3070], [313, 334]], 
-            temperature=2,
+            # bad_words_ids=[[3070], [313, 334]], 
+            logits_processor=self.logits_processor,
+            temperature=1,
             top_k=40,
         )
         return prediction
@@ -43,4 +45,19 @@ class StopAtPeriodCriteria(StoppingCriteria):
         # Decode the last generated token to text
         last_token_text = self.tokenizer.decode(input_ids[:, -1], skip_special_tokens=True)
         # Check if the decoded text ends with a period
-        return '.' in last_token_text
+        return '.' in last_token_text
+    
+class BanSpecificTokensLogitsProcessor(LogitsProcessor):
+    """
+    Logits processor that sets the logits of specific tokens to -inf,
+    effectively banning them from being generated.
+    """
+    def __init__(self, tokenizer, banned_tokens_ids):
+        self.tokenizer = tokenizer
+        self.banned_tokens_ids = banned_tokens_ids
+
+    def __call__(self, input_ids, scores):
+        # Set logits of banned tokens to -inf
+        for token_id in self.banned_tokens_ids:
+            scores[:, token_id] = float('-inf')
+        return scores

sample.py

@@ -4,7 +4,7 @@ from handler import EndpointHandler
 my_handler = EndpointHandler(path=".")
 
 # prepare sample payload
-payload = {"inputs": "I am so happy to reach in my pocket and find a"}
+payload = {"inputs": "This is the format for a"}
 
 # test the handler
 payload=my_handler(payload)