Changing the logic in spaces_model_predict

app.py

@@ -41,26 +41,20 @@ class StopOnSequence(StoppingCriteria):
     def __call__(
         self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs
     ) -> bool:
-        if input_ids.shape[1] < self.sequence_len:
-            return False
-        return (
-            (
-                input_ids[0, -self.sequence_len :]
-                == torch.tensor(self.sequence_ids, device=input_ids.device)
-            )
-            .all()
-            .item()
-        )
+        for i in range(input_ids.shape[0]):
+            if input_ids[i, -self.sequence_len:].tolist() == self.sequence_ids:
+                return True
+        return False
 
 
 @spaces.GPU(duration=54)
 def spaces_model_predict(message: str, history: list[tuple[str, str]]):
     history_transformer_format = history + [[message, ""]]
-    stop = StopOnSequence("<|human|>", tokenizer)
+    stop = StopOnSequence("<|user|>", tokenizer)
 
     messages = "".join(
         [
-            "".join(["\n<human>:" + item[0], "\n<ai>:" + item[1]])
+            f"<|user|>\n{item[0]}\n<|assistant|>\n{item[1]}"
             for item in history_transformer_format
         ]
     )
@@ -85,9 +79,8 @@ def spaces_model_predict(message: str, history: list[tuple[str, str]]):
 
     partial_message = ""
     for new_token in streamer:
-        if new_token != "<":
-            partial_message += new_token
-            return partial_message
+        partial_message += new_token
+        yield partial_message
 
 
 def predict(