Update app.py

app.py

@@ -100,19 +100,7 @@ class StopOnTokens(StoppingCriteria):
         return False
 
 
-def predict(message, history):
-
-    global source_text
-    global assess_rag
-    #For now, we only query the vector database once, at the start.
-    if len(history) == 0:
-        assess_rag = classification_chatrag(message)
-        if assess_rag:
-            source_text = vector_search(message)
-        else:
-            source_text = "Albert-Tchap n'utilise pas de sources comme votre requête n'a pas l'air d'en recueillir."
-    
-    history_transformer_format = history + [[message, ""]]
+def predict(history_transformer_format):
 
     print(history_transformer_format)
     stop = StopOnTokens()
@@ -141,6 +129,8 @@ def predict(message, history):
 
     messages = system_prompt + messages
 
+    print(messages)
+
     model_inputs = tokenizer([messages], return_tensors="pt").to("cuda")
     streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
@@ -155,12 +145,27 @@ def predict(message, history):
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
 
-    partial_message = ""
+    history_transformer_format[-1][1] = ""
     for new_token in streamer:
         if new_token != '<':
-            partial_message += new_token
-            yield partial_message
-    return messages
+            history_transformer_format[-1][1] += new_token
+            yield history_transformer_format
+
+def user(message, history):
+    global source_text
+    global assess_rag
+    #For now, we only query the vector database once, at the start.
+    if len(history) == 0:
+        assess_rag = classification_chatrag(message)
+        if assess_rag:
+            source_text = vector_search(message)
+        else:
+            source_text = "Albert-Tchap n'utilise pas de sources comme votre requête n'a pas l'air d'en recueillir."
+    
+    history_transformer_format = history + [[message, ""]]
+
+    print(history_transformer_format)
+    return source_text, history_transformer_format
 
 # Define the Gradio interface
 title = "Tchap"
@@ -171,9 +176,21 @@ examples = [
         0.7  # temperature
     ]
 ]
-demo = gr.Blocks()
+
 with gr.Blocks() as demo:
-    gr.ChatInterface(predict)
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox()
+    clear = gr.Button("Clear")
+    
+    user_output = gr.HTML()  # To display the user's message
+    history = gr.State()
+
+    msg.submit(user, inputs=[msg, history], outputs=[user_output, history], queue=False).then(
+        predict, chatbot, chatbot
+    )
+    
+    clear.click(lambda: None, None, chatbot, queue=False)
+
 
-if __name__ == "__main__":
-    demo.queue().launch()
+demo.queue()
+demo.launch()