Streaming output

app.py

@@ -1,3 +1,4 @@
+'''
 import os
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
@@ -78,43 +79,6 @@ def generate_response(chat_history, max_new_tokens, model, tokenizer):
 with gr.Blocks() as demo:
     gr.Markdown("# Zamba2 Model Selector")
     with gr.Tabs():
-        with gr.TabItem("2.7B Instruct Model"):
-            gr.Markdown("### Zamba2-2.7B Instruct Model")
-            with gr.Column():
-                chat_history_2_7B_instruct = gr.State([])  
-                chatbot_2_7B_instruct = gr.Chatbot()
-                message_2_7B_instruct = gr.Textbox(lines=2, placeholder="Enter your message...", label="Your Message")
-            with gr.Accordion("Generation Parameters", open=False):
-                max_new_tokens_2_7B_instruct = gr.Slider(50, 1000, step=50, value=500, label="Max New Tokens")
-                # temperature_2_7B_instruct = gr.Slider(0.1, 1.5, step=0.1, value=0.2, label="Temperature")
-                # top_k_2_7B_instruct = gr.Slider(1, 100, step=1, value=50, label="Top K")
-                # top_p_2_7B_instruct = gr.Slider(0.1, 1.0, step=0.1, value=1.0, label="Top P")
-                # repetition_penalty_2_7B_instruct = gr.Slider(1.0, 2.0, step=0.1, value=1.2, label="Repetition Penalty")
-                # num_beams_2_7B_instruct = gr.Slider(1, 10, step=1, value=1, label="Number of Beams")
-                # length_penalty_2_7B_instruct = gr.Slider(0.0, 2.0, step=0.1, value=1.0, label="Length Penalty")
-
-            def user_message_2_7B_instruct(message, chat_history):
-                chat_history = chat_history + [[message, None]]
-                return gr.update(value=""), chat_history, chat_history
-
-            def bot_response_2_7B_instruct(chat_history, max_new_tokens):
-                response = generate_response(chat_history, max_new_tokens, model_2_7B_instruct, tokenizer_2_7B_instruct)
-                chat_history[-1][1] = response
-                return chat_history, chat_history
-
-            send_button_2_7B_instruct = gr.Button("Send")
-            send_button_2_7B_instruct.click(
-                fn=user_message_2_7B_instruct,
-                inputs=[message_2_7B_instruct, chat_history_2_7B_instruct],
-                outputs=[message_2_7B_instruct, chat_history_2_7B_instruct, chatbot_2_7B_instruct]
-            ).then(
-                fn=bot_response_2_7B_instruct,
-                inputs=[
-                    chat_history_2_7B_instruct,
-                    max_new_tokens_2_7B_instruct
-                ],
-                outputs=[chat_history_2_7B_instruct, chatbot_2_7B_instruct]
-            )
         with gr.TabItem("7B Instruct Model"):
             gr.Markdown("### Zamba2-7B Instruct Model")
             with gr.Column():
@@ -152,19 +116,58 @@ with gr.Blocks() as demo:
                 ],
                 outputs=[chat_history_7B_instruct, chatbot_7B_instruct]
             )
+        with gr.TabItem("2.7B Instruct Model"):
+            gr.Markdown("### Zamba2-2.7B Instruct Model")
+            with gr.Column():
+                chat_history_2_7B_instruct = gr.State([])  
+                chatbot_2_7B_instruct = gr.Chatbot()
+                message_2_7B_instruct = gr.Textbox(lines=2, placeholder="Enter your message...", label="Your Message")
+            with gr.Accordion("Generation Parameters", open=False):
+                max_new_tokens_2_7B_instruct = gr.Slider(50, 1000, step=50, value=500, label="Max New Tokens")
+                # temperature_2_7B_instruct = gr.Slider(0.1, 1.5, step=0.1, value=0.2, label="Temperature")
+                # top_k_2_7B_instruct = gr.Slider(1, 100, step=1, value=50, label="Top K")
+                # top_p_2_7B_instruct = gr.Slider(0.1, 1.0, step=0.1, value=1.0, label="Top P")
+                # repetition_penalty_2_7B_instruct = gr.Slider(1.0, 2.0, step=0.1, value=1.2, label="Repetition Penalty")
+                # num_beams_2_7B_instruct = gr.Slider(1, 10, step=1, value=1, label="Number of Beams")
+                # length_penalty_2_7B_instruct = gr.Slider(0.0, 2.0, step=0.1, value=1.0, label="Length Penalty")
+
+            def user_message_2_7B_instruct(message, chat_history):
+                chat_history = chat_history + [[message, None]]
+                return gr.update(value=""), chat_history, chat_history
+
+            def bot_response_2_7B_instruct(chat_history, max_new_tokens):
+                response = generate_response(chat_history, max_new_tokens, model_2_7B_instruct, tokenizer_2_7B_instruct)
+                chat_history[-1][1] = response
+                return chat_history, chat_history
+
+            send_button_2_7B_instruct = gr.Button("Send")
+            send_button_2_7B_instruct.click(
+                fn=user_message_2_7B_instruct,
+                inputs=[message_2_7B_instruct, chat_history_2_7B_instruct],
+                outputs=[message_2_7B_instruct, chat_history_2_7B_instruct, chatbot_2_7B_instruct]
+            ).then(
+                fn=bot_response_2_7B_instruct,
+                inputs=[
+                    chat_history_2_7B_instruct,
+                    max_new_tokens_2_7B_instruct
+                ],
+                outputs=[chat_history_2_7B_instruct, chatbot_2_7B_instruct]
+            )
 
 if __name__ == "__main__":
     demo.queue().launch(max_threads=1)
-
 '''
 
 import os
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 import torch
+import threading
+import re
 
 model_name_2_7B_instruct = "Zyphra/Zamba2-2.7B-instruct"
 model_name_7B_instruct = "Zyphra/Zamba2-7B-instruct"
+max_context_length = 4096
 
 tokenizer_2_7B_instruct = AutoTokenizer.from_pretrained(model_name_2_7B_instruct)
 model_2_7B_instruct = AutoModelForCausalLM.from_pretrained(
@@ -176,131 +179,47 @@ model_7B_instruct = AutoModelForCausalLM.from_pretrained(
     model_name_7B_instruct, device_map="cuda", torch_dtype=torch.bfloat16
 )
 
-def extract_assistant_response(generated_text):
-    assistant_token = '<|im_start|> assistant'
-    end_token = '<|im_end|>'
-    start_idx = generated_text.rfind(assistant_token)
-    if start_idx == -1:
-        # Assistant token not found
-        return generated_text.strip()
-    start_idx += len(assistant_token)
-    end_idx = generated_text.find(end_token, start_idx)
-    if end_idx == -1:
-        # End token not found, return from start_idx to end
-        return generated_text[start_idx:].strip()
-    else:
-        return generated_text[start_idx:end_idx].strip()
-
-def generate_response_2_7B_instruct(chat_history, max_new_tokens):
+def generate_response(chat_history, max_new_tokens, model, tokenizer):
     sample = []
     for turn in chat_history:
         if turn[0]:
             sample.append({'role': 'user', 'content': turn[0]})
         if turn[1]:
             sample.append({'role': 'assistant', 'content': turn[1]})
-    chat_sample = tokenizer_2_7B_instruct.apply_chat_template(sample, tokenize=False)
-    input_ids = tokenizer_2_7B_instruct(chat_sample, return_tensors='pt', add_special_tokens=False).input_ids.to(model_2_7B_instruct.device)
+    chat_sample = tokenizer.apply_chat_template(sample, tokenize=False)
+    input_ids = tokenizer(chat_sample, return_tensors='pt', add_special_tokens=False).to(model.device)
 
-    # Handle context length limit
-    max_context_length = 4096
     max_new_tokens = int(max_new_tokens)
     max_input_length = max_context_length - max_new_tokens
-    if input_ids.size(1) > max_input_length:
-        input_ids = input_ids[:, -max_input_length:]  # Truncate from the left (oldest tokens)
+    if input_ids['input_ids'].size(1) > max_input_length:
+        input_ids['input_ids'] = input_ids['input_ids'][:, -max_input_length:]
+        if 'attention_mask' in input_ids:
+            input_ids['attention_mask'] = input_ids['attention_mask'][:, -max_input_length:]
 
-    with torch.no_grad():
-        outputs = model_2_7B_instruct.generate(
-            input_ids=input_ids,
-            max_new_tokens=max_new_tokens,
-            return_dict_in_generate=False,
-            output_scores=False,
-            use_cache=True,
-            num_beams=1,
-            do_sample=False
-        )
-
-    generated_text = tokenizer_2_7B_instruct.decode(outputs[0])
-    assistant_response = extract_assistant_response(generated_text)
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    generation_kwargs = dict(**input_ids, max_new_tokens=int(max_new_tokens), streamer=streamer)
 
-    del input_ids
-    del outputs
-    torch.cuda.empty_cache()
-    return assistant_response
+    thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
 
-def generate_response_7B_instruct(chat_history, max_new_tokens):
-    sample = []
-    for turn in chat_history:
-        if turn[0]:
-            sample.append({'role': 'user', 'content': turn[0]})
-        if turn[1]:
-            sample.append({'role': 'assistant', 'content': turn[1]})
-    chat_sample = tokenizer_7B_instruct.apply_chat_template(sample, tokenize=False)
-    input_ids = tokenizer_7B_instruct(chat_sample, return_tensors='pt', add_special_tokens=False).input_ids.to(model_7B_instruct.device)
+    assistant_response = ""
 
-    # Handle context length limit
-    max_context_length = 4096
-    max_new_tokens = int(max_new_tokens)
-    max_input_length = max_context_length - max_new_tokens
-    if input_ids.size(1) > max_input_length:
-        input_ids = input_ids[:, -max_input_length:]  # Truncate from the left (oldest tokens)
-
-    with torch.no_grad():
-        outputs = model_7B_instruct.generate(
-            input_ids=input_ids,
-            max_new_tokens=max_new_tokens,
-            return_dict_in_generate=False,
-            output_scores=False,
-            use_cache=True,
-            num_beams=1,
-            do_sample=False
-        )
-
-    generated_text = tokenizer_7B_instruct.decode(outputs[0])
-    assistant_response = extract_assistant_response(generated_text)
+    for new_text in streamer:
+        new_text = re.sub(r'^\s*(?i:assistant)[:\s]*', '', new_text)
+        assistant_response += new_text
+        yield assistant_response
 
+    thread.join()
     del input_ids
-    del outputs
     torch.cuda.empty_cache()
-    return assistant_response
 
 with gr.Blocks() as demo:
     gr.Markdown("# Zamba2 Model Selector")
     with gr.Tabs():
-        with gr.TabItem("2.7B Instruct Model"):
-            gr.Markdown("### Zamba2-2.7B Instruct Model")
-            with gr.Column():
-                chat_history_2_7B_instruct = gr.State([])  
-                chatbot_2_7B_instruct = gr.Chatbot()
-                message_2_7B_instruct = gr.Textbox(lines=2, placeholder="Enter your message...", label="Your Message")
-            with gr.Accordion("Generation Parameters", open=False):
-                max_new_tokens_2_7B_instruct = gr.Slider(50, 1000, step=50, value=500, label="Max New Tokens")
-
-            def user_message_2_7B_instruct(message, chat_history):
-                chat_history = chat_history + [[message, None]]
-                return gr.update(value=""), chat_history, chat_history
-
-            def bot_response_2_7B_instruct(chat_history, max_new_tokens):
-                response = generate_response_2_7B_instruct(chat_history, max_new_tokens)
-                chat_history[-1][1] = response
-                return chat_history, chat_history
-
-            send_button_2_7B_instruct = gr.Button("Send")
-            send_button_2_7B_instruct.click(
-                fn=user_message_2_7B_instruct,
-                inputs=[message_2_7B_instruct, chat_history_2_7B_instruct],
-                outputs=[message_2_7B_instruct, chat_history_2_7B_instruct, chatbot_2_7B_instruct]
-            ).then(
-                fn=bot_response_2_7B_instruct,
-                inputs=[
-                    chat_history_2_7B_instruct,
-                    max_new_tokens_2_7B_instruct
-                ],
-                outputs=[chat_history_2_7B_instruct, chatbot_2_7B_instruct]
-            )
         with gr.TabItem("7B Instruct Model"):
             gr.Markdown("### Zamba2-7B Instruct Model")
             with gr.Column():
-                chat_history_7B_instruct = gr.State([])  
+                chat_history_7B_instruct = gr.State([])
                 chatbot_7B_instruct = gr.Chatbot()
                 message_7B_instruct = gr.Textbox(lines=2, placeholder="Enter your message...", label="Your Message")
             with gr.Accordion("Generation Parameters", open=False):
@@ -311,9 +230,10 @@ with gr.Blocks() as demo:
                 return gr.update(value=""), chat_history, chat_history
 
             def bot_response_7B_instruct(chat_history, max_new_tokens):
-                response = generate_response_7B_instruct(chat_history, max_new_tokens)
-                chat_history[-1][1] = response
-                return chat_history, chat_history
+                assistant_response_generator = generate_response(chat_history, max_new_tokens, model_7B_instruct, tokenizer_7B_instruct)
+                for assistant_response in assistant_response_generator:
+                    chat_history[-1][1] = assistant_response
+                    yield chat_history
 
             send_button_7B_instruct = gr.Button("Send")
             send_button_7B_instruct.click(
@@ -322,13 +242,39 @@ with gr.Blocks() as demo:
                 outputs=[message_7B_instruct, chat_history_7B_instruct, chatbot_7B_instruct]
             ).then(
                 fn=bot_response_7B_instruct,
-                inputs=[
-                    chat_history_7B_instruct,
-                    max_new_tokens_7B_instruct
-                ],
-                outputs=[chat_history_7B_instruct, chatbot_7B_instruct]
+                inputs=[chat_history_7B_instruct, max_new_tokens_7B_instruct],
+                outputs=chatbot_7B_instruct,
+            )
+
+        with gr.TabItem("2.7B Instruct Model"):
+            gr.Markdown("### Zamba2-2.7B Instruct Model")
+            with gr.Column():
+                chat_history_2_7B_instruct = gr.State([])
+                chatbot_2_7B_instruct = gr.Chatbot()
+                message_2_7B_instruct = gr.Textbox(lines=2, placeholder="Enter your message...", label="Your Message")
+            with gr.Accordion("Generation Parameters", open=False):
+                max_new_tokens_2_7B_instruct = gr.Slider(50, 1000, step=50, value=500, label="Max New Tokens")
+
+            def user_message_2_7B_instruct(message, chat_history):
+                chat_history = chat_history + [[message, None]]
+                return gr.update(value=""), chat_history, chat_history
+
+            def bot_response_2_7B_instruct(chat_history, max_new_tokens):
+                assistant_response_generator = generate_response(chat_history, max_new_tokens, model_2_7B_instruct, tokenizer_2_7B_instruct)
+                for assistant_response in assistant_response_generator:
+                    chat_history[-1][1] = assistant_response
+                    yield chat_history
+
+            send_button_2_7B_instruct = gr.Button("Send")
+            send_button_2_7B_instruct.click(
+                fn=user_message_2_7B_instruct,
+                inputs=[message_2_7B_instruct, chat_history_2_7B_instruct],
+                outputs=[message_2_7B_instruct, chat_history_2_7B_instruct, chatbot_2_7B_instruct]
+            ).then(
+                fn=bot_response_2_7B_instruct,
+                inputs=[chat_history_2_7B_instruct, max_new_tokens_2_7B_instruct],
+                outputs=chatbot_2_7B_instruct,
             )
 
 if __name__ == "__main__":
-    demo.queue().launch()
-'''
+    demo.queue().launch(max_threads=1)