Views
No views yet
1import gradio as gr
2from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig
3import torch
4from transformers import LlamaForCausalLM, LlamaTokenizer
5
6
7class Chat:
8 def __init__(self, model, tokenizer, conv_prompt, user_alias='User', character_name='Chatbot', message_history=[], chat_buffer_size=10):
9 self.model = model
10 self.tokenizer = tokenizer
11 self.conv_prompt = conv_prompt
12 self.user_alias = user_alias
13 self.character_name = character_name
14 self.chat_buffer_size = chat_buffer_size
15 self.message_history = message_history
16 self.display_messages = []
17 for message_pairs in message_history:
18 message1, message2 = message_pairs
19 self.display_messages.append([message1['text'], message2['text']])
20
21 def evaluate(self, message, temperature=0.6, top_p=0.75, top_k=50, num_beams=5, max_new_tokens=256, repetition_penalty=1.4, **kwargs):
22 prompt = self.prompt_gen_chat(self.message_history, message)
23 inputs = self.tokenizer(prompt, return_tensors="pt")
24 input_ids = inputs["input_ids"].to(self.model.device)
25 generation_config = GenerationConfig(
26 temperature=temperature,
27 top_p=top_p,
28 top_k=top_k,
29 num_beams=num_beams,
30 early_stopping=True,
31 repetition_penalty=repetition_penalty,
32 **kwargs,
33 )
34 with torch.no_grad():
35 generation_output = self.model.generate(
36 input_ids=input_ids,
37 generation_config=generation_config,
38 return_dict_in_generate=True,
39 output_scores=True,
40 max_new_tokens=max_new_tokens,
41 )
42 s = generation_output.sequences[0]
43 output = self.tokenizer.decode(s, skip_special_tokens=True)
44 split_str = """### Response:\n{self.character_name}:"""
45 output = output.split(split_str)[1].strip()
46 return output
47
48 def gradio_helper(self, message):
49 # make response
50 response = self.evaluate(message)
51 # update message history
52 self.message_history.append(
53 (
54 {"speaker": self.user_alias, "text": message},
55 {"speaker": self.character_name, "text": response},
56 )
57 )
58 if len(self.message_history) > self.chat_buffer_size:
59 self.message_history = self.message_history[-self.chat_buffer_size:]
60 # update display messages
61 self.display_messages.append([message, response])
62 return self.display_messages
63
64 def prompt_gen_chat(self, message_history, message):
65 past_dialogue = []
66 for message_pairs in message_history:
67 message1, message2 = message_pairs
68 past_dialogue.append(f"{message1['speaker']}: {message1['text']}")
69 past_dialogue.append(f"{message2['speaker']}: {message2['text']}")
70 past_dialogue_formatted = "\n".join(past_dialogue)
71
72 prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
73
74### Instruction:
75{self.conv_prompt}
76
77This is the conversation between {self.user_alias} and {self.character_name} till now:
78{past_dialogue_formatted}
79
80Continuing from the previous conversation, write what {self.character_name} says to {self.user_alias}:
81### Input:
82{self.user_alias}: {message}
83### Response:
84{self.character_name}:"""
85
86 return prompt
87
88 def launch_gradio(self):
89 with gr.Blocks(theme="JohnSmith9982/small_and_pretty") as demo:
90 chatbot = gr.Chatbot(elem_id="chatbot")
91 with gr.Row():
92 txt = gr.Textbox(show_label=False,
93 placeholder="Enter text and press enter")
94 txt.submit(self.gradio_helper, txt, chatbot)
95 txt.submit(lambda: "", None, txt)
96
97 demo.launch(debug=True, share=True)
98
99
100if __name__ == "__main__":
101 model_path = "Xilabs/calypso-3b-alpha-v2"
102 load_in_8bit = False
103 model = LlamaForCausalLM.from_pretrained(
104 model_path, device_map="auto", load_in_8bit=load_in_8bit)
105 tokenizer = LlamaTokenizer.from_pretrained(model_path)
106 conv_prompt = "Two people are texting each other on a messaging platform."
107 message_history = [
108 (
109 {
110 "speaker": "Bob",
111 "text": "Hey, Alice! How are you doing? What's the status on those reports?",
112 },
113 {
114 "speaker": "Alice",
115 "text": "Hey, Bob! I'm doing well. I'm almost done with the reports. I'll send them to you by the end of the day.",
116 },
117 ),
118 (
119 {
120 "speaker": "Bob",
121 "text": "That's great! Thanks, Alice. I'll be waiting for them. Btw, I have approved your leave for next week.",
122 },
123 {
124 "speaker": "Alice",
125 "text": "Oh, thanks, Bob! I really appreciate it. I will be sure to send you the reports before I leave. Anything else you need from me?",
126 },
127 )
128 ]
129
130 chat_instance = Chat(model, tokenizer, conv_prompt, user_alias='Bob',
131 character_name='Alice', message_history=message_history)
132 chat_instance.launch_gradio()| Metric | Value |
|---|---|
| Avg. | 37.52 |
| ARC (25-shot) | 41.55 |
| HellaSwag (10-shot) | 71.48 |
| MMLU (5-shot) | 25.82 |
| TruthfulQA (0-shot) | 35.73 |
| Winogrande (5-shot) | 65.27 |
| GSM8K (5-shot) | 0.68 |
| DROP (3-shot) | 22.08 |