Remove custom stoppingcriteria and trust generate

app.py

@@ -6,15 +6,9 @@ import spaces
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 import threading
 
-from transformers import TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList
+from transformers import TextIteratorStreamer
 import threading
 
-class StopOnEos(StoppingCriteria):
-    def __init__(self, eos_token_id):
-        self.eos_token_id = eos_token_id
-
-    def __call__(self, input_ids, scores, **kwargs):
-        return input_ids[0, -1].item() == self.eos_token_id
 
 @spaces.GPU
 def chat_with_model(messages):
@@ -33,14 +27,12 @@ def chat_with_model(messages):
     inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
 
     streamer = TextIteratorStreamer(current_tokenizer, skip_prompt=True, skip_special_tokens=False)
-    stopping_criteria = StoppingCriteriaList([StopOnEos(current_tokenizer.eos_token_id)])
 
     generation_kwargs = dict(
         **inputs,
         max_new_tokens=256,
         do_sample=True,
         streamer=streamer,
-        stopping_criteria=stopping_criteria,
         eos_token_id=current_tokenizer.eos_token_id,
         pad_token_id=pad_id
     )