jon-tow commited on
Commit
9ad7e40
1 Parent(s): 0bf4d3d

Add application file

Browse files
Files changed (1) hide show
  1. app.py +84 -0
app.py ADDED
@@ -0,0 +1,84 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Model by @duyphung for @carperai
3
+ Dumb Simple Gradio by @jon-tow
4
+ """
5
+ from string import Template
6
+
7
+ import torch
8
+ import gradio as gr
9
+ from transformers import AutoTokenizer, AutoModelForCausalLM
10
+
11
+
12
+ tokenizer = AutoTokenizer.from_pretrained("CarperAI/vicuna-13b-fine-tuned-rlhf")
13
+ model = AutoModelForCausalLM.from_pretrained(
14
+ "CarperAI/vicuna-13b-fine-tuned-rlhf",
15
+ torch_dtype=torch.bfloat16,
16
+ )
17
+ model.cuda()
18
+ max_context_length = model.config.max_position_embeddings
19
+ max_new_tokens = 256
20
+
21
+
22
+ prompt_template = Template("""\
23
+ ### Human: $human
24
+ ### Assistant: $bot\
25
+ """)
26
+
27
+
28
+ def bot(history):
29
+ history = history or []
30
+
31
+ # Hack to inject prompt formatting into the history
32
+ prompt_history = []
33
+ for human, bot in history:
34
+ prompt_history.append(
35
+ prompt_template.substitute(
36
+ human=human, bot=bot if bot is not None else "")
37
+ )
38
+
39
+ prompt = "\n\n".join(prompt_history)
40
+ prompt = prompt.rstrip()
41
+ inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
42
+ # Use only the most recent context up to the maximum context length with room left over
43
+ # for the max new tokens
44
+ inputs = {k: v[:, -max_context_length + max_new_tokens:] for k, v in inputs.items()}
45
+ inputs_length = inputs['input_ids'].shape[1]
46
+
47
+ # Generate the response
48
+ tokens = model.generate(
49
+ **inputs,
50
+ # Only allow the model to generate up to 512 tokens
51
+ max_new_tokens=max_new_tokens,
52
+ num_return_sequences=1,
53
+ do_sample=True,
54
+ temperature=1.0,
55
+ top_p=1.0,
56
+ )
57
+ # Strip the initial prompt
58
+ tokens = tokens[:, inputs_length:]
59
+
60
+ # Process response
61
+ response = tokenizer.decode(tokens[0], skip_special_tokens=True)
62
+ response = response.split("###")[0].strip()
63
+
64
+ # Add the response to the history
65
+ history[-1][1] = response
66
+ return history
67
+
68
+
69
+ def user(user_message, history):
70
+ return "", history + [[user_message, None]]
71
+
72
+
73
+ with gr.Blocks() as demo:
74
+ gr.Markdown("""Vicuna-13B RLHF Chatbot""")
75
+ chatbot = gr.Chatbot([], elem_id="chatbot").style(height=512)
76
+ msg = gr.Textbox()
77
+ clear = gr.Button("Clear")
78
+ state = gr.State([])
79
+
80
+ msg.submit(user, inputs=[msg, chatbot], outputs=[msg, chatbot], queue=False).then(
81
+ bot, chatbot, chatbot)
82
+ clear.click(lambda: None, None, chatbot, queue=False)
83
+
84
+ demo.launch(share=True)