Update app.py

app.py

@@ -70,14 +70,58 @@ def generate_text(message, history):
     history = ["init", input_prompt]
 
 
-demo = gr.ChatInterface(
-    generate_text,
-    title="A N I M A",
-    description="ANIMA is an expert in various scientific disciplines.",
-    examples=["tell me everything about biomimicry"],
-    cache_examples=True,
-    retry_btn=None,
-    undo_btn="Delete Previous",
-    clear_btn="Clear",
-)
+# User message function
+def user(message, history):
+    new_history = history + [[message, None]]
+    return "", new_history
+
+# Bot message function
+def bot(history, system_prompt):
+    tokens = model.tokenize(system_prompt.encode("utf-8"))
+    
+    for user_message, bot_message in history[:-1]:
+        message_tokens = model.tokenize(user_message.encode("utf-8"))
+        tokens.extend(message_tokens)
+        
+        if bot_message:
+            message_tokens = model.tokenize(bot_message.encode("utf-8"))
+            tokens.extend(message_tokens)
+
+    last_user_message = history[-1][0]
+    message_tokens = model.tokenize(last_user_message.encode("utf-8"))
+    tokens.extend(message_tokens)
+
+    generator = model.generate(tokens)
+
+    partial_text = ""
+    for i, token in enumerate(generator):
+        if token == model.token_eos() or (max_new_tokens is not None and i >= max_new_tokens):
+            break
+        partial_text += model.detokenize([token]).decode("utf-8", "ignore")
+    history[-1][1] = partial_text
+    return history
+
+# Gradio UI
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    favicon = '<img src="https://cdn-uploads.huggingface.co/production/uploads/64740cf7485a7c8e1bd51ac9/seIR5ErFdX5Snr4O7r7tY.png" width="68px" style="display: inline">'
+    gr.Markdown(f"""<h1><center>{favicon}A N I M A</center></h1>ANIMA is an expert in various scientific disciplines.""")
+    with gr.Row():
+        with gr.Column(scale=5):
+            
+            chatbot = gr.Chatbot(label="Dialogue").style(height=400)
+    with gr.Row():
+        with gr.Column():
+            msg = gr.Textbox(label="Send Message", placeholder="Send Message", show_label=False).style(container=False)
+        with gr.Column():
+            with gr.Row():
+                submit = gr.Button("Send")
+                stop = gr.Button("Stop")
+                clear = gr.Button("Clear")
+    with gr.Row():
+        gr.Markdown("""WARNING: The model may generate factually or ethically incorrect texts. We are not responsible for this.""")
+    
+    submit_event = msg.submit(fn=user, inputs=[msg, chatbot], outputs=[msg, chatbot], queue=False).success(fn=bot, inputs=[chatbot, system_prompt], outputs=chatbot, queue=True)
+    submit_click_event = submit.click(fn=user, inputs=[msg, chatbot], outputs=[msg, chatbot], queue=False).success(fn=bot, inputs=[chatbot, system_prompt], outputs=chatbot, queue=True)
+    stop.click(fn=None, inputs=None, outputs=None, cancels=[submit_event, submit_click_event], queue=False)
+    clear.click(lambda: None, None, chatbot, queue=False)
 demo.launch(max_threads=10)