Update app.py

app.py

@@ -100,7 +100,19 @@ class StopOnTokens(StoppingCriteria):
         return False
 
 
-def predict(history_transformer_format):
+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, ""]]
 
     print(history_transformer_format)
     stop = StopOnTokens()
@@ -129,8 +141,6 @@ def predict(history_transformer_format):
 
     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(
@@ -145,27 +155,12 @@ def predict(history_transformer_format):
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
 
-    history_transformer_format[-1][1] = ""
+    partial_message = ""
     for new_token in streamer:
         if new_token != '<':
-            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 "", history_transformer_format
+            partial_message += new_token
+            yield partial_message
+    return messages
 
 # Define the Gradio interface
 title = "Tchap"
@@ -176,17 +171,9 @@ examples = [
         0.7  # temperature
     ]
 ]
-
+demo = gr.Blocks()
 with gr.Blocks() as demo:
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox()
-    clear = gr.Button("Clear")
-
-    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        predict, chatbot, chatbot
-    )
-    clear.click(lambda: None, None, chatbot, queue=False)
-
+    gr.ChatInterface(predict)
 
-demo.queue()
-demo.launch()
+if __name__ == "__main__":
+    demo.queue().launch()