Update chatbot.py

chatbot.py

@@ -1,18 +1,17 @@
-from knowledge_base import get_answer_from_knowledge_base, load_qa_model
+from transformers import pipeline
 
-def generate_answer(question):
-    # Check if the question is in the custom answers dictionary
-    custom_answers = {
-        "Who is the chief minister of AP?": "Y. S. Jagan Mohan Reddy",
-        # Add other custom answers if necessary
-    }
+# Load the GPT-2 model using Hugging Face's pipeline for text generation
+generator = pipeline("text-generation", model="gpt2")
 
-    if question in custom_answers:
-        return custom_answers[question]
+def generate_answer(question):
+    # Generate an answer using the GPT-2 model
+    response = generator(question, max_length=100, num_return_sequences=1)
+    
+    # Extract the generated text and return it as the answer
+    answer = response[0]['generated_text']
+    
+    # Clean the answer to remove the prompt (question)
+    answer = answer.replace(question, '').strip()
 
-    # If not, get the answer from the knowledge base using Hugging Face model
-    answer = get_answer_from_knowledge_base(question)
-    if answer:
-        return answer
-    else:
-        return "Sorry, I couldn't find an answer to that question. Please try again later."
+    # Return the cleaned answer
+    return answer