Add modification of CePO configs through yaml and cli arguments

This will allow users to do the following:
1. Cli where if they can pass in anything that's "cepo_<name-of-attribute>"
2. Yaml file where if they pass it in as "<name-of-attribute>"
3. If none of them have a specific attribute, we use the default setting
4. If both of them have the specific attribute, we error out

Author: erich-cerebras

configs/cepo_config.yaml

@@ -0,0 +1,14 @@
+bestofn_n: 4
+bestofn_temperature: 0.2324214
+bestofn_max_tokens: 4096
+bestofn_rating_type: "absolute"  # or "pairwise"
+planning_n: 2
+planning_m: 2
+planning_temperature_step1: 0.00055
+planning_temperature_step2: 0.25
+planning_temperature_step3: 0.999
+planning_temperature_step4: 1.2
+planning_max_tokens_step1: 96
+planning_max_tokens_step2: 2
+planning_max_tokens_step3: 8000
+planning_max_tokens_step4: 0

optillm.py

@@ -28,7 +28,7 @@
 from optillm.plansearch import plansearch
 from optillm.leap import leap
 from optillm.reread import re2_approach
-from optillm.cepo import cepo
+from optillm.cepo import cepo, CepoConfig, init_cepo_config
 
 # Setup logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -292,7 +292,9 @@ def execute_single_approach(approach, system_prompt, initial_query, client, mode
         elif approach == 're2':
             return re2_approach(system_prompt, initial_query, client, model, n=server_config['n'])
         elif approach == 'cepo':
-            return cepo(system_prompt, initial_query, client, model)
+            # build the cepo config based on the cmd line arguments and the 
+            logger.debug(f"Calling with {cepo_config}")
+            return cepo(system_prompt, initial_query, client, model, cepo_config)
     elif approach in plugin_approaches:
         return plugin_approaches[approach](system_prompt, initial_query, client, model)
     else:
@@ -701,6 +703,13 @@ def parse_args():
     parser.add_argument("--base-url", "--base_url", dest="base_url", type=str, default=base_url_default,
                         help="Base url for OpenAI compatible endpoint")
 
+    # Special handling of all the Cepo Configurations
+    for key, value in CepoConfig.__dict__.items():
+        if not key.startswith('__'):
+            parser.add_argument(f"--cepo_{key}", dest=f"cepo_{key}", type=type(value), default=None, help=f"CePO configuration for {key}")
+
+    parser.add_argument(f"--cepo_config_file", dest=f"cepo_config_file", type=str, default=None, help="Path to CePO configuration file")
+
     args = parser.parse_args()
 
     # Convert argument names to match server_config keys
@@ -714,6 +723,7 @@ def parse_args():
 
 def main():
     global server_config
+    global cepo_config
     # Call this function at the start of main()
     args = parse_args()
     # Update server_config with all argument values
@@ -728,6 +738,11 @@ def main():
     if logging_level in logging_levels.keys():
         logger.setLevel(logging_levels[logging_level])
 
+    # set and log the cepo configs
+    cepo_config = init_cepo_config(server_config)
+    if args.approach == 'cepo':
+        logger.info(f"CePO Config: {cepo_config}")
+    
     logger.info(f"Starting server with approach: {server_config['approach']}")
     server_config_clean = server_config.copy()
     if server_config_clean['optillm_api_key']:

optillm/cepo.py

@@ -3,27 +3,65 @@
 import openai
 
 from dataclasses import dataclass
-from typing import Optional
+from typing import Optional, Literal
 
+import yaml
 
 @dataclass
 class CepoConfig:
     bestofn_n: int = 3
-    bestofn_temperature: int = 0.1
+    bestofn_temperature: float = 0.1
     bestofn_max_tokens: int = 4096
-    bestofn_rating_type: str = "absolute"  # or "pairwise"
+    bestofn_rating_type: Literal["absolute", "pairwise"] = "absolute"
     planning_n: int = 3
     planning_m: int = 6
-    planning_temperature_step1: int = 0.55
-    planning_temperature_step2: int = 0.25
-    planning_temperature_step3: int = 0.1
-    planning_temperature_step4: int = 0
+    planning_temperature_step1: float = 0.55
+    planning_temperature_step2: float = 0.25
+    planning_temperature_step3: float = 0.1
+    planning_temperature_step4: float = 0
     planning_max_tokens_step1: int = 4096
     planning_max_tokens_step2: int = 4096
     planning_max_tokens_step3: int = 4096
     planning_max_tokens_step4: int = 4096
 
 
+# given command line arguments which includes a yaml file path, initialize a CePO configuration
+def init_cepo_config(cmd_line_args: dict) -> CepoConfig:
+    # get the command line arguments
+    cepo_args = {
+        key.split("cepo_")[1]: value
+        for key, value in cmd_line_args.items()
+        if "cepo" in key and "cepo_config_file" != key and value is not None
+    }
+
+    # get the yaml file arguments
+    cepo_config_yaml = {}
+    if "cepo_config_file" in cmd_line_args.keys():
+        with open(cmd_line_args["cepo_config_file"], "r") as yaml_file:
+            cepo_config_yaml = yaml.safe_load(yaml_file)
+
+    # check if any of the keys overlap, and if they do, error out
+    for key in cepo_config_yaml.keys():
+        if key in cepo_args.keys():
+            raise RuntimeError(f"Key {key} is found in both yaml file and command line arguments")
+
+    # if not, then we take both of them and add them to the cepo config
+    cepo_config = CepoConfig()
+    cepo_attrs = [key for key, _ in cepo_config.__dict__.items() if not key.startswith('__')]
+
+    # add command line arguments
+    for key, value in cepo_args.items():
+        # this assert should not be raised as the cli parser should catch this
+        assert key in cepo_attrs, f"Command line argument {key} is not found in CepoConfig"
+        setattr(cepo_config, key, value)
+
+    # add yaml arguments
+    for key, value in cepo_config_yaml.items():
+        assert key in cepo_attrs, f"Yaml argument {key} is not found in CepoConfig"
+        setattr(cepo_config, key, value)
+
+    return cepo_config
+
 def extract_question_only(task: str) -> str:
     """We noticed that sometimes if the task includes specific formatting instructions, they may interfere with the reasoning flow. This
     is a temporary workaround to extract the question only from the task. Work in progress.