Update app_chat.py

app_chat.py

@@ -6,7 +6,7 @@ import gradio as gr
 import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-# from transformers import StoppingCriteria, StoppingCriteriaList, StopStringCriteria
+from transformers import StoppingCriteria, StoppingCriteriaList, StopStringCriteria
 
 import subprocess
 # useradd -m -u 1000 user
@@ -33,19 +33,19 @@ tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
 tokenizer.chat_template = "{{'<extra_id_0>System'}}{% for message in messages %}{% if message['role'] == 'system' %}{{'\n' + message['content'].strip()}}{% if tools or contexts %}{{'\n'}}{% endif %}{% endif %}{% endfor %}{% if tools %}{% for tool in tools %}{{ '\n<tool> ' + tool|tojson + ' </tool>' }}{% endfor %}{% endif %}{% if contexts %}{% if tools %}{{'\n'}}{% endif %}{% for context in contexts %}{{ '\n<context> ' + context.strip() + ' </context>' }}{% endfor %}{% endif %}{{'\n\n'}}{% for message in messages %}{% if message['role'] == 'user' %}{{ '<extra_id_1>User\n' + message['content'].strip() + '\n' }}{% elif message['role'] == 'assistant' %}{{ '<extra_id_1>Assistant\n' + message['content'].strip() + '\n' }}{% elif message['role'] == 'tool' %}{{ '<extra_id_1>Tool\n' + message['content'].strip() + '\n' }}{% endif %}{% endfor %}{%- if add_generation_prompt %}{{'<extra_id_1>Assistant\n'}}{%- endif %}"
 #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
+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
+    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(
@@ -66,7 +66,7 @@ def generate(
 
     input_ids = tokenizer.apply_chat_template(conversation,  tokenize=True, add_generation_prompt=True, return_tensors="pt").to('cuda')
     
-    # stopping_criteria = StoppingCriteriaList([StopStringCriteria(tokenizer=tokenizer, stop_strings="</s>")])
+    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:]
@@ -84,7 +84,7 @@ def generate(
         temperature=temperature,
         num_beams=1,
         repetition_penalty=repetition_penalty,
-        # "stopping_criteria": stopping_criteria,
+        "stopping_criteria": stopping_criteria,
     )
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
@@ -98,11 +98,11 @@ def generate(
 chat_interface = gr.ChatInterface(
     fn=generate,
     additional_inputs=[
-        gr.Textbox(label="System prompt", lines=6, value="You are a helpful assistant. Your name is Hymba-1.5B-Instruct-8K. \
-                                               You are a new family of small language models featuring a hybrid-head architecture that strategically integrates attention mechanisms with state space models (SSMs). \
-                                               You are developed by Deep Learning Efficiency Research (DLER) team at NVIDIA Research. \
-                                               Nvidia Corporation is an American multinational corporation and technology company headquartered in Santa Clara, California. Nvidia was founded on April 5, 1993 by Jensen Huang. \
-                                               The above is just a background context. You can answer any questions not limited to the above background context."),
+        # gr.Textbox(label="System prompt", lines=6, value="You are a helpful assistant. Your name is Hymba-1.5B-Instruct-8K. \
+        #                                        You are a new family of small language models featuring a hybrid-head architecture that strategically integrates attention mechanisms with state space models (SSMs). \
+        #                                        You are developed by Deep Learning Efficiency Research (DLER) team at NVIDIA Research. \
+        #                                        The above is just a background context. You can answer any questions not limited to the above background context."),
+        gr.Textbox(label="System prompt", lines=6, value="You are a helpful assistant. Your name is Hymba-1.5B-Instruct-8K. "),
         gr.Slider(
             label="Max new tokens",
             minimum=1,