Update app.py

app.py

@@ -173,15 +173,52 @@ with gr.Blocks() as app:
         combine_button.click(combine_datasets, outputs=[combined_output, download_output])
 
     with gr.Tab("Train and Chat"):
-        gr.Markdown("Train a chatbot with a selected dataset and interact with it.")
+    gr.Markdown("**Train a chatbot with a selected dataset and interact with it.**")
 
-        chat_dataset = gr.Textbox(label="Dataset for Training", placeholder="Paste or load a dataset for training.")
-        train_button = gr.Button("Train Chatbot")
-        chatbot = gr.Chatbot(label="Chat with Trained Bot")
+    chat_dataset = gr.Textbox(
+        label="Dataset for Training",
+        placeholder="Paste or load a dataset for training.",
+        lines=5,
+    )
+    train_button = gr.Button("Train Chatbot")
+    chatbot = gr.Chatbot(label="Chat with Trained Bot", type="messages")
+    user_input = gr.Textbox(
+        label="Your Message",
+        placeholder="Type a message and press Enter...",
+        lines=1,
+    )
+
+    # Persistent system message with dataset knowledge
+    system_message = {"system": "You are a bot trained on the following dataset:"}
+    bot_knowledge = {"dataset": None}
+
+    # Train the chatbot by setting the dataset
+    def train_chatbot(dataset):
+        bot_knowledge["dataset"] = dataset
+        return "Chatbot trained successfully!"
 
-        system_message = {"system": "You are a bot trained on the following dataset:"}
+    # Chat function for handling user messages
+    def chat_with_bot(history, user_message):
+        if not bot_knowledge["dataset"]:
+            return history + [{"role": "bot", "content": "No dataset loaded. Please train the bot first."}]
 
-        train_button.click(train_chatbot, inputs=[chat_dataset], outputs=None)
-        chatbot.click(chat_with_bot, inputs=[chatbot, gr.Textbox(label="User Input")], outputs=chatbot)
+        # Append user input to history
+        history.append({"role": "user", "content": user_message})
+
+        # Generate bot response based on the dataset
+        prompt = f"{bot_knowledge['dataset']} {user_message}"
+        response = client.text_generation(prompt=prompt, max_new_tokens=128)["generated_text"]
+
+        # Append bot response to history
+        history.append({"role": "bot", "content": response})
+        return history
+
+    # Train button event
+    train_button.click(train_chatbot, inputs=[chat_dataset], outputs=None)
+
+    # User input submission event
+    user_input.submit(
+        chat_with_bot, inputs=[chatbot, user_input], outputs=chatbot
+    )
 
 app.launch()