feat: batch_size in basic augmentation

Author: Darinochka

autointent/generation/utterances/basic/utterance_generator.py

@@ -21,11 +21,12 @@ class UtteranceGenerator:
     punctuation, and length of the desired generations.
     """
 
-    def __init__(self,
-                 generator: Generator,
-                 prompt_maker: Callable[[Intent, int], list[Message]],
-                 async_mode: bool = False
-        ) -> None:
+    def __init__(
+        self,
+        generator: Generator,
+        prompt_maker: Callable[[Intent, int], list[Message]],
+        async_mode: bool = False
+    ) -> None:
         """Initialize."""
         self.generator = generator
         self.prompt_maker = prompt_maker
@@ -49,27 +50,33 @@ def augment(
         split_name: str = Split.TRAIN,
         n_generations: int = 5,
         update_split: bool = True,
+        batch_size: int | None = None
     ) -> list[Sample]:
         """
         Augment some split of dataset.
 
-        Note that for now it supports only single-label datasets.
+        :param dataset: Dataset object
+        :param split_name: Dataset split (default is TRAIN)
+        :param n_generations: Number of utterances to generate per intent
+        :param update_split: Whether to update the dataset split
+        :param batch_size: Batch size for async generation (None means all at once)
+        :return: List of generated samples
         """
         if self.async_mode:
-            return asyncio.run(self._augment_async(dataset, split_name, n_generations, update_split))
+            return asyncio.run(self._augment_async(dataset, split_name, n_generations, update_split, batch_size))
+
         original_split = dataset[split_name]
         new_samples = []
         for intent in dataset.intents:
-            generated_utterances = self(
-                intent_data=intent,
-                n_generations=n_generations,
-            )
+            generated_utterances = self(intent_data=intent, n_generations=n_generations)
             new_samples.extend(
                 [{Dataset.label_feature: intent.id, Dataset.utterance_feature: ut} for ut in generated_utterances]
             )
+
         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]
 
     async def _augment_async(
@@ -78,19 +85,32 @@ async def _augment_async(
         split_name: str = Split.TRAIN,
         n_generations: int = 5,
         update_split: bool = True,
+        batch_size: int | None = None
     ) -> list[Sample]:
         """
-        Augment some split of dataset asynchronously.
-
-        Note that for now it supports only single-label datasets.
+        Augment some split of dataset asynchronously in batches.
+
+        :param dataset: Dataset object
+        :param split_name: Dataset split (default is TRAIN)
+        :param n_generations: Number of utterances to generate per intent
+        :param update_split: Whether to update the dataset split
+        :param batch_size: Batch size for async generation (None means all at once)
+        :return: List of generated samples
         """
         original_split = dataset[split_name]
         new_samples = []
-        tasks = []
 
-        tasks = [self._call_async(intent_data=intent, n_generations=n_generations) for intent in dataset.intents]
+        if not batch_size:
+            tasks = [self._call_async(intent_data=intent, n_generations=n_generations) for intent in dataset.intents]
+            results = await asyncio.gather(*tasks)
 
-        results = await asyncio.gather(*tasks)
+        else:
+            results = []
+            for start_idx in range(0, len(dataset.intents), batch_size):
+                batch_intents = dataset.intents[start_idx:start_idx + batch_size]
+                tasks = [self._call_async(intent_data=intent, n_generations=n_generations) for intent in batch_intents]
+                batch_results = await asyncio.gather(*tasks)
+                results.extend(batch_results)
 
         for i, generated_utterances in enumerate(results):
             intent = dataset.intents[i]
@@ -113,4 +133,4 @@ def _extract_utterances(response_text: str) -> list[str]:
     """
     raw_utterances = response_text.split("\n")
     # remove enumeration
-    return [ut[ut.find(" ") + 1 :] for ut in raw_utterances]
+    return [ut[ut.find(" ") + 1:] if " " in ut else ut for ut in raw_utterances]

tests/generation/utterances/test_basic_synthesizer.py

@@ -74,3 +74,24 @@ def test_on_dataset_async(dataset):
 
     assert n_before + len(new_samples) == n_after
     assert len(new_samples) == len(dataset.intents)
+
+def test_on_dataset_async_with_batch_size(dataset):
+    mock_llm = AsyncMock()
+    mock_llm.get_chat_completion_async.return_value = "1. LLM answer"
+
+    split_name = "train_0"
+
+    template = SynthesizerChatTemplate(dataset, split=split_name)
+    augmenter = UtteranceGenerator(mock_llm, template, async_mode=True)
+
+    batch_size = 2
+    new_samples = augmenter.augment(dataset, split_name=split_name, update_split=False, batch_size=batch_size)
+
+    assert len(new_samples) == len(dataset.intents)
+    assert all(sample.utterance == "LLM answer" for sample in new_samples)
+
+    batch_size = len(dataset.intents) + 5
+    new_samples = augmenter.augment(dataset, split_name=split_name, update_split=False, batch_size=batch_size)
+
+    assert len(new_samples) == len(dataset.intents)
+    assert all(sample.utterance == "LLM answer" for sample in new_samples)

tests/generation/utterances/test_evolver.py

@@ -56,3 +56,36 @@ def test_on_dataset_evolver_async(dataset):
 
     assert n_before + len(new_samples) == n_after
     assert len(new_samples) == n_before
+
+def test_on_dataset_evolver_async_with_batch_size(dataset):
+    mock_llm = AsyncMock()
+    mock_llm.get_chat_completion_async.return_value = "LLM answer"
+
+    split_name = "train_0"
+
+    template = AbstractEvolution()
+    augmenter = UtteranceEvolver(mock_llm, [template], async_mode=True)
+
+    batch_size = 2
+    new_samples = augmenter.augment(
+        dataset,
+        split_name=split_name,
+        n_evolutions=1,
+        update_split=False,
+        batch_size=batch_size
+    )
+
+    assert len(new_samples) == len(dataset[split_name])
+    assert all(sample.utterance == "LLM answer" for sample in new_samples)
+
+    batch_size = len(dataset[split_name]) + 5
+    new_samples = augmenter.augment(
+        dataset,
+        split_name=split_name,
+        n_evolutions=1,
+        update_split=False,
+        batch_size=batch_size
+    )
+
+    assert len(new_samples) == len(dataset[split_name])
+    assert all(sample.utterance == "LLM answer" for sample in new_samples)