using .generate, returning ids, custom bad_words

handler.py

@@ -2,17 +2,16 @@ from typing import Dict, List, Any
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, StoppingCriteria, StoppingCriteriaList
 
 
-
 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.model = AutoModelForCausalLM.from_pretrained(path)
+        self.tokenizer = tokenizer
         self.stopping_criteria = StoppingCriteriaList([StopAtPeriodCriteria(tokenizer)])
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+    def __call__(self, data: Dict[str, Any], additional_bad_words_ids: List[List[int]] = None) -> List[Dict[str, Any]]:
         """
        data args:
             inputs (:obj: `str`)
@@ -22,16 +21,29 @@ class EndpointHandler():
         """
         inputs = data.pop("inputs", data)
 
-        # Bad word: id 3070 corresponds to "(*", and we do not want to output a comment
-        prediction = self.pipeline(
-            inputs, 
-            stopping_criteria=self.stopping_criteria, 
-            max_new_tokens=50, 
-            return_full_text=False, 
-            bad_words_ids=[[3070], [313, 334], [10456]], 
+
+        # Bad word: id 3070, 10456 corresponds to "(*", and we do not want to output a comment
+        bad_words_ids = [[3070], [313, 334], [10456]]
+        if additional_bad_words_ids:
+            bad_words_ids.extend(additional_bad_words_ids)
+
+        input_ids = self.tokenizer.encode(inputs, return_tensors="pt")
+
+        # Generate text using model.generate
+        generated_ids = self.model.generate(
+            input_ids,
+            max_length=input_ids.shape[1] + 50,  # 50 new tokens
+            bad_words_ids=bad_words_ids,
             temperature=1,
             top_k=40,
+            stopping_criteria=self.stopping_criteria,
         )
+
+        # Slice the generated_ids to only include the new tokens generated, excluding the input tokens
+        generated_ids = generated_ids[:, input_ids.shape[1]:]
+
+        generated_text = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+        prediction = [{"generated_text": generated_text, "generated_ids": generated_ids[0].tolist()}]
         return prediction