Notes

# Define the environment

class CustomEnv:
def __init__(self, data,tokenizer):
self.data = data # Your data containing questions, correct, and incorrect responses
self.tokenizer = tokenizer
self.current_index = 0
self.max_steps = len(data) # Assuming one episode per data entry

def reset(self):
# Initialize the environment at the beginning of an episode
self.current_index = 0
state = {
"question": self.tokenizer.tokenize(self.data.iloc[self.current_index]['question']),
"correct_response": self.tokenizer.tokenize(self.data.iloc[self.current_index]['correct_response']),
"incorrect_response": self.tokenizer.tokenize(self.data.iloc[self.current_index]['incorrect_response']),
}
return state

def step(self, action):
# Take an action (e.g., select a response) and update the environment
# Calculate the reward based on the action and state
# Return the next state, reward, and whether the episode is done
if self.current_index < self.max_steps - 1:
self.current_index += 1
next_state = {
"question": self.tokenizer.tokenize(self.data.iloc[self.current_index]['question']),
"correct_response": self.tokenizer.tokenize(self.data.iloc[self.current_index]['correct_response']),
"incorrect_response": self.tokenizer.tokenize(self.data.iloc[self.current_index]['incorrect_response']),
}
reward = calculate_reward(action, next_state) # Implement your custom reward function
done = False
else:
# End of episode
next_state = None
reward = 0
done = True

return next_state, reward, done



pip install --upgrade torch

import torch.nn as nn
#import torch.distributions.Categorical as Categorical

class PolicyNetwork(nn.Module):
def __init__(self, input_size, action_dim, hidden_size=64):
super(PolicyNetwork, self).__init__()
self.rnn_question = nn.LSTM(input_size, hidden_size, num_layers=1, batch_first=True)
self.rnn_correct_response = nn.LSTM(input_size, hidden_size, num_layers=1, batch_first=True)
self.rnn_incorrect_response = nn.LSTM(input_size, hidden_size, num_layers=1, batch_first=True)
self.fc = nn.Sequential(
nn.Linear(3 * hidden_size, 64),
nn.ReLU(),
nn.Linear(64, action_dim),
nn.Softmax(dim=-1)
)

def forward(self, question, correct_response, incorrect_response):
_, (question_hidden, _) = self.rnn_question(question)
_, (correct_response_hidden, _) = self.rnn_correct_response(correct_response)
_, (incorrect_response_hidden, _) = self.rnn_incorrect_response(incorrect_response)
combined_state = torch.cat((question_hidden, correct_response_hidden, incorrect_response_hidden), dim=-1)
return self.fc(combined_state)

def select_action(self, state):
# Forward pass through the policy network
action_probs = self.forward(state['question'], state['correct_response'], state['incorrect_response'])
# Create a categorical distribution over the action probabilities
dist = Categorical(action_probs)
# Sample an action from the distribution
action = dist.sample()
# Return the selected action as an integer
return action.item()

def calculate_reward(action, state):
# Implement your custom reward logic here
# Example: Reward +1 for selecting the correct response, -1 for selecting the incorrect response, 0 otherwise
if action == "correct":
reward = 1
elif action == "incorrect":
reward = -1
else:
reward = 0
return reward

from pandas_ods_reader import read_ods
data = read_ods("/content/drive/MyDrive/sample_data/Fire_security.ods", columns=["question", "correct_response","incorrect_response"])


data = data.head(3)
data

max_question_length = max(len(tokenizer.tokenize(row['question'])) for _, row in data.iterrows())
max_correct_response_length = max(len(tokenizer.tokenize(row['correct_response'])) for _, row in data.iterrows())
max_incorrect_response_length = max(len(tokenizer.tokenize(row['incorrect_response'])) for _, row in data.iterrows())

input_size = max_question_length + max_correct_response_length + max_incorrect_response_length

input_size

# Initialize the environment and tokenizer
tokenizer = Tokenizer()
env = CustomEnv(data, tokenizer)
policy = PolicyNetwork(input_size, 2) # Adapt input_size and action_dim
# Define the optimizer
optimizer = optim.Adam(policy.parameters(), lr=0.01)

question_sequence = "".join([tokenizer.token_to_character(token) for token in question_sequence])


num_episodes = 10
max_question_length,max_correct_response_length,max_incorrect_response_length
# Training loop
for episode in range(num_episodes):
episode_states = []
episode_actions = []
episode_rewards = []

state = env.reset()

while True:

# # Extract the question from the state (assuming state contains the question)
#question_sequence = state["question"]
# question_sequence = "".join([tokenizer.token_to_character(token) for token in question_sequence])
# response_text = generate_response(question_text) # Call your generate_response function here
question_sequence = state["question"]
correct_response_sequence = state["correct_response"]
incorrect_response_sequence = state["incorrect_response"]
question_length = len(question_sequence)
correct_response_length = len(correct_response_sequence)
incorrect_response_length = len(incorrect_response_sequence)
# input_size = max_question_length + max_correct_response_length + max_incorrect_response_length

question = torch.nn.functional.pad(question, (0, input_size - question.size(2)))
correct_response = torch.nn.functional.pad(correct_response, (0, input_size - correct_response.size(2)))
incorrect_response = torch.nn.functional.pad(incorrect_response, (0, input_size - incorrect_response.size(2)))
# Perform actions, get rewards, and transition to the next state
action = policy.select_action({'question' : question ,
'correct_response' : correct_response,
'incorrect_response' : incorrect_response})
# Assuming you have a reward function that calculates reward given the action and state
reward = calculate_reward(action, state)
next_state, _, done = env.step(action)
episode_states.append(state)
episode_actions.append(action)
episode_rewards.append(reward)
state = next_state
if done:
break

# Compute discounted returns
discounted_returns = []
running_add = 0
for r in reversed(episode_rewards):
running_add = r + 0.99 * running_add # Discount factor: 0.99
discounted_returns.insert(0, running_add)

# Calculate loss and perform policy gradient update
action_probs = policy(torch.cat(episode_states), response_text)


selected_action_probs = torch.gather(action_probs, 1, torch.tensor(episode_actions).unsqueeze(1))
loss = -torch.sum(torch.log(selected_action_probs) * torch.FloatTensor(discounted_returns))
optimizer.zero_grad()
loss.backward()
optimizer.step()

# Print episode information
if episode % 10 == 0:
print(f"Episode {episode}, Total Reward: {np.sum(episode_rewards)}")

# Use the trained policy for inference
with torch.no_grad():
test_state = env.reset()
while True:
# Extract the question from the state (assuming state contains the question)
question_text = test_state["question"]
question_text = "".join([tokenizer.token_to_character(token) for token in question_text])

# Generate a response based on the input question_text
generate_response(question_text) # Call your generate_response function here

test_state = torch.tensor([test_state], dtype=torch.float32)
action_probs = policy(test_state)
action = np.argmax(action_probs.detach().numpy()[0])
next_state, _, done = env.step(action)
test_state = next_state
if done:
break