Add stopping criteria

handler.py

@@ -1,9 +1,32 @@
 from typing import Dict, List, Any, Union
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria
 import torch
 from peft import PeftModel
 
 
+class MyStoppingCriteria(StoppingCriteria):
+    def __init__(self, target_sequence, prompt, tokenizer):
+        self.target_sequence = target_sequence
+        self.prompt=prompt
+        self.tokenizer = tokenizer
+
+    def __call__(self, input_ids, scores, **kwargs):
+        # Get the generated text as a string
+        generated_text = self.tokenizer.decode(input_ids[0])
+        generated_text = generated_text.replace(self.prompt,'')
+        # Check if the target sequence appears in the generated text
+        if self.target_sequence in generated_text:
+            return True  # Stop generation
+
+        return False  # Continue generation
+
+    def __len__(self):
+        return 1
+
+    def __iter__(self):
+        yield self
+
+
 class EndpointHandler:
     def __init__(self, path=""):
         # load model and processor from path
@@ -69,9 +92,15 @@ class EndpointHandler:
 
         # pass inputs with all kwargs in data
         if parameters is not None:
-            outputs = self.model.generate(input_ids=input_ids, **parameters)
+            outputs = self.model.generate(
+                input_ids=input_ids,
+                stopping_criteria=MyStoppingCriteria("<|endoftext|>", inputs, self.tokenizer),
+                **parameters)
         else:
-            outputs = self.model.generate(input_ids=input_ids, max_new_tokens=20)
+            outputs = self.model.generate(
+                input_ids=input_ids, max_new_tokens=32,
+                stopping_criteria=MyStoppingCriteria("<|endoftext|>", inputs, self.tokenizer)
+            )
 
         # postprocess the prediction
         prediction = self.tokenizer.decode(outputs[0][input_ids.shape[1]:]) #, skip_special_tokens=True)