Spaces:
Sleeping
Sleeping
| import os | |
| import sys | |
| import torch | |
| import torch.nn as nn | |
| import torch.optim as optim | |
| from torch.distributions import Categorical | |
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) | |
| from code_review_env import CodeReviewAction | |
| from code_review_env.server.environment import CodeReviewEnvironment | |
| class PolicyNetwork(nn.Module): | |
| def __init__(self, input_dim, hidden_dim, output_dim): | |
| super(PolicyNetwork, self).__init__() | |
| self.fc1 = nn.Linear(input_dim, hidden_dim) | |
| self.fc2 = nn.Linear(hidden_dim, hidden_dim) | |
| self.fc3 = nn.Linear(hidden_dim, output_dim) | |
| self.relu = nn.ReLU() | |
| def forward(self, x): | |
| x = self.relu(self.fc1(x)) | |
| x = self.relu(self.fc2(x)) | |
| x = self.fc3(x) | |
| return torch.softmax(x, dim=-1) | |
| def extract_state(obs): | |
| state = [ | |
| obs.churn_score / 100.0, | |
| obs.complexity_score / 100.0, | |
| obs.todo_score / 100.0, | |
| obs.recency_score / 100.0, | |
| obs.files_flagged / max(1.0, float(obs.review_budget)), | |
| obs.review_budget / 20.0 | |
| ] | |
| return torch.FloatTensor(state).unsqueeze(0) | |
| def main(): | |
| input_dim = 6 | |
| hidden_dim = 32 | |
| output_dim = 2 | |
| learning_rate = 0.01 | |
| policy_net = PolicyNetwork(input_dim, hidden_dim, output_dim) | |
| optimizer = optim.Adam(policy_net.parameters(), lr=learning_rate) | |
| num_episodes = 50 | |
| print(f"training for {num_episodes} episodes") | |
| env = CodeReviewEnvironment() | |
| import random | |
| difficulties = ["easy", "medium", "hard"] | |
| for episode in range(num_episodes): | |
| difficulty = random.choice(difficulties) | |
| obs = env.reset(difficulty=difficulty) | |
| saved_log_probs = [] | |
| rewards = [] | |
| while not obs.done: | |
| state = extract_state(obs) | |
| probs = policy_net(state) | |
| m = Categorical(probs) | |
| action = m.sample() | |
| saved_log_probs.append(m.log_prob(action)) | |
| decision = "flag" if action.item() == 1 else "skip" | |
| obs = env.step(CodeReviewAction(decision=decision)) | |
| rewards.append(obs.reward) | |
| gamma = 0.99 | |
| R = 0 | |
| policy_loss = [] | |
| returns = [] | |
| for r in rewards[::-1]: | |
| R = r + gamma * R | |
| returns.insert(0, R) | |
| returns = torch.tensor(returns) | |
| if len(returns) > 1 and returns.std() > 0: | |
| returns = (returns - returns.mean()) / (returns.std() + 1e-8) | |
| for log_prob, R_val in zip(saved_log_probs, returns): | |
| policy_loss.append(-log_prob * R_val) | |
| optimizer.zero_grad() | |
| if policy_loss: | |
| policy_loss = torch.cat(policy_loss).sum() | |
| policy_loss.backward() | |
| optimizer.step() | |
| if (episode + 1) % 10 == 0 or episode == 0: | |
| print(f"Episode {episode + 1}/{num_episodes} | " | |
| f"Repo: {obs.repo_name} | " | |
| f"Loss: {policy_loss.item() if policy_loss else 0:.2f} | " | |
| f"Reward: {sum(rewards):.1f} | " | |
| f"F1: {obs.f1_score:.2f} (P: {obs.precision:.2f}, R: {obs.recall:.2f})") | |
| print("\ndone.") | |
| if __name__ == "__main__": | |
| main() | |