modify dataset_generator.py to support small_model_training

src/fabricator/dataset_generator.py

@@ -1,6 +1,7 @@
 import json
 import time
 
+from importlib import import_module
 from pathlib import Path
 from collections import defaultdict
 from typing import Any, Callable, Dict, Optional, Union, Tuple, List
@@ -16,7 +17,6 @@
 from .samplers import single_label_stratified_sample
 from .utils import log_dir, create_timestamp_path
 
-
 class DatasetGenerator:
     """The DatasetGenerator class is the main class of the fabricator package.
     It generates datasets based on a prompt template. The main function is generate()."""
@@ -59,6 +59,8 @@ def generate(
         return_unlabeled_dataset: bool = False,
         max_prompt_calls: int = 10,
         num_samples_to_generate: int = 10,
+        small_model_training: Optional[str] = None,
+        train_small_model_every_X_generations: Optional[int] = None,
         timeout_per_prompt: Optional[int] = None,
         log_every_n_api_calls: int = 25,
         dummy_response: Optional[Union[str, Callable]] = None
@@ -82,6 +84,9 @@ def generate(
             return_unlabeled_dataset (bool, optional): Whether to return the original dataset. Defaults to False.
             max_prompt_calls (int, optional): Maximum number of prompt calls. Defaults to 10.
             num_samples_to_generate (int, optional): Number of samples to generate. Defaults to 10.
+            small_model_training (str, optional): Task to perform small model training on. Defaults to None.
+            train_small_model_every_X_generations (int, optional): Number of generations between small model
+                training iterations. Defaults to None.
             timeout_per_prompt (Optional[int], optional): Timeout per prompt call. Defaults to None.
             log_every_n_api_calls (int, optional): Log every n api calls. Defaults to 25.
             dummy_response (Optional[Union[str, Callable]], optional): Dummy response for dry runs. Defaults to None.
@@ -99,6 +104,9 @@ def generate(
         if fewshot_dataset and not fewshot_sampling_column:
             fewshot_sampling_column = prompt_template.generate_data_for_column[0]
 
+        assert small_model_training in [None, "text-classification"], \
+            "Task for small model training must be available in 'src/small_model_training' e.g. 'text-classification'"
+
         generated_dataset, original_dataset = self._inner_generate_loop(
             prompt_template,
             fewshot_dataset,
@@ -109,6 +117,8 @@ def generate(
             return_unlabeled_dataset,
             max_prompt_calls,
             num_samples_to_generate,
+            small_model_training,
+            train_small_model_every_X_generations,
             timeout_per_prompt,
             log_every_n_api_calls,
             dummy_response
@@ -170,6 +180,8 @@ def _inner_generate_loop(
         return_unlabeled_dataset: bool,
         max_prompt_calls: int,
         num_samples_to_generate: int,
+        small_model_training: Optional[str],
+        train_small_model_every_x_generations: Optional[int],
         timeout_per_prompt: Optional[int],
         log_every_n_api_calls: int = 25,
         dummy_response: Optional[Union[str, Callable]] = None
@@ -284,6 +296,12 @@ def _inner_generate_loop(
             if timeout_per_prompt is not None:
                 time.sleep(timeout_per_prompt)
 
+            if train_small_model_every_x_generations > 0:
+                if prompt_call_idx % train_small_model_every_x_generations == 0:
+                    small_model = import_module(small_model_training,"src.small_model_training")
+                    inf_subset = small_model.get_influential_subset(generated_dataset)
+                    fewshot_dataset = inf_subset
+
         generated_dataset = Dataset.from_dict(generated_dataset)
 
         if return_unlabeled_dataset: