Merge pull request #9 from VectorInstitute/develop

Implement code skeleton for the entire new capability generation pipeline

Author: xeon27

pyproject.toml

@@ -125,6 +125,8 @@ ignore = [
 [tool.ruff.lint.per-file-ignores]
 "__init__.py" = ["E402", "F401", "F403", "F811"]
 "tests/src/seed_capabilities/math/math_competition_algebra/capability.py" = ["D100", "D101", "D102"]
+"src/run.py" = ["ERA001"]
+"src/lbo.py" = ["ERA001"]
 
 [tool.ruff.lint.pep8-naming]
 ignore-names = ["X*", "setUp"]

src/capability.py

@@ -3,7 +3,7 @@
 import os
 import sys
 from collections import defaultdict
-from typing import Any, Dict
+from typing import Any, Dict, List
 
 from src.model import Model
 from src.utils.capability_utils import parse_python_class_str, read_score_inspect_json
@@ -252,10 +252,38 @@ def encode(self, encoder_model: Any) -> None:
         self.encoding = None
         raise NotImplementedError
 
-    def evaluate_using_inspect(self, model: Model) -> None:  # noqa: D102
-        # evaluate the capability using inspect-evals
+    def _create_inspect_file(self) -> None:
+        """
+        Implement pipeline to evaluate the capability using the inspect framework.
+
+        This involves converting the METR format to inspect solvers and scorers.
+        """
+        raise NotImplementedError
+
+    def _evaluate_using_inspect(self, subject_llm: Model) -> None:  # noqa: D102
+        """
+        Evaluate subject LLM on the capability using the inspect framework.
+
+        Args
+        ----
+        subject_llm : Model
+            The LLM to use for evaluation.
+        """
         raise NotImplementedError
 
+    def evaluate(self, subject_llms: List[Model]) -> None:
+        """
+        Evaluate the provided subject LLMs on the capability.
+
+        Args
+        ----
+        subject_llms : List[Model]
+            The list of LLMs to use for evaluation.
+        """
+        # TODO: Run asynchronosly
+        for model in subject_llms:
+            self._evaluate_using_inspect(model)
+
 
 def _import_from_path(module_name: str, file_path: str) -> Any:
     """

src/cfg/run_cfg.yaml

@@ -1,10 +1,10 @@
-generator_model:
+scientist_llm:
   name: gpt-4o-mini
   gen_cfg:
     temperature: 0.7
     max_tokens: 64
 
-candidate_model:
+subject_llm:
   name: Meta-Llama-3.1-70B-Instruct
 
 prompt_cfg:
@@ -17,6 +17,16 @@ capabilities_cfg:
   num_seed_capabilities: -1
   num_gen_capabilities: 4
   num_gen_capabilities_per_run: 2
+  num_gen_tasks_per_capability: 2
+
+lbo_cfg:
+  # Number of capabilities to generate using LBO
+  num_lbo_runs: 1
+  # Type of LBO pipeline to use
+  pipeline_id: "nearest_neighbor" # "nearest_neighbor" or "discover_new"
+  # Train args for 'nearest_neighbor' pipeline
+  train_frac: 0.5
+  min_train_size: 10
 
 exp_cfg:
   # Set this flag to true to run test experiments during development

src/generate_capabilities.py

@@ -18,7 +18,7 @@
 def _sample_seed_capabilities(
     seed_capability_dir: str,
     num_seed_capabilities: int = -1,
-    include_capabilities: List[str] | None = None,
+    include_capability_names: List[str] | None = None,
     random_seed: int = 42,
 ) -> List[Capability]:
     """
@@ -31,7 +31,8 @@ def _sample_seed_capabilities(
     ----
         seed_capability_dir (str): The directory containing the seed capabilities.
         num_seed_capabilities (int): The number of seed capabilities to sample.
-        include_capabilities (List[str] | None): A list of capability names to include.
+        include_capability_names (List[str] | None): A list of
+            capability names to include.
         random_seed (int): The seed for the random number generator.
 
     Returns
@@ -46,21 +47,21 @@ def _sample_seed_capabilities(
     # Select all capabilities if num_seed_capabilities is -1
     if num_seed_capabilities == -1:
         num_seed_capabilities = len(all_seed_capability_paths)
-        include_capabilities = None
+        include_capability_names = None
 
     # Force include some capabilities
-    if include_capabilities is not None:
-        assert num_seed_capabilities >= len(include_capabilities), (
+    if include_capability_names is not None:
+        assert num_seed_capabilities >= len(include_capability_names), (
             "Number of seed capabilities is less than the number of capabilities to include."
         )
-        for capability_name in include_capabilities:
+        for capability_name in include_capability_names:
             assert os.path.exists(os.path.join(seed_capability_dir, capability_name)), (
                 f"{capability_name} does not exist in {seed_capability_dir}."
             )
             capability = Capability(os.path.join(seed_capability_dir, capability_name))
             sampled_seed_capabilities.append(capability)
             all_seed_capability_paths.remove(capability_name)
-        num_seed_capabilities -= len(include_capabilities)
+        num_seed_capabilities -= len(include_capability_names)
 
     # TODO: Enhance the selection criterion
     for capability_path in random.sample(
@@ -121,10 +122,10 @@ def generate_capabilities_using_llm(
     sys_prompt: str,
     user_prompt: str,
     num_seed_capabilities: int,
-    prev_capabilities: List[str],
+    prev_capabilities: List[Capability],
     scientist_llm_gen_cfg: Dict[str, Any],
     base_capability_dir: str,
-    include_seed_capabilities: Optional[List[str]] = None,
+    include_seed_capability_names: Optional[List[str]] = None,
     **kwargs: Any,
 ) -> Dict[str, Any]:
     """
@@ -142,25 +143,27 @@ def generate_capabilities_using_llm(
         sys_prompt (str): The system prompt.
         user_prompt (str): The user prompt.
         num_seed_capabilities (int): The number of seed capabilities to use.
-        prev_capabilities (List[str]): The list of previously
-            generated capability names.
+        prev_capabilities (List[Capability]): The list of previously
+            generated capabilities.
         scientist_llm_gen_cfg (Dict[str, Any]): The generation configuration
             for the scientist LLM.
         base_capability_dir (str): The base directory to store
             the generated capabilities for the specified domain.
-        include_seed_capabilities (List[str] | None): A list of seed capability
+        include_seed_capability_names (List[str] | None): A list of seed capability
             names to include in the generation process.
+        **kwargs (Any): Additional keyword arguments.
 
     Returns
     -------
-        List[str]: The generated capability names.
+        Dict[str, Any]: A dictionary containing the generated capabilities
+        and metadata about the generation process.
     """
     # Select seed capabilities
     seed_capability_dir = os.path.join(BASE_ARTIFACTS_DIR, "seed_capabilities", domain)
     seed_capabilities = _sample_seed_capabilities(
         seed_capability_dir=seed_capability_dir,
         num_seed_capabilities=num_seed_capabilities,
-        include_capabilities=include_seed_capabilities,
+        include_capability_names=include_seed_capability_names,
     )
     # Get capability JSON strings (without scores)
     seed_capabilities_repr = [
@@ -170,7 +173,7 @@ def generate_capabilities_using_llm(
     # LLM input
     user_prompt = user_prompt.format(
         seed_capabilities="\n".join(seed_capabilities_repr),
-        prev_capabilities="\n".join(prev_capabilities),
+        prev_capabilities="\n".join([elm.name for elm in prev_capabilities]),
         domain=domain,
         num_gen_capabilities=num_capabilities,
     )
@@ -193,10 +196,9 @@ def generate_capabilities_using_llm(
         Capability.from_dict(capability_dict=capability, base_dir=base_capability_dir)
         for capability in gen_capabilities
     ]
-    gen_capabilities_names = [elm.name for elm in gen_capabilities]
 
     return {
-        "capabilities": gen_capabilities_names,
+        "capabilities": gen_capabilities,
         "metadata": {
             "model": scientist_llm.get_model_name(),
             "thought": parsed_response["thought"],
@@ -206,20 +208,20 @@ def generate_capabilities_using_llm(
 
 
 def filter_capabilities(
-    capabilities: List[str],
-) -> List[str]:
+    capabilities: List[Capability],
+) -> List[Capability]:
     """
     Filter capabilities based on multiple criterion.
 
     Remove repeated, irrelevant, and ill-formed capabilities.
 
     Args
     ----
-        capabilities (List[str]): The list of capabilities.
+        capabilities (List[Capability]): The list of capabilities.
 
     Returns
     -------
-        List[str]: The filtered capability names.
+        List[Capability]: The list of remaining capabilities.
     """
     # TODO: Implement capability filtering
     return capabilities
@@ -232,9 +234,9 @@ def generate_capabilities(
     scientist_llm: Model,
     num_seed_capabilities: int,
     scientist_llm_gen_cfg: Dict[str, Any],
-    include_seed_capabilities: Optional[List[str]] = None,
+    include_seed_capability_names: Optional[List[str]] = None,
     **kwargs: Any,
-) -> List[str]:
+) -> List[Capability]:
     """
     Generate initial capabilities for the specified domain.
 
@@ -247,12 +249,12 @@ def generate_capabilities(
         num_seed_capabilities (int): The number of seed capabilities to use.
         scientist_llm_gen_cfg (Dict[str, Any]): The generation configuration
             for the scientist LLM.
-        include_seed_capabilities (List[str] | None): A list of seed capability
+        include_seed_capability_names (List[str] | None): A list of seed capability
             names to include in the generation process.
 
     Returns
     -------
-        List[str]: The generated capability names.
+        List[Capability]: The generated capabilities.
     """
     num_runs = int(np.ceil(num_capabilities / num_capabilities_per_run))
     gen_capabilities = []
@@ -268,10 +270,7 @@ def generate_capabilities(
         base_capability_dir = os.path.join(BASE_ARTIFACTS_DIR, "capabilities", domain)
 
     # Fetch previously generated capabilities, if any
-    prev_capabilities = [
-        elm.name
-        for elm in _get_previous_capabilities(capability_dir=base_capability_dir)
-    ]
+    prev_capabilities = _get_previous_capabilities(capability_dir=base_capability_dir)
 
     for run_id in range(num_runs):
         print("Run ID:", run_id)
@@ -286,7 +285,7 @@ def generate_capabilities(
             prev_capabilities=prev_capabilities,
             scientist_llm_gen_cfg=scientist_llm_gen_cfg,
             base_capability_dir=base_capability_dir,
-            include_seed_capabilities=include_seed_capabilities,
+            include_seed_capability_names=include_seed_capability_names,
             **kwargs,
         )
         gen_capabilities.extend(response["capabilities"])

src/generate_tasks.py

@@ -0,0 +1,58 @@
+from typing import Any, Dict  # noqa: D100
+
+from capability import Capability
+from model import Model
+
+
+def generate_tasks_using_llm(
+    capability: Capability,
+    scientist_llm: Model,
+    sys_prompt: str,
+    user_prompt: str,
+    num_tasks: int,
+    scientist_llm_gen_cfg: Dict[str, Any],
+) -> None:
+    """
+    Generate `num_tasks` tasks for the given capability.
+
+    Generate tasks for the given capability
+    using the scientist LLM model based on the following approach:
+    <Approach>
+
+    Args
+    ----
+        capability (Capability): The capability to generate tasks for.
+        scientist_llm (Model): The scientist LLM model.
+        sys_prompt (str): The system prompt for generating tasks.
+        user_prompt (str): The user prompt for generating tasks.
+        num_tasks (int): The number of tasks to generate.
+        scientist_llm_gen_cfg (Dict[str, Any]): The generation configuration
+            for the scientist LLM.
+    """
+    # TODO: Implement the function with the following components
+    # # Approach 1
+    # 1. First generate task questions. This can be done in two ways:
+    #   a. Single run to generate all `num_tasks` (Nt) questions
+    #       - input tokens: Pt
+    #       - output tokens: Nt * Qt, where Qt is the mean # tokens in a question
+    #   b. Multiple runs to generate `num_tasks` (Nt)
+    #      questions in batches of `num_tasks_per_run` (Ntr)
+    #       - input tokens: Pt * Ntr
+    #       - output tokens: Nt * Qt
+    # 2. Filter out similar/ill-formatted questions
+    # 3. Then obtain task answers by:
+    #   a. prompting the scientist LLM to solve these selected questions
+    #   b. using a group of (less capable) models to solve
+    #      these questions and then selecting the majority answer
+    #   c. using a scoring function
+    #
+    # # Approach 2
+    # 1. Generate task questions and answers together in a single run.
+    #    Again, this can be done in two ways described above.
+    # 2. Filter out similar/ill-formatted question/asnwer pairs
+    # 3. Verify each pair by:
+    #   a. prompting the scientist LLM to function as a judge
+    #   b. using a group of (less capable) models to judge and
+    #      then selecting the majority answer
+
+    raise NotImplementedError

src/get_seed_capability_results.py

@@ -105,7 +105,7 @@ def main(cfg: DictConfig) -> None:
     3. Reads the capability configuration from the "capability.json" file.
     4. Determines the dataset name and capability details
     from the capability configuration.
-    5. Iterates over results for all candidate models
+    5. Iterates over results for all subject models
     in the seed datasets log directory.
     6. For each log file that matches the dataset name,
     processes the log file based on the dataset type:
@@ -137,23 +137,23 @@ def main(cfg: DictConfig) -> None:
         if dataset_name == "math":
             subject = capability_json["capability_subject"]
 
-        # Iterate over results for all candidate models
-        for candidate_model_dir in os.listdir(seed_datasets_log_dir):
-            candidate_model_log_path = os.path.join(
-                seed_datasets_log_dir, candidate_model_dir
+        # Iterate over results for all subject models
+        for subject_model_dir in os.listdir(seed_datasets_log_dir):
+            subject_model_log_path = os.path.join(
+                seed_datasets_log_dir, subject_model_dir
             )
-            for log_file in os.listdir(candidate_model_log_path):
+            for log_file in os.listdir(subject_model_log_path):
                 if dataset_name not in log_file:
                     continue
 
-                out_dir = os.path.join(seed_capability_result_dir, candidate_model_dir)
+                out_dir = os.path.join(seed_capability_result_dir, subject_model_dir)
                 out_dir = os.path.join(out_dir, domain)
                 os.makedirs(out_dir, exist_ok=True)
 
                 # For math dataset, extract math capability logs
                 if "math" in log_file:
                     extract_math_capability_logs(
-                        log_file=os.path.join(candidate_model_log_path, log_file),
+                        log_file=os.path.join(subject_model_log_path, log_file),
                         capability_name=capability_name,
                         subject=subject,
                         out_dir=out_dir,
@@ -163,15 +163,15 @@ def main(cfg: DictConfig) -> None:
                 elif "gsm8k" in log_file:
                     # No changes to log file, just copy it to output directory
                     shutil.copyfile(
-                        src=os.path.join(candidate_model_log_path, log_file),
+                        src=os.path.join(subject_model_log_path, log_file),
                         dst=os.path.join(
                             out_dir,
                             f"{capability_name}.json",
                         ),
                     )
 
                 print(
-                    f"Extracted {candidate_model_dir} result for {capability_name} capability."
+                    f"Extracted {subject_model_dir} result for {capability_name} capability."
                 )