Create utils.py

utils.py

@@ -0,0 +1,14 @@
+from transformers import StoppingCriteria
+
+
+# Define a stopping condition for generation
+class SpecificStringStoppingCriteria(StoppingCriteria):
+    def __init__(self, tokenizer, stop_strings, input_len):
+        self.tokenizer = tokenizer
+        self.stop_strings = stop_strings
+        self.input_len = input_len
+
+    def __call__(self, input_ids, scores, **kwargs):
+        current_text = self.tokenizer.decode(input_ids[0], skip_special_tokens=True)[self.input_len:]
+        
+        return any(stop_string in current_text for stop_string in self.stop_strings)