Update app.py

app.py

@@ -4,18 +4,20 @@ import torch
 import gradio as gr
 
 # checkpoint = "tiiuae/falcon-40b-instruct"
-# checkpoint = "tiiuae/falcon-7b-instruct"
+checkpoint = "tiiuae/falcon-7b-instruct"
 # checkpoint = "tiiuae/falcon-7b"
-checkpoint = "gpt2"
+# checkpoint = "gpt2"
 
 
 model = AutoModelForCausalLM.from_pretrained(
     checkpoint, device_map="auto", offload_folder="offload", 
     trust_remote_code=True,
+    # torch_dtype="auto",
 )
 
 tokenizer = AutoTokenizer.from_pretrained(checkpoint, 
                                           trust_remote_code=True,
+                                          torch_dtype="auto"
                                          )
 
 pipeline = transformers.pipeline(
@@ -28,41 +30,103 @@ pipeline = transformers.pipeline(
     device_map="auto",
 )
 
-def format_chat_prompt(message, chat_history):
-    prompt = ""
+# def format_chat_prompt(message, chat_history):
+#     prompt = ""
+#     for turn in chat_history:
+#         user_message, bot_message = turn
+#         prompt = f"{prompt}\nUser: {user_message}\nAssistant: {bot_message}"
+#     prompt = f"{prompt}\nUser: {message}\nAssistant:"
+#     return prompt
+
+# def respond(message, chat_history):
+#         formatted_prompt = format_chat_prompt(message, chat_history)
+#         # print(formatted_prompt)
+#         bot_message = generate_seqs(prompt = formatted_prompt,
+#                                      max_new_tokens=1024,
+#                                      stop_sequence=["\nUser:", "<|endoftext|>"]).split('Assistant: ')[-1]
+
+#         chat_history.append((message, bot_message))
+#         return "", chat_history
+
+# def generate_seqs(prompt, max_new_tokens=None, stop_sequence=None):
+#     output = pipeline(prompt,
+#       max_length=200,
+#       max_new_tokens = max_new_tokens,
+#       stop_sequence = stop_sequence,
+#       do_sample=True,
+#       top_k=10,
+#       num_return_sequences=1,
+#       eos_token_id=tokenizer.eos_token_id)
+#     return output[0]['generated_text']
+
+# with gr.Blocks() as demo:
+#     chatbot = gr.Chatbot(height=240) #just to fit the notebook
+#     msg = gr.Textbox(label="Prompt")
+#     btn = gr.Button("Submit")
+#     clear = gr.ClearButton(components=[msg, chatbot], value="Clear console")
+
+#     btn.click(respond, inputs=[msg, chatbot], outputs=[msg, chatbot])
+#     msg.submit(respond, inputs=[msg, chatbot], outputs=[msg, chatbot]) #Press enter to submit
+# demo.launch()
+
+
+def format_chat_prompt(message, chat_history, instruction):
+    prompt = f"System:{instruction}"
     for turn in chat_history:
         user_message, bot_message = turn
         prompt = f"{prompt}\nUser: {user_message}\nAssistant: {bot_message}"
     prompt = f"{prompt}\nUser: {message}\nAssistant:"
     return prompt
 
-def respond(message, chat_history):
-        formatted_prompt = format_chat_prompt(message, chat_history)
-        # print(formatted_prompt)
-        bot_message = generate_seqs(prompt = formatted_prompt,
-                                     max_new_tokens=1024,
-                                     stop_sequence=["\nUser:", "<|endoftext|>"]).split('Assistant: ')[-1]
-
-        chat_history.append((message, bot_message))
-        return "", chat_history
-
-def generate_seqs(prompt, max_new_tokens=None, stop_sequence=None):
+def generate_seqs(prompt, max_new_tokens=None, stop_sequence=None, temperature=None):
     output = pipeline(prompt,
       max_length=200,
       max_new_tokens = max_new_tokens,
       stop_sequence = stop_sequence,
+      temperature=temperature,
       do_sample=True,
       top_k=10,
       num_return_sequences=1,
       eos_token_id=tokenizer.eos_token_id)
     return output[0]['generated_text']
 
+def respond(message, chat_history, instruction, temperature=0.7):
+    prompt = format_chat_prompt(message, chat_history, instruction)
+    chat_history = chat_history + [[message, ""]]
+    stream = generate_seqs(prompt = prompt,
+                                    max_new_tokens=1024,
+                                    stop_sequence=["\nUser:", "<|endoftext|>"],
+                                    temperature=temperature).split('Assistant: ')[-1]
+                                    #stop_sequence to not generate the user answer
+    acc_text = ""
+    #Streaming the tokens
+    for idx, response in enumerate(stream):
+            # text_token = response.token.text
+            text_token = response
+
+            # if response.details:
+            #     return
+
+            if idx == 0 and text_token.startswith(" "):
+                text_token = text_token[1:]
+
+            acc_text += text_token
+            last_turn = list(chat_history.pop(-1))
+            last_turn[-1] += acc_text
+            chat_history = chat_history + [last_turn]
+            yield "", chat_history
+            acc_text = ""
+
 with gr.Blocks() as demo:
     chatbot = gr.Chatbot(height=240) #just to fit the notebook
     msg = gr.Textbox(label="Prompt")
+    with gr.Accordion(label="Advanced options",open=False):
+        system = gr.Textbox(label="System message", lines=2, value="A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.")
+        temperature = gr.Slider(label="temperature", minimum=0.1, maximum=1, value=0.7, step=0.1)
     btn = gr.Button("Submit")
     clear = gr.ClearButton(components=[msg, chatbot], value="Clear console")
 
-    btn.click(respond, inputs=[msg, chatbot], outputs=[msg, chatbot])
-    msg.submit(respond, inputs=[msg, chatbot], outputs=[msg, chatbot]) #Press enter to submit
-demo.launch()
+    btn.click(respond, inputs=[msg, chatbot, system], outputs=[msg, chatbot])
+    msg.submit(respond, inputs=[msg, chatbot, system], outputs=[msg, chatbot]) #Press enter to submit
+gr.close_all()
+demo.queue().launch()