add custom experimenter

expo/MCTS.py

@@ -205,6 +205,8 @@ class Node():
         if self.state["low_is_better"]:
             # normalized the score to be between 0 and 1, and higher is better
             def normalize_score(score):
+                if score == -1:
+                    return 0
                 return 1 / (1 + score)
             score_dict = {k: normalize_score(v) for k, v in score_dict.items()}
         self.normalized_reward = score_dict

expo/README.md

@@ -25,8 +25,9 @@ ### Run DI RandExp
 ### Run DI MCTS
 `python run_experiment.py --exp_mode mcts --task titanic --rollout 5`
 
-If the dataset has reg metric, remember to use `--low_is_better`
-`python run_experiment.py --exp_mode mcts --task househouse_prices --rollout 5 --low_is_better`
+If the dataset has reg metric, remember to use `--low_is_better`:
+
+- `python run_experiment.py --exp_mode mcts --task househouse_prices --rollout 5 --low_is_better`
 
 
 ## Code and Configs Explanation

expo/dataset.py

@@ -27,7 +27,7 @@ TASK_PROMPT = """\
 - Both files should contain a single column named `target` with the predicted values.
 2. Make sure the prediction results are in the same format as the target column in the training set. 
 - The labels should be transformed back to the original format if any transformation was applied during training.
-- If the original target column was categorical, the predictions should be in the same format.
+- If the original target column was categorical or string, the predictions MUST be in the same format.
 
 ## Output Training Set Performance
 Make sure the performance of the model is printed in python in the last step even if it has been printed in the previous steps. The value should be a float number.

expo/experimenter/__init__.py

@@ -1,3 +1,4 @@
 from .experimenter import Experimenter
 from .mcts import MCTSExperimenter
-from .aug import AugExperimenter
+from .aug import AugExperimenter
+from .custom import CustomExperimenter

expo/experimenter/custom.py

@@ -0,0 +1,49 @@
+from expo.experimenter import Experimenter
+from expo.MCTS import create_initial_state
+from expo.evaluation.evaluation import evaluate_score
+import pandas as pd
+import os
+
+class CustomExperimenter(Experimenter):
+    result_path : str = "results/custom"
+    
+    def __init__(self, args, **kwargs):
+        super().__init__(args, **kwargs)
+        self.framework = kwargs["framework"]
+        self.name = kwargs.get("name", "")
+        self.result_path = f"results/custom_{self.name}"
+        self.state = create_initial_state(self.args.task, start_task_id=1, data_config=self.data_config, low_is_better=self.args.low_is_better, name="")
+    
+    async def run_experiment(self):
+        user_requirement = self.state["requirement"]
+        preds = await self.framework.run(user_requirement)
+        test_preds = preds["test_preds"]
+        dev_preds = preds["dev_preds"]
+        score_dict = {
+            "dev_score": self.evaluate_predictions(dev_preds, "dev"),
+            "test_score": self.evaluate_predictions(test_preds, "test")
+        }
+        results = {
+            "score_dict": score_dict,
+            "user_requirement": user_requirement,
+            "args": vars(self.args)
+        }
+        self.save_result(results)
+
+    def evaluate_predictions(self, preds, split):
+        metric = self.state["dataset_config"]["metric"]
+        gt_path = os.path.join(self.state["datasets_dir"][f"{split}_target"])
+        gt = pd.read_csv(gt_path)["target"]
+        score = evaluate_score(preds, gt, metric)
+        return score        
+    
+
+    def load_datasets(self):
+        train_path = self.state["datasets_dir"]["train"]
+        dev_path = self.state["datasets_dir"]["dev"]
+        test_path = self.state["datasets_dir"]["test"]
+        train = pd.read_csv(train_path)
+        dev = pd.read_csv(dev_path)
+        test = pd.read_csv(test_path)
+        return train, dev, test
+    

expo/run_experiment.py

@@ -1,4 +1,4 @@
-from expo.experimenter import MCTSExperimenter, Experimenter, AugExperimenter
+from expo.experimenter import MCTSExperimenter, Experimenter, AugExperimenter, CustomExperimenter
 import asyncio
 import argparse
 
@@ -6,7 +6,7 @@ import argparse
 def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--name", type=str, default="")
-    parser.add_argument("--exp_mode", type=str, default="mcts", choices=["mcts", "aug", "base"])
+    parser.add_argument("--exp_mode", type=str, default="mcts", choices=["mcts", "aug", "base", "custom"])
     get_di_args(parser)
     get_mcts_args(parser)
     get_aug_exp_args(parser)
@@ -17,7 +17,7 @@ def get_mcts_args(parser):
     parser.add_argument("--load_tree", dest="load_tree", action="store_true")
     parser.add_argument("--no_load_tree", dest="load_tree", action="store_false")
     parser.set_defaults(load_tree=False)
-    parser.add_argument("--rollouts", type=int, default=3)
+    parser.add_argument("--rollouts", type=int, default=5)
 
 def get_aug_exp_args(parser):
     parser.add_argument("--aug_mode", type=str, default="single", choices=["single", "set"])
@@ -40,6 +40,8 @@ async def main(args):
         experimenter = AugExperimenter(args)
     elif args.exp_mode == "base":
         experimenter = Experimenter(args)
+    elif args.exp_mode == "custom":
+        experimenter = CustomExperimenter(args)
     else:
         raise ValueError(f"Invalid exp_mode: {args.exp_mode}")
     await experimenter.run_experiment()