refactor

Author: Darinochka

autointent/generation/utterances/evolution/cli.py

@@ -3,9 +3,10 @@
 import logging
 from argparse import ArgumentParser, Namespace
 from pathlib import Path
-from typing import Any
+from typing import Any, Optional
 
 from datasets import concatenate_datasets
+from datasets import Dataset as HFDataset
 
 from autointent import Dataset, Pipeline, load_dataset
 from autointent.configs import EmbedderConfig
@@ -63,24 +64,22 @@ def _choose_search_space(search_space: str | None) -> list[dict[str, Any]] | Pat
 
 
 def _optimize_n_evolutions(
+    generator: Generator,
     input_path: str,
+    dataset: Dataset,
     max_n_evolutions: int,
-    evolutions: list,
-    seed: int,
     split_train: str,
-    async_mode: bool,
     batch_size: int,
-    search_space: list[dict[str, Any]] | Path | str,
+    search_space: Optional[str],
 ) -> Dataset:
     emb_config = EmbedderConfig(batch_size=16, device="cuda")
+    search_space = _choose_search_space(search_space)
 
     best_result = 0
     best_n = 0
-    dataset = load_dataset(input_path)
     merge_dataset = load_dataset(input_path)
 
     for n in range(max_n_evolutions):
-        generator = UtteranceEvolver(Generator(), evolutions, seed, async_mode)
         new_samples_dataset = generator.augment(
             dataset, split_name=split_train, n_evolutions=1, update_split=False, batch_size=batch_size
         )
@@ -101,32 +100,6 @@ def _optimize_n_evolutions(
     logger.info("# optimal n evolutions: %s", best_n)
     return dataset
 
-
-def _generate_fixed_evolutions(
-    input_path: str,
-    n_evolutions: int,
-    evolutions: list,
-    seed: int,
-    split: str,
-    async_mode: bool,
-    batch_size: int,
-    *args,
-    **kwargs,
-) -> Dataset:
-    dataset = load_dataset(input_path)
-    n_before = len(dataset[split])
-
-    generator = UtteranceEvolver(Generator(), evolutions, seed, async_mode)
-    new_samples = generator.augment(dataset, split_name=split, n_evolutions=n_evolutions, batch_size=batch_size)
-    n_after = len(dataset[split])
-
-    logger.info("# samples before %s", n_before)
-    logger.info("# samples generated %s", len(new_samples))
-    logger.info("# samples after %s", n_after)
-
-    return dataset
-
-
 def _parse_args() -> Namespace:
     parser = ArgumentParser()
     parser.add_argument(
@@ -188,22 +161,30 @@ def main() -> None:
         logger.warning("No evolutions selected. Exiting.")
         return
 
-    search_space = _choose_search_space(args.search_space)
 
-    process_func = _generate_fixed_evolutions
+    generator = UtteranceEvolver(Generator(), evolutions, args.seed, args.async_mode)
+    dataset = load_dataset(args.input_path)
+
     if args.decide_for_me:
-        process_func = _optimize_n_evolutions
-
-    dataset = process_func(
-        args.input_path,
-        args.n_evolutions,
-        evolutions,
-        args.seed,
-        args.split,
-        args.async_mode,
-        args.batch_size,
-        search_space,
-    )
+        dataset = _optimize_n_evolutions(
+            generator,
+            args.input_path,
+            dataset,
+            args.n_evolutions,
+            args.split,
+            args.batch_size,
+            args.search_space,
+        )
+    else:
+        n_before = len(dataset[args.split])
+    
+        new_samples = generator.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)
+        logger.info("# samples generated %s", len(new_samples))
+        logger.info("# samples after %s", n_after)
+        
     dataset.to_json(args.output_path)
 
     if args.output_repo is not None: