Update q_learning_chatbot.py

q_learning_chatbot.py

@@ -1,7 +1,7 @@
 import os
 import numpy as np
 import pandas as pd
-from xgb_mental_health import MentalHealthClassifier
+
 import pickle
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
@@ -11,12 +11,11 @@ import torch
 
 
 class QLearningChatbot:
-    def __init__(self, states, actions, learning_rate=0.9, discount_factor=0.1):
+    def __init__(self, states, learning_rate=0.9, discount_factor=0.1):
         self.states = states
-        self.actions = actions
         self.learning_rate = learning_rate
         self.discount_factor = discount_factor
-        self.q_values = np.random.rand(len(states), len(actions))
+        self.q_values = np.random.rand(len(states))
         self.mood = "Neutral"
         self.mood_history = [] 
         self.mood_history_int = []
@@ -44,22 +43,16 @@ class QLearningChatbot:
         self.mood = highest_sentiment
         return highest_sentiment
 
-    def get_action(self, current_state):
-        current_state_index = self.states.index(current_state)
-        # print(np.argmax(self.q_values[current_state_index, :]))
-        return self.actions[np.argmax(self.q_values[current_state_index, :])]
-
-    def update_q_values(self, current_state, action, reward, next_state):
+    def update_q_values(self, current_state, reward, next_state):
         # print(f"state-reward: {current_state} - {reward} -- (b)")
         current_state_index = self.states.index(current_state)
-        action_index = self.actions.index(action)
         next_state_index = self.states.index(next_state)
 
-        current_q_value = self.q_values[current_state_index, action_index]
+        current_q_value = self.q_values[current_state_index]
         max_next_q_value = np.max(self.q_values[next_state_index, :])
 
         new_q_value = current_q_value + self.learning_rate * (reward + self.discount_factor * max_next_q_value - current_q_value)
-        self.q_values[current_state_index, action_index] = new_q_value
+        self.q_values[current_state_index] = new_q_value
 
     def update_mood_history(self):
         st.session_state.entered_mood.append(self.mood)