Update knowledge_base.py

knowledge_base.py

@@ -1,18 +1,33 @@
 import requests
+from transformers import pipeline
+
+# Load the model and tokenizer using Hugging Face's pipeline
+qa_pipeline = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad")
+
+def load_qa_model():
+    # Initialize any model settings or variables, if needed
+    return qa_pipeline
 
 def get_answer_from_knowledge_base(question):
     """
-    This function queries a knowledge base or a public API for the answer.
-    You can replace this with more complex logic or a dynamic database query.
+    This function uses a Hugging Face model (DistilBERT) for question answering
+    from a given context.
     """
-    # Example using a mock API for current affairs and general knowledge.
-    # Replace with real-time news or current affairs API as needed.
-    
-    api_url = f"https://api.duckduckgo.com/?q={question}&format=json"
-    response = requests.get(api_url)
-    data = response.json()
-    
-    if 'AbstractText' in data and data['AbstractText']:
-        return data['AbstractText']
-    else:
-        return None
+    try:
+        # Example: Using a static context for the model to search for answers
+        context = """
+        Andhra Pradesh (AP) is a state in southern India. The current Chief Minister of Andhra Pradesh is Y. S. Jagan Mohan Reddy.
+        He assumed office in May 2019. Jagan Mohan Reddy is from the Yuvajana Sramika Rythu Congress Party (YSRCP).
+        """
+
+        # Using the question-answering pipeline from Hugging Face
+        result = qa_pipeline(question=question, context=context)
+
+        # Extract and return the answer from the result
+        if result['score'] > 0.1:  # Only return answer if the model is confident
+            return result['answer']
+        else:
+            return None
+    except Exception as e:
+        print(f"Error occurred: {e}")
+        return "Sorry, I couldn't fetch data at the moment. Please try again later."