Cleanups

app.py

@@ -1,286 +1,3 @@
-'''
-import os
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-
-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(
- model_name_2_7B_instruct, device_map="cuda", torch_dtype=torch.bfloat16
-)
-
-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(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.apply_chat_template(sample, tokenize=False)
- input_ids = tokenizer(chat_sample, return_tensors='pt', add_special_tokens=False).to(model.device)
-
- max_new_tokens = int(max_new_tokens)
- max_input_length = max_context_length - max_new_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.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.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,
- repetition_penalty=repetition_penalty,
- num_beams=int(num_beams),
- length_penalty=length_penalty,
- num_return_sequences=1
- )
- """
- generated_text = tokenizer.decode(outputs[0])
- assistant_response = extract_assistant_response(generated_text)
-
- 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("7B Instruct Model"):
- gr.Markdown("### Zamba2-7B Instruct Model")
- 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_7B_instruct(message, chat_history):
- chat_history = chat_history + [[message, None]]
- return gr.update(value=""), chat_history, chat_history
-
- def bot_response_7B_instruct(chat_history, max_new_tokens):
- response = generate_response(chat_history, max_new_tokens, model_7B_instruct, tokenizer_7B_instruct)
- chat_history[-1][1] = response
- return chat_history, 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]
- )
- 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, 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(
- model_name_2_7B_instruct, device_map="cuda", torch_dtype=torch.bfloat16
-)
-
-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 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.apply_chat_template(sample, tokenize=False)
- input_ids = tokenizer(chat_sample, return_tensors='pt', add_special_tokens=False).to(model.device)
-
- max_new_tokens = int(max_new_tokens)
- max_input_length = max_context_length - max_new_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:]
-
- streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
- generation_kwargs = dict(**input_ids, max_new_tokens=int(max_new_tokens), streamer=streamer)
-
- thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
- thread.start()
-
- assistant_response = ""
-
- 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
- torch.cuda.empty_cache()
-
-with gr.Blocks() as demo:
- gr.Markdown("# Zamba2 Model Selector")
- with gr.Tabs():
- with gr.TabItem("7B Instruct Model"):
- gr.Markdown("### Zamba2-7B Instruct Model")
- 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")
-
- def user_message_7B_instruct(message, chat_history):
- chat_history = chat_history + [[message, None]]
- return gr.update(value=""), chat_history, chat_history
-
- def bot_response_7B_instruct(chat_history, max_new_tokens):
- 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(
- 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=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(max_threads=1)
-'''
-
 import os
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer