Update app.py

app.py

@@ -1,32 +1,34 @@
 import gradio as gr
 from transformers import pipeline
 
-# Load a text generation model
-generator = pipeline(
+# Load the model
+chatbot_pipeline = pipeline(
     "text-generation",
-    model="microsoft/DialoGPT-small",
-    max_length=200,
-    pad_token_id=50256
+    model="microsoft/DialoGPT-small"
 )
 
-def chat_fn(message, history):
-    # Flatten history into a single prompt
-    prompt_text = ""
-    for user, bot in history:
-        prompt_text += f"User: {user}\nBot: {bot}\n"
-    prompt_text += f"User: {message}\nBot:"
+# Define a function to handle the conversation
+def respond(message, history=None):
+    if history is None:
+        history = []
 
-    # Generate reply
-    response = generator(prompt_text, num_return_sequences=1)
-    bot_reply = response[0]['generated_text'].split("Bot:")[-1].strip()
-    return bot_reply
+    # Merge previous messages into a single conversation string
+    conversation = ""
+    for user_msg, bot_msg in history:
+        conversation += f"User: {user_msg}\nBot: {bot_msg}\n"
 
-demo = gr.ChatInterface(
-    fn=chat_fn,
-    title="🤖 My Hugging Face Chatbot",
-    description="A chatbot running on Hugging Face Spaces with conversation memory.",
-    examples=["Hello!", "Tell me a joke", "What's the capital of France?"]
-)
+    conversation += f"User: {message}\nBot:"
+
+    # Generate response
+    output = chatbot_pipeline(
+        conversation,
+        max_length=200,
+        pad_token_id=50256
+    )
+
+    reply = output[0]['generated_text'].split("Bot:")[-1].strip()
+    history.append((message, reply))
+    return reply, history
 
-if __name__ == "__main__":
-    demo.launch()
+# Create the Gradio interface
+gr.ChatInterface(fn=respond).launch()