Update handler.py

handler.py

@@ -20,21 +20,33 @@ class EndpointHandler():
 
         self.tokenizer = transformers.AutoTokenizer.from_pretrained('EleutherAI/gpt-neox-20b')
         print("tokenizer created  ", datetime.now())
+
+        
+    stop_token_ids = self.tokenizer.convert_tokens_to_ids(["<|endoftext|>"])
+
+    class StopOnTokens(StoppingCriteria):
+        def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs):
+            for stop_id in stop_token_ids:
+                if input_ids[0][-1] == stop_id:
+                    return True
+            return False
+
+    stopping_criteria = StoppingCriteriaList([StopOnTokens()])
+        
         self.generate_text = transformers.pipeline(
             model=self.model,
             tokenizer=self.tokenizer,
+            stopping_criteria=stopping_criteria,
             task='text-generation',
             return_full_text=True,
             temperature=0.1,
             top_p=0.15,
             top_k=0,
-            # max_new_tokens=64,
+            max_new_tokens=64,
             repetition_penalty=1.1
         )
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        print("iiiiiiiiii   ",  data)
-        inputs = data.pop("inputs  ", data)
-        print(inputs)
-        res = self.generate_text("Explain to me the difference between nuclear fission and fusion." , max_length= 60)
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:    
+        res = self.generate_text("Explain to me the difference between nuclear fission and fusion.")
         return res