Update app.py

app.py

@@ -13,6 +13,18 @@ from transformers import (
     TextIteratorStreamer,
 )
 
+class StoppingCriteriaSub(StoppingCriteria):
+    def __init__(self, stops = [], encounters=1):
+        super().__init__()
+        self.stops = [stop.to("cuda") for stop in stops]
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
+        last_token = input_ids[0][-1]
+        for stop in self.stops:
+            if tokenizer.decode(stop) == tokenizer.decode(last_token):
+                return True
+        return False
+
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))