Update src/chat.py

src/chat.py

@@ -16,6 +16,9 @@ class SchoolChatbot:
         """
         model_id = MY_MODEL if MY_MODEL else BASE_MODEL # define MY_MODEL in config.py if you create a new model in the HuggingFace Hub
         self.client = InferenceClient(model=model_id, token=HF_TOKEN)
+        self.df = pd.read_csv("bps_data.csv")
+        with open("keyword_to_column_map.json") as f:
+            self.keyword_map = json.load(f)
         
     def format_prompt(self, user_input):
         """
@@ -42,6 +45,25 @@ class SchoolChatbot:
             f"<|user|>{user_input}<|end|>\n"
             "<|assistant|>"
         )
+
+    def lookup_structured_data(self, query, school_name=None):
+        """Search the structured BPS dataset for relevant information"""
+        results = []
+        df_filtered = self.df
+
+        if school_name:
+            df_filtered = self.df[self.df["BPS_School_Name"].str.contains(school_name, case=False, na=False)]
+
+        if df_filtered.empty:
+            return None
+
+        row = df_filtered.iloc[0]
+        for key, col in self.keyword_map.items():
+            if key in query.lower():
+                val = row.get(col, "N/A")
+                if pd.notna(val):
+                    results.append(f"{key.title()}: {val}")
+        return "\n".join(results) if results else None
         
     def get_response(self, user_input):
         """
@@ -62,6 +84,17 @@ class SchoolChatbot:
         - Use self.format_prompt() to format the user's input
         - Use self.client to generate responses
         """
+        matched_school = None
+        for name in self.df["BPS_School_Name"].dropna():
+            if name.lower() in user_input.lower():
+                matched_school = name
+                break
+
+        structured_response = self.lookup_structured_data(user_input, matched_school)
+
+        if structured_response:
+            return f"Here’s what I found based on school data:\n{structured_response}"
+
         prompt = self.format_prompt(user_input)
         response = self.client.text_generation(
             prompt,