feat: added optimizer n evolutions (#126)

* feat: added optimizer n evolutions

* refactor: effective way to generate evolutions

* refactor: effective way to generate evolutions

* feat: added search space

* refactor

* refactor

* feat: added tests and fix mypy

* feat: IncrementalUtteranceEvolver

* feat: deleted embedder config

* fix: mypy

* fix: import

* feat: updated cli and search space

* feat: updated tests

* feat: updated __init__

* fix: mypy

Author: Darinochka

autointent/generation/utterances/__init__.py

@@ -6,6 +6,7 @@
     FormalEvolution,
     FunnyEvolution,
     GoofyEvolution,
+    IncrementalUtteranceEvolver,
     InformalEvolution,
     ReasoningEvolution,
     UtteranceEvolver,
@@ -20,6 +21,7 @@
     "FunnyEvolution",
     "Generator",
     "GoofyEvolution",
+    "IncrementalUtteranceEvolver",
     "InformalEvolution",
     "ReasoningEvolution",
     "SynthesizerChatTemplate",

autointent/generation/utterances/evolution/__init__.py

@@ -9,6 +9,7 @@
     ReasoningEvolution,
 )
 from .evolver import UtteranceEvolver
+from .incremental_evolver import IncrementalUtteranceEvolver
 
 __all__ = [
     "AbstractEvolution",
@@ -17,6 +18,7 @@
     "FormalEvolution",
     "FunnyEvolution",
     "GoofyEvolution",
+    "IncrementalUtteranceEvolver",
     "InformalEvolution",
     "ReasoningEvolution",
     "UtteranceEvolver",

autointent/generation/utterances/evolution/cli.py

@@ -1,16 +1,15 @@
 """CLI for evolutionary augmenter."""
 
 import logging
-from argparse import ArgumentParser
+from argparse import ArgumentParser, Namespace
 
 from autointent import load_dataset
-from autointent.generation.utterances.evolution.evolver import UtteranceEvolver
+from autointent.generation.utterances.evolution import IncrementalUtteranceEvolver, UtteranceEvolver
 from autointent.generation.utterances.generator import Generator
 
 from .chat_templates import (
     AbstractEvolution,
     ConcreteEvolution,
-    EvolutionChatTemplate,
     FormalEvolution,
     FunnyEvolution,
     GoofyEvolution,
@@ -22,8 +21,7 @@
 logger = logging.getLogger(__name__)
 
 
-def main() -> None:
-    """CLI endpoint."""
+def _parse_args() -> Namespace:
     parser = ArgumentParser()
     parser.add_argument(
         "--input-path",
@@ -46,6 +44,7 @@ def main() -> None:
     )
     parser.add_argument("--private", action="store_true", help="Publish privately if --output-repo option is used")
     parser.add_argument("--n-evolutions", type=int, default=1, help="Number of utterances to generate for each intent")
+    parser.add_argument("--decide-for-me", action="store_true")
     parser.add_argument("--reasoning", action="store_true", help="Whether to use `Reasoning` evolution")
     parser.add_argument("--concretizing", action="store_true", help="Whether to use `Concretizing` evolution")
     parser.add_argument("--abstract", action="store_true", help="Whether to use `Abstract` evolution")
@@ -55,34 +54,46 @@ def main() -> None:
     parser.add_argument("--informal", action="store_true", help="Whether to use `Informal` evolution")
     parser.add_argument("--async-mode", action="store_true", help="Enable asynchronous generation")
     parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--batch-size", type=int, default=4)
+    parser.add_argument("--search-space", type=str, default=None)
+
+    return parser.parse_args()
+
 
-    args = parser.parse_args()
-
-    evolutions: list[EvolutionChatTemplate] = []
-    if args.reasoning:
-        evolutions.append(ReasoningEvolution())
-    if args.concretizing:
-        evolutions.append(ConcreteEvolution())
-    if args.abstract:
-        evolutions.append(AbstractEvolution())
-    if args.formal:
-        evolutions.append(FormalEvolution())
-    if args.funny:
-        evolutions.append(FunnyEvolution())
-    if args.goofy:
-        evolutions.append(GoofyEvolution())
-    if args.informal:
-        evolutions.append(InformalEvolution())
+def main() -> None:
+    """CLI endpoint."""
+    mapping = {
+        "reasoning": ReasoningEvolution,
+        "concretizing": ConcreteEvolution,
+        "abstract": AbstractEvolution,
+        "formal": FormalEvolution,
+        "funny": FunnyEvolution,
+        "goofy": GoofyEvolution,
+        "informal": InformalEvolution,
+    }
+    args = _parse_args()
+    evolutions = []
+
+    for arg_name, evolution_cls in mapping.items():
+        if getattr(args, arg_name):
+            evolutions.append(evolution_cls())  # type: ignore[abstract]
 
     if not evolutions:
         logger.warning("No evolutions selected. Exiting.")
         return
 
+    utterance_evolver: UtteranceEvolver
+    if args.decide_for_me:
+        utterance_evolver = IncrementalUtteranceEvolver(Generator(), evolutions, args.seed, args.async_mode)
+    else:
+        utterance_evolver = UtteranceEvolver(Generator(), evolutions, args.seed, args.async_mode)
     dataset = load_dataset(args.input_path)
+
     n_before = len(dataset[args.split])
 
-    generator = UtteranceEvolver(Generator(), evolutions, args.seed, async_mode=args.async_mode)
-    new_samples = generator.augment(dataset, split_name=args.split, n_evolutions=args.n_evolutions)
+    new_samples = utterance_evolver.augment(
+        dataset, split_name=args.split, n_evolutions=args.n_evolutions, batch_size=args.batch_size
+    )
     n_after = len(dataset[args.split])
 
     logger.info("# samples before %s", n_before)

autointent/generation/utterances/evolution/evolver.py

@@ -15,7 +15,7 @@
 from autointent.custom_types import Split
 from autointent.generation.utterances.generator import Generator
 from autointent.generation.utterances.schemas import Message
-from autointent.schemas import Intent, Sample
+from autointent.schemas import Intent
 
 
 class UtteranceEvolver:
@@ -62,7 +62,7 @@ def augment(
         n_evolutions: int = 1,
         update_split: bool = True,
         batch_size: int = 4,
-    ) -> list[Sample]:
+    ) -> HFDataset:
         """
         Augment some split of dataset.
 
@@ -90,11 +90,11 @@ def augment(
                 [{Dataset.label_feature: intent_data.id, Dataset.utterance_feature: ut} for ut in generated_utterances]
             )
 
+        generated_split = HFDataset.from_list(new_samples)
         if update_split:
-            generated_split = HFDataset.from_list(new_samples)
             dataset[split_name] = concatenate_datasets([original_split, generated_split])
 
-        return [Sample(**sample) for sample in new_samples]
+        return generated_split
 
     async def _augment_async(
         self,
@@ -103,7 +103,7 @@ async def _augment_async(
         n_evolutions: int = 1,
         update_split: bool = True,
         batch_size: int = 4,
-    ) -> list[Sample]:
+    ) -> HFDataset:
         original_split = dataset[split_name]
         new_samples = []
 
@@ -124,8 +124,8 @@ async def _augment_async(
             for result, intent_id in zip(batch_results, batch_labels, strict=False):
                 new_samples.append({Dataset.label_feature: intent_id, Dataset.utterance_feature: result})
 
+        generated_split = HFDataset.from_list(new_samples)
         if update_split:
-            generated_split = HFDataset.from_list(new_samples)
             dataset[split_name] = concatenate_datasets([original_split, generated_split])
 
-        return [Sample(**sample) for sample in new_samples]
+        return generated_split

autointent/generation/utterances/evolution/incremental_evolver.py

@@ -0,0 +1,101 @@
+"""
+Evolutionary strategy to augmenting utterances.
+
+Deeply inspired by DeepEval evolutions.
+"""
+
+import copy
+from collections.abc import Callable, Sequence
+from pathlib import Path
+from typing import Any
+
+from datasets import Dataset as HFDataset
+from datasets import concatenate_datasets
+
+from autointent import Dataset, Pipeline
+from autointent.custom_types import Split
+from autointent.generation.utterances.evolution.evolver import UtteranceEvolver
+from autointent.generation.utterances.generator import Generator
+from autointent.generation.utterances.schemas import Message
+from autointent.schemas import Intent
+
+SEARCH_SPACE = [
+    {
+        "node_type": "scoring",
+        "target_metric": "scoring_roc_auc",
+        "metrics": ["scoring_accuracy"],
+        "search_space": [
+            {
+                "module_name": "linear",
+                "embedder_config": ["sentence-transformers/all-MiniLM-L6-v2"],
+            }
+        ],
+    },
+    {
+        "node_type": "decision",
+        "target_metric": "decision_accuracy",
+        "search_space": [
+            {"module_name": "argmax"},
+        ],
+    },
+]
+
+
+class IncrementalUtteranceEvolver(UtteranceEvolver):
+    """Incremental evolutionary strategy to augmenting utterances."""
+
+    def __init__(
+        self,
+        generator: Generator,
+        prompt_makers: Sequence[Callable[[str, Intent], list[Message]]],
+        seed: int = 0,
+        async_mode: bool = False,
+        search_space: str | None = None,
+    ) -> None:
+        """Initialize."""
+        super().__init__(generator, prompt_makers, seed, async_mode)
+        self.search_space = self._choose_search_space(search_space)
+
+    def _choose_search_space(self, search_space: str | None) -> list[dict[str, Any]] | Path | str:
+        if search_space is None:
+            return SEARCH_SPACE
+        return search_space
+
+    def augment(
+        self,
+        dataset: Dataset,
+        split_name: str = Split.TRAIN,
+        n_evolutions: int = 1,
+        update_split: bool = True,
+        batch_size: int = 4,
+    ) -> HFDataset:
+        """
+        Augment some split of dataset.
+
+        Note that for now it supports only single-label datasets.
+        """
+        best_result = 0
+        merge_dataset = copy.deepcopy(dataset)
+        generated_samples = []
+
+        for _ in range(n_evolutions):
+            new_samples_dataset = super().augment(
+                dataset, split_name=split_name, n_evolutions=1, update_split=False, batch_size=batch_size
+            )
+            merge_dataset[split_name] = concatenate_datasets([merge_dataset[split_name], new_samples_dataset])
+            generated_samples.append(new_samples_dataset)
+
+            pipeline_optimizer = Pipeline.from_search_space(self.search_space)
+            ctx = pipeline_optimizer.fit(merge_dataset)
+            results = ctx.optimization_info.dump_evaluation_results()
+            decision_metric = results["metrics"]["decision"][0]
+
+            if decision_metric > best_result:
+                best_result = decision_metric
+            else:
+                break
+
+        if update_split:
+            dataset[split_name] = merge_dataset[split_name]
+
+        return concatenate_datasets(generated_samples)