Do instruct models for 2.7 and 7 b sizes

app.py

@@ -3,26 +3,50 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# Define models as None to delay loading
-model, model_instruct = None, None
-tokenizer, tokenizer_instruct = None, None
+model_name_2_7B_instruct = "Zyphra/Zamba2-2.7B-instruct"
+model_name_7B_instruct = "Zyphra/Zamba2-7B-instruct"
 
-def generate_response_base(input_text, max_new_tokens, temperature, top_k, top_p, repetition_penalty, num_beams, length_penalty):
-    global model, tokenizer
-    if model is None:
-        tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba2-7B")
-        model = AutoModelForCausalLM.from_pretrained(
-            "Zyphra/Zamba2-7B", device_map="cuda", torch_dtype=torch.bfloat16
-        )
-    selected_model = model
-    selected_tokenizer = tokenizer
+tokenizer_2_7B_instruct = AutoTokenizer.from_pretrained(model_name_2_7B_instruct)
+model_2_7B_instruct = AutoModelForCausalLM.from_pretrained(
+    model_name_2_7B_instruct, device_map="cuda", torch_dtype=torch.bfloat16
+)
 
-    # Tokenize and generate response
-    input_ids = selected_tokenizer(input_text, return_tensors="pt").input_ids.to(selected_model.device)
-    outputs = selected_model.generate(
+tokenizer_7B_instruct = AutoTokenizer.from_pretrained(model_name_7B_instruct)
+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):
+    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).to(model_2_7B_instruct.device)
+    outputs = model_2_7B_instruct.generate(**input_ids, max_new_tokens=int(max_new_tokens), return_dict_in_generate=False, output_scores=False, use_cache=True, num_beams=1, do_sample=False)
+    """
+    outputs = model_2_7B_instruct.generate(
         input_ids=input_ids,
         max_new_tokens=int(max_new_tokens),
         do_sample=True,
+        use_cache=True,
         temperature=temperature,
         top_k=int(top_k),
         top_p=top_p,
@@ -31,34 +55,27 @@ def generate_response_base(input_text, max_new_tokens, temperature, top_k, top_p
         length_penalty=length_penalty,
         num_return_sequences=1
     )
-    response = selected_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+    """
+    generated_text = tokenizer_2_7B_instruct.decode(outputs[0])
+    assistant_response = extract_assistant_response(generated_text)
+    return assistant_response
 
-def generate_response_instruct(chat_history, max_new_tokens, temperature, top_k, top_p, repetition_penalty, num_beams, length_penalty):
-    global model_instruct, tokenizer_instruct
-    if model_instruct is None:
-        tokenizer_instruct = AutoTokenizer.from_pretrained("Zyphra/Zamba2-7B-instruct")
-        model_instruct = AutoModelForCausalLM.from_pretrained(
-            "Zyphra/Zamba2-7B-instruct", device_map="cuda", torch_dtype=torch.bfloat16
-        )
-    selected_model = model_instruct
-    selected_tokenizer = tokenizer_instruct
-
-    # Build the sample
+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]})
-    # Format the chat sample
-    chat_sample = selected_tokenizer.apply_chat_template(sample, tokenize=False)
-    # Tokenize input and generate output
-    input_ids = selected_tokenizer(chat_sample, return_tensors='pt', add_special_tokens=False).input_ids.to(selected_model.device)
-    outputs = selected_model.generate(
+    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).to(model_7B_instruct.device)
+    outputs = model_7B_instruct.generate(**input_ids, max_new_tokens=int(max_new_tokens), return_dict_in_generate=False, output_scores=False, use_cache=True, num_beams=1, do_sample=False)
+    """
+    outputs = model_7B_instruct.generate(
         input_ids=input_ids,
         max_new_tokens=int(max_new_tokens),
         do_sample=True,
+        use_cache=True,
         temperature=temperature,
         top_k=int(top_k),
         top_p=top_p,
@@ -67,53 +84,88 @@ def generate_response_instruct(chat_history, max_new_tokens, temperature, top_k,
         length_penalty=length_penalty,
         num_return_sequences=1
     )
-    response = selected_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
-
-def clear_text():
-    return ""
+    """
+    generated_text = tokenizer_7B_instruct.decode(outputs[0])
+    assistant_response = extract_assistant_response(generated_text)
+    return assistant_response
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Zamba2-7B Model Selector")
+    gr.Markdown("# Zamba2 Model Selector")
     with gr.Tabs():
-        with gr.TabItem("Base Model"):
-            gr.Markdown("### Zamba2-7B Base Model")
-            input_text = gr.Textbox(lines=2, placeholder="Enter your input text...", label="Input Text")
-            output_text = gr.Textbox(label="Generated Response")
-            max_new_tokens = gr.Slider(50, 1000, step=50, value=500, label="Max New Tokens")
-            temperature = gr.Slider(0.1, 1.5, step=0.1, value=0.7, label="Temperature")
-            top_k = gr.Slider(1, 100, step=1, value=50, label="Top K")
-            top_p = gr.Slider(0.1, 1.0, step=0.1, value=0.9, label="Top P")
-            repetition_penalty = gr.Slider(1.0, 2.0, step=0.1, value=1.2, label="Repetition Penalty")
-            num_beams = gr.Slider(1, 10, step=1, value=5, label="Number of Beams")
-            length_penalty = gr.Slider(0.0, 2.0, step=0.1, value=1.0, label="Length Penalty")
-            submit_button = gr.Button("Generate Response")
-            submit_button.click(fn=generate_response_base, inputs=[input_text, max_new_tokens, temperature, top_k, top_p, repetition_penalty, num_beams, length_penalty], outputs=output_text)
-            submit_button.click(fn=clear_text, outputs=input_text)
-        with gr.TabItem("Instruct Model"):
+        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_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")
-            chat_history = gr.Chatbot()
-            message = gr.Textbox(lines=2, placeholder="Enter your message...", label="Your Message")
-            max_new_tokens_instruct = gr.Slider(50, 1000, step=50, value=500, label="Max New Tokens")
-            temperature_instruct = gr.Slider(0.1, 1.5, step=0.1, value=0.7, label="Temperature")
-            top_k_instruct = gr.Slider(1, 100, step=1, value=50, label="Top K")
-            top_p_instruct = gr.Slider(0.1, 1.0, step=0.1, value=0.9, label="Top P")
-            repetition_penalty_instruct = gr.Slider(1.0, 2.0, step=0.1, value=1.2, label="Repetition Penalty")
-            num_beams_instruct = gr.Slider(1, 10, step=1, value=5, label="Number of Beams")
-            length_penalty_instruct = gr.Slider(0.0, 2.0, step=0.1, value=1.0, label="Length Penalty")
+            with gr.Column():
+                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):
+                max_new_tokens_7B_instruct = gr.Slider(50, 1000, step=50, value=500, label="Max New Tokens")
+                # temperature_7B_instruct = gr.Slider(0.1, 1.5, step=0.1, value=0.2, label="Temperature")
+                # top_k_7B_instruct = gr.Slider(1, 100, step=1, value=50, label="Top K")
+                # top_p_7B_instruct = gr.Slider(0.1, 1.0, step=0.1, value=1.0, label="Top P")
+                # repetition_penalty_7B_instruct = gr.Slider(1.0, 2.0, step=0.1, value=1.2, label="Repetition Penalty")
+                # num_beams_7B_instruct = gr.Slider(1, 10, step=1, value=1, label="Number of Beams")
+                # length_penalty_7B_instruct = gr.Slider(0.0, 2.0, step=0.1, value=1.0, label="Length Penalty")
 
-            def user_message(message, chat_history):
+            def user_message_7B_instruct(message, chat_history):
                 chat_history = chat_history + [[message, None]]
-                return "", chat_history
+                return gr.update(value=""), chat_history, chat_history
 
-            def bot_response(chat_history):
-                response = generate_response_instruct(chat_history, max_new_tokens_instruct, temperature_instruct, top_k_instruct, top_p_instruct, repetition_penalty_instruct, num_beams_instruct, length_penalty_instruct)
+            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
+                return chat_history, chat_history
 
-            message.submit(user_message, [message, chat_history], [message, chat_history], queue=False).then(
-                bot_response, inputs=[chat_history], outputs=[chat_history]
+            send_button_7B_instruct = gr.Button("Send")
+            send_button_7B_instruct.click(
+                fn=user_message_7B_instruct,
+                inputs=[message_7B_instruct, chat_history_7B_instruct],
+                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]
             )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue().launch()