adjustments for running small model during generation

src/fabricator/dataset_generator.py

@@ -210,16 +210,16 @@ def _inner_generate_loop(
                 # require a second parameter for sample from label options and not from fewshot examples
                 prompt_labels = prompt_template.label_options
 
-            if fewshot_dataset:
+            # label-conditioned generation with label options
+            if fewshot_dataset is None and isinstance(prompt_labels,list) or prompt_call_idx > train_small_model_every_x_generations:
+                prompt_labels = choice(prompt_labels, 1)[0]
+
+            if fewshot_dataset and prompt_labels in fewshot_dataset['label']:
                 prompt_labels, fewshot_examples = self._sample_fewshot_examples(
-                    prompt_template, fewshot_dataset, fewshot_sampling_strategy, fewshot_examples_per_class,
+                    prompt_labels, fewshot_dataset, fewshot_sampling_strategy, fewshot_examples_per_class,
                     fewshot_sampling_column
                 )
 
-            # label-conditioned generation with label options
-            if fewshot_dataset is None and isinstance(prompt_labels, list):
-                prompt_labels = choice(prompt_labels,1)[0]
-
             prompt_text = prompt_template.get_prompt_text(prompt_labels, fewshot_examples)
 
             if unlabeled_dataset:
@@ -302,9 +302,11 @@ def _inner_generate_loop(
 
             if train_small_model_every_x_generations is not None and train_small_model_every_x_generations > 0:
                 if prompt_call_idx % train_small_model_every_x_generations == 0:
+                    logger.info("Commencing small model training.")
                     small_model = import_module(small_model_training,"src.small_model_training")
                     inf_subset = small_model.get_influential_subset(generated_dataset)
                     fewshot_dataset = inf_subset
+                    logger.info("Continuing generation.")
 
         generated_dataset = Dataset.from_dict(generated_dataset)
 
@@ -339,29 +341,30 @@ def _convert_prediction(self, prediction: str, target_type: type) -> Any:
 
     @staticmethod
     def _sample_fewshot_examples(
-        prompt_template: BasePrompt,
+        prompt_labels,
         fewshot_dataset: Dataset,
         fewshot_sampling_strategy: str,
         fewshot_examples_per_class: int,
         fewshot_sampling_column: str
     ) -> Tuple[Union[List[str], str], Dataset]:
 
         if fewshot_sampling_strategy == "uniform":
-            prompt_labels = choice(prompt_template.label_options, 1)[0]
+            prompt_labels = choice(prompt_labels, 1)[0]
             fewshot_examples = fewshot_dataset.filter(
                 lambda example: example[fewshot_sampling_column] == prompt_labels
-            ).shuffle().select(range(fewshot_examples_per_class))
+            )
+            fewshot_examples = fewshot_examples.shuffle().select(
+                range(fewshot_examples_per_class) if fewshot_examples_per_class is not None else range(len(fewshot_examples))
+            )
 
         elif fewshot_sampling_strategy == "stratified":
-            prompt_labels = prompt_template.label_options
             fewshot_examples = single_label_stratified_sample(
                 fewshot_dataset,
                 fewshot_sampling_column,
                 fewshot_examples_per_class
             )
 
         else:
-            prompt_labels = prompt_template.label_options if prompt_template.label_options else None
             if fewshot_examples_per_class:
                 fewshot_examples = fewshot_dataset.shuffle().select(range(fewshot_examples_per_class))
             else:

src/fabricator/prompts/base.py

@@ -190,6 +190,12 @@ def get_prompt_text(self, labels: Union[str, List[str]] = None, examples: Option
             task_description = self.task_description
 
         if examples:
+            # due to small_model training fewshot_examples appear after some iterations
+            if self.relevant_columns_for_fewshot_examples is None:
+                self.relevant_columns_for_fewshot_examples = examples.column_names
+                self.fewshot_prompt = self.inner_fewshot_example_separator.join(
+                    [f"{var}: {{{var}}}" for var in self.relevant_columns_for_fewshot_examples]
+                )
             examples = self.filter_examples_by_columns(examples, self.relevant_columns_for_fewshot_examples)
             formatted_examples = [self.fewshot_prompt.format(**example) for example in examples]
             prompt_text = self.fewshot_example_separator.join(

src/small_model_training/iterative_dataset_synthesis_movies.py

@@ -20,8 +20,7 @@
     max_prompt_calls=100,
     num_samples_to_generate=100,
     small_model_training='text_classification',
-    train_small_model_every_X_generations=10
-
+    train_small_model_every_X_generations=10,
 )
 
 generated_dataset.push_to_hub("your-first-generated-dataset")

src/small_model_training/text_classification.py

@@ -1,3 +1,4 @@
+import math
 import torch.nn
 import numpy as np
 from dataclasses import dataclass, field
@@ -20,8 +21,15 @@ def get_influential_subset(dataset):
     # setup preprocessing
     tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path)
 
-    def preprocess_function(examples):
-        return tokenizer(examples["text"], truncation=True, padding=True)
+    # define labels
+    # TO-DO: dynamic labels
+    id2label = {0: "negative", 1: "positive"}
+    label2id = {"negative": 0, "positive": 1}
+
+    def preprocess_function(batch):
+        preprocessed_tokens = tokenizer(batch["text"], truncation=True, padding=True)
+        preprocessed_tokens["label"] = [label2id[label] for label in batch["label"]]
+        return preprocessed_tokens
 
     # setup compute_metrics
     accuracy = evaluate.load("accuracy")
@@ -43,11 +51,6 @@ def compute_metrics(eval_pred):
 
     data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
 
-    # define labels
-    # TO-DO: dynamic labels
-    id2label = {0: "negative", 1: "positive"}
-    label2id = {"negative": 0, "positive": 1}
-
     model = AutoModelForSequenceClassification.from_pretrained(
         model_args.model_name_or_path, num_labels=2, id2label=id2label, label2id=label2id
     )
@@ -68,16 +71,17 @@ def compute_metrics(eval_pred):
 
     inf_subset = []
 
-    indices = calculate_candidates_knn(outputs)
+    indices = calculate_candidates_knn(outputs, math.ceil(len(dataset["test"])/2))
 
     for elem in indices:
         inf_subset.append(dataset["test"][elem.item()])
 
-    # TO-DO: check for pre-processing
+    inf_subset = Dataset.from_list(inf_subset)
+
     return inf_subset
 
 #TO-DO: dynamic variant for calculating inf_subset
-def calculate_candidates_knn(model_outputs):
+def calculate_candidates_knn(model_outputs, num_candidates):
     logits = model_outputs[0]
 
     logits = torch.from_numpy(logits)
@@ -88,6 +92,7 @@ def calculate_candidates_knn(model_outputs):
     second_values = scores[:, 1]
 
     distance = torch.abs(first_values - second_values)
-    _, knn_indices = distance.topk(5, largest=False)
+    #make dynamic
+    _, knn_indices = distance.topk(num_candidates, largest=False)
 
     return knn_indices