Update app.py

app.py

@@ -49,6 +49,16 @@ except Exception:
 
 from transformers import StoppingCriteria, StoppingCriteriaList
 
+class StopOnTokens(StoppingCriteria):
+    def __init__(self, stop_token_ids: List[int]):
+        self.stop_token_ids = stop_token_ids
+
+    def __call__(self, input_ids: "torch.LongTensor", scores: "torch.FloatTensor", **kwargs) -> bool:
+        for stop_id in self.stop_token_ids:
+            if input_ids[0][-1] == stop_id:
+                return True
+        return False
+
 try:
     import faiss
 except Exception: