Update app_chat.py

app_chat.py

@@ -6,6 +6,7 @@ import gradio as gr
 import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+# from transformers import StoppingCriteria, StoppingCriteriaList, StopStringCriteria
 
 MAX_MAX_NEW_TOKENS = 1024
 DEFAULT_MAX_NEW_TOKENS = 512
@@ -21,6 +22,19 @@ model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat1
 tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
 #tokenizer.use_default_system_prompt = False
 
+# class StoppingCriteriaSub(StoppingCriteria):
+#     def __init__(self, tokenizer, stops = [], encounters=1):
+#         super().__init__()
+#         self.stops = [stop.to("cuda") for stop in stops]
+#         self.tokenizer = tokenizer
+#         self.num_mamba_stop_ids = 8
+
+#     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
+#         last_token = input_ids[0][-self.num_mamba_stop_ids:]
+#         for stop in self.stops:
+#             if self.tokenizer.decode(stop) in self.tokenizer.decode(last_token):
+#                 return True
+#         return False
 
 @spaces.GPU
 def generate(
@@ -39,7 +53,10 @@ def generate(
     conversation += chat_history
     conversation.append({"role": "User", "content": message})
 
-    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
+    input_ids = tokenizer.apply_chat_template(conversation,  tokenize=True, add_generation_prompt=True, return_tensors="pt")
+    
+    # stopping_criteria = StoppingCriteriaList([StopStringCriteria(tokenizer=tokenizer, stop_strings="</s>")])
+    
     if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
         input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
         gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
@@ -56,6 +73,7 @@ def generate(
         temperature=temperature,
         num_beams=1,
         repetition_penalty=repetition_penalty,
+        # "stopping_criteria": stopping_criteria,
     )
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()