Update chatbot.py

chatbot.py

@@ -1,10 +1,10 @@
 from transformers import pipeline
 
-# Load the GPT-2 model using Hugging Face's pipeline for text generation
+# Initialize the GPT-2 model using Hugging Face's pipeline for text generation
 generator = pipeline("text-generation", model="gpt2")
 
 def generate_answer(question):
-    # Generate an answer using the GPT-2 model
+    # Generate an answer using GPT-2
     response = generator(question, max_length=100, num_return_sequences=1)
     
     # Extract the generated text and return it as the answer
@@ -13,5 +13,10 @@ def generate_answer(question):
     # Clean the answer to remove the prompt (question)
     answer = answer.replace(question, '').strip()
 
-    # Return the cleaned answer
-    return answer
+    # Format the answer for Gradio's chatbot component using the 'messages' format
+    messages = [
+        {"role": "user", "content": question},
+        {"role": "assistant", "content": answer}
+    ]
+    
+    return messages