fix: main

Author: Darinochka

autointent/generation/utterances/basic/cli.py

@@ -13,7 +13,7 @@
 logger = logging.getLogger(__name__)
 
 
-async def main() -> None:
+def main() -> None:
     """CLI endpoint."""
     parser = ArgumentParser()
     parser.add_argument(

autointent/generation/utterances/evolution/evolver.py

@@ -52,7 +52,7 @@ async def _evolve_async(self, utterance: str, intent_data: Intent) -> str:
         return await self.generator.get_chat_completion_async(chat)
 
     def __call__(self, utterance: str, intent_data: Intent, n_evolutions: int = 1) -> list[str]:
-        """Apply evolutions multiple times."""
+        """Apply evolutions multiple times (synchronously)."""
         return [self._evolve(utterance, intent_data) for _ in range(n_evolutions)]
 
     async def _call_async(self, utterance: str, intent_data: Intent, n_evolutions: int = 1) -> list[str]:
@@ -61,15 +61,29 @@ async def _call_async(self, utterance: str, intent_data: Intent, n_evolutions: i
         return await asyncio.gather(*tasks)
 
     def augment(
-        self, dataset: Dataset, split_name: str = Split.TRAIN, n_evolutions: int = 1, update_split: bool = True
+        self,
+        dataset: Dataset,
+        split_name: str = Split.TRAIN,
+        n_evolutions: int = 1,
+        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.
         """
         if self.async_mode:
-            return asyncio.run(self._augment_async(dataset, split_name, n_evolutions, update_split))
+            return asyncio.run(
+                self._augment_async(
+                    dataset=dataset,
+                    split_name=split_name,
+                    n_evolutions=n_evolutions,
+                    update_split=update_split,
+                    batch_size=batch_size
+                )
+            )
+
         original_split = dataset[split_name]
         new_samples = []
         for sample in original_split:
@@ -80,37 +94,70 @@ def augment(
             new_samples.extend(
                 [{Dataset.label_feature: intent_data.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(
-        self, dataset: Dataset, split_name: str = Split.TRAIN, n_evolutions: int = 1, update_split: bool = True
+        self,
+        dataset: Dataset,
+        split_name: str = Split.TRAIN,
+        n_evolutions: int = 1,
+        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.
-        """
         original_split = dataset[split_name]
         new_samples = []
-        tasks = []
 
-        for sample in original_split:
-            utterance = sample[Dataset.utterance_feature]
-            label = sample[Dataset.label_feature]
-            intent_data = next(intent for intent in dataset.intents if intent.id == label)
-            tasks.append(self._call_async(utterance=utterance, intent_data=intent_data, n_evolutions=n_evolutions))
-
-        results = await asyncio.gather(*tasks)
-
-        for i, generated_utterances in enumerate(results):
-            intent_data = next(intent for intent in dataset.intents if intent.id == original_split[i][
-                Dataset.label_feature])
-            new_samples.extend(
-                [{Dataset.label_feature: intent_data.id, Dataset.utterance_feature: ut} for ut in generated_utterances]
-            )
+        if not batch_size:
+            tasks = []
+            for sample in original_split:
+                utterance = sample[Dataset.utterance_feature]
+                label = sample[Dataset.label_feature]
+                intent_data = next(intent for intent in dataset.intents if intent.id == label)
+                tasks.append(
+                    self._call_async(utterance=utterance, intent_data=intent_data, n_evolutions=n_evolutions)
+                )
+
+            results = await asyncio.gather(*tasks)
+
+            for i, generated_utterances in enumerate(results):
+                intent_data = next(
+                    intent for intent in dataset.intents if intent.id == original_split[i][Dataset.label_feature]
+                )
+                new_samples.extend(
+                    [{Dataset.label_feature: intent_data.id, Dataset.utterance_feature: ut}
+                     for ut in generated_utterances]
+                )
+
+        else:
+            total_samples = len(original_split)
+            for start_idx in range(0, total_samples, batch_size):
+                batch = original_split[start_idx : start_idx + batch_size]
+                tasks = []
+                for utterance, label in zip(
+                    batch[Dataset.utterance_feature],
+                    batch[Dataset.label_feature],
+                    strict=False
+                ):
+                    intent_data = next(intent for intent in dataset.intents if intent.id == label)
+                    tasks.append(
+                        self._call_async(utterance=utterance, intent_data=intent_data, n_evolutions=n_evolutions)
+                    )
+
+                batch_results = await asyncio.gather(*tasks)
+
+                for i, generated_utterances in enumerate(batch_results):
+                    intent_data = next(
+                        intent for intent in dataset.intents if intent.id == batch[Dataset.label_feature][i]
+                    )
+                    new_samples.extend(
+                        [{Dataset.label_feature: intent_data.id, Dataset.utterance_feature: ut}
+                         for ut in generated_utterances]
+                    )
 
         if update_split:
             generated_split = HFDataset.from_list(new_samples)