update prompt (specify whether each set has target label)

expo/MCTS.py

@@ -240,7 +240,7 @@ class MCTS():
         all_children = [child for children in self.children.values() for child in children]
         return max(all_children, key=uct)
 
-    async def expand(self, node : Node, max_children=4):
+    async def expand(self, node : Node, max_children=5):
         await node.expand(max_children)
         if node not in self.children or not self.children[node]:
             self.children[node] = node.children
@@ -273,7 +273,7 @@ class MCTS():
                 return best_score, best_child
             for child in self.children[node]:
                 score = child.normalized_reward[split]
-                print(child.id, score)
+                print(child.id, split, score)
                 if score > best_score:
                     best_score = score
                     best_child = child

expo/dataset.py

@@ -20,7 +20,7 @@ TASK_PROMPT = """\
 **Attention** 
 1. Please do not leak the target label in any form during training.
 2. Dev and Test sets do not have the target column.
-3. You should perform transformations on all sets at the same step.
+3. You should perform transformations on train, dev, and test sets at the same time (it's a good idea to define functions for this and avoid code repetition).
 4. If labels are transformed during training, they should be transformed back to the original format before saving the predictions.
 
 ## Saving Dev and Test Predictions
@@ -38,9 +38,9 @@ print("Train score:", train_score)
 ```
 
 # Data dir
-training: {train_path}
-dev: {dev_path}
-testing: {test_path}
+training (with labels): {train_path}
+dev (without labels): {dev_path}
+testing (without labels): {test_path}
 
 # Output dir
 {output_dir}

expo/evaluation/evaluation.py

@@ -5,7 +5,7 @@ def evaluate_score(pred, gt, metric):
     if metric == "accuracy":
         return accuracy_score(gt, pred)
     elif metric == "f1":
-        unique_classes = np.unique(gt)
+        unique_classes = sorted(list(np.unique(gt)))
         if 1 in unique_classes and 0 in unique_classes:
             pos_label = 1
         else:

expo/evaluation/visualize_mcts.py

@@ -48,7 +48,9 @@ def get_tree_text(node : Node):
             for child in node.children:
                 text += textwrap.indent(visualize_tree(child, depth+1, previous_plans), "\t")
         return text
-    
-    return visualize_tree(node), len(code_set)
+    num_simulations = node.visited
+    text = f"Number of simulations: {num_simulations}\n"
+    text += visualize_tree(node)
+    return text, len(code_set)
 
 

expo/experimenter/mcts.py

@@ -22,19 +22,18 @@ class MCTSExperimenter(Experimenter):
         text += f"Best node: {best_node}, score: {best_node.raw_reward}\n"
         text += f"Dev best node: {dev_best_node}, score: {dev_best_node.raw_reward}\n"
         print(text)
-        if self.args.rollouts > 0:
-            self.save_tree(text)
+        self.save_tree(text)
 
-            results = {
-                "best_node": best_node.id,
-                "best_node_score": best_node.raw_reward,
-                "dev_best_node": dev_best_node.id,
-                "dev_best_node_score": dev_best_node.raw_reward,
-                "num_generated_codes": num_generated_codes,
-                "user_requirement": best_node.state["requirement"],
-                "args": vars(self.args)
-            }
-            self.save_result(results)
+        results = {
+            "best_node": best_node.id,
+            "best_node_score": best_node.raw_reward,
+            "dev_best_node": dev_best_node.id,
+            "dev_best_node_score": dev_best_node.raw_reward,
+            "num_generated_codes": num_generated_codes,
+            "user_requirement": best_node.state["requirement"],
+            "args": vars(self.args)
+        }
+        self.save_result(results)