Spaces:
Runtime error
Runtime error
File size: 2,289 Bytes
b308128 937be2f 7f877a9 b308128 8eb0f9a 7eaa7b0 937be2f 04fc021 853ae4d edb0bcd 853ae4d edb0bcd 853ae4d edb0bcd 853ae4d a450a5f 853ae4d a450a5f 853ae4d 8eb0f9a 853ae4d 8eb0f9a 853ae4d 8eb0f9a ac4f141 853ae4d a450a5f 853ae4d 5e8be56 853ae4d 8c245db |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import time
# Load the Vicuna 7B v1.3 LMSys model and tokenizer
model_name = "lmsys/vicuna-7b-v1.3"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
def create_chatbot_tab(description, prompt_placeholder, strategy_labels):
tab = gr.Tab(description)
prompt_textbox = gr.Textbox(show_label=False, placeholder=prompt_placeholder)
chatbots = []
for strategy_label in strategy_labels:
with gr.Row():
vicuna_chatbot = gr.Chatbot(label="vicuna-7b")
llama_chatbot = gr.Chatbot(label="llama-7b")
gpt_chatbot = gr.Chatbot(label="gpt-3.5")
chatbots.append(vicuna_chatbot)
clear_button = gr.ClearButton([prompt_textbox] + chatbots)
return tab, prompt_textbox, chatbots, clear_button
def create_submit_function(prompt_textbox, chatbots):
def respond(message, chat_history):
input_ids = tokenizer.encode(message, return_tensors="pt")
output_ids = model.generate(input_ids, max_length=50, num_beams=5, no_repeat_ngram_size=2)
bot_message = tokenizer.decode(output_ids[0], skip_special_tokens=True)
chat_history.append((message, bot_message))
time.sleep(2)
return "", chat_history
for chatbot in chatbots:
prompt_textbox.submit(respond, [prompt_textbox, chatbot], [prompt_textbox, chatbot])
# Create POS and Chunk tabs
pos_tab, pos_prompt_textbox, pos_chatbots, pos_clear_button = create_chatbot_tab("POS", "Write a prompt and press enter", ["Strategy 1 QA", "Strategy 2 Instruction", "Strategy 3 Structured Prompting"])
chunk_tab, chunk_prompt_textbox, chunk_chatbots, chunk_clear_button = create_chatbot_tab("Chunk", "Write a prompt and press enter", ["Strategy 1 QA", "Strategy 2 Instruction", "Strategy 3 Structured Prompting"])
# Create submit functions for POS and Chunk tabs
create_submit_function(pos_prompt_textbox, pos_chatbots)
create_submit_function(chunk_prompt_textbox, chunk_chatbots)
# Launch the demo with POS and Chunk tabs and separate "Clear" buttons
demo = gr.Blocks()
demo.append(pos_tab)
demo.append(pos_clear_button)
demo.append(chunk_tab)
demo.append(chunk_clear_button)
demo.launch()
|