testing .generate instead of pipeline

handler.py

@@ -1,5 +1,5 @@
 from typing import Dict, List, Any
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, StoppingCriteria, StoppingCriteriaList, LogitsProcessor, LogitsProcessorList
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, StoppingCriteria, StoppingCriteriaList
 
 
 
@@ -7,11 +7,10 @@ class EndpointHandler():
     def __init__(self, path=""):
         # Preload all the elements you are going to need at inference.
         tokenizer = AutoTokenizer.from_pretrained(path)
-        model = AutoModelForCausalLM.from_pretrained(path)
-        tokenizer.pad_token = tokenizer.eos_token
-        self.pipeline = pipeline('text-generation', model=model, tokenizer=tokenizer)
+        self.tokenizer = tokenizer
+        self.model = AutoModelForCausalLM.from_pretrained(path)
+        self.tokenizer.pad_token = tokenizer.eos_token
         self.stopping_criteria = StoppingCriteriaList([StopAtPeriodCriteria(tokenizer)])
-        self.logits_processor = LogitsProcessorList([BanSpecificTokensLogitsProcessor(tokenizer, [3070])])
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
@@ -22,19 +21,24 @@ class EndpointHandler():
             A :obj:`list` | `dict`: will be serialized and returned
         """
         inputs = data.pop("inputs", data)
+        input_ids = self.tokenizer.encode(inputs, return_tensors="pt")
 
         # Bad word: id 3070 corresponds to "(*", and we do not want to output a comment
-        prediction = self.pipeline(
-            inputs, 
+        prediction_ids = self.model.generate(
+            input_ids, 
+            max_length=input_ids.shape[1] + 50, 
             stopping_criteria=self.stopping_criteria, 
-            max_new_tokens=50, 
-            return_full_text=False, 
-            # bad_words_ids=[[3070], [313, 334]], 
-            logits_processor=self.logits_processor,
+            bad_words_ids=[[3070], [313, 334]], 
             temperature=1,
             top_k=40,
+            # pad_token_id=self.tokenizer.eos_token_id,
+            # return_dict_in_generate=True,  # To get more detailed output (optional)
         )
-        return prediction
+
+        # Decode the generated ids to text
+        # Exclude the input_ids length to get only the new tokens
+        prediction_text = self.tokenizer.decode(prediction_ids[0, input_ids.shape[1]:], skip_special_tokens=True)
+        return [{"generated_text": prediction_text}]
 
 
 class StopAtPeriodCriteria(StoppingCriteria):
@@ -45,19 +49,4 @@ 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
-    
-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
+        return '.' in last_token_text

sample.py

@@ -4,7 +4,7 @@ from handler import EndpointHandler
 my_handler = EndpointHandler(path=".")
 
 # prepare sample payload
-payload = {"inputs": "This is the format for a"}
+payload = {"inputs": "I want to turn the next page of the"}
 
 # test the handler
 payload=my_handler(payload)