Refactor augmentations (#147)

* update chat template structure

* update chat template structure

* refactor chat templates

* lint

* fix imports

Author: Samoed

autointent/_dataset/_validation.py

@@ -101,7 +101,7 @@ def _validate_classes(self, splits: list[list[Sample]]) -> int:
             )
             raise ValueError(message)
         if not n_classes[0]:
-            message = "Number of classes is zero or undefined. " "Ensure at least one class is present in the splits."
+            message = "Number of classes is zero or undefined. Ensure at least one class is present in the splits."
             raise ValueError(message)
         return n_classes[0]
 
@@ -120,8 +120,7 @@ def _validate_intents(self, n_classes: int) -> "DatasetReader":
         intent_ids = [intent.id for intent in self.intents]
         if intent_ids != list(range(len(self.intents))):
             message = (
-                f"Invalid intent IDs. Expected sequential IDs from 0 to {len(self.intents) - 1}, "
-                f"but got {intent_ids}."
+                f"Invalid intent IDs. Expected sequential IDs from 0 to {len(self.intents) - 1}, but got {intent_ids}."
             )
             raise ValueError(message)
         return self

autointent/context/data_handler/_data_handler.py

@@ -236,7 +236,7 @@ def _split_cv(self) -> None:
                 random_seed=self.random_seed,
                 allow_oos_in_train=True,
             )
-        self.dataset[f"{Split.TRAIN}_{self.config.n_folds-1}"] = self.dataset.pop(Split.TRAIN)
+        self.dataset[f"{Split.TRAIN}_{self.config.n_folds - 1}"] = self.dataset.pop(Split.TRAIN)
 
     def _split_validation_from_train(self, size: float) -> None:
         if Split.TRAIN in self.dataset:

autointent/generation/utterances/basic/chat_templates/_base.py

@@ -66,10 +66,10 @@ def __call__(self, intent_data: Intent, n_examples: int) -> list[Message]:
         ]
 
     def _create_final_message(self, intent_data: Intent, n_examples: int, sample_utterances: list[str]) -> Message:
-        content = f"{self._INTENT_NAME_LABEL}: {intent_data.name}\n\n" f"{self._EXAMPLE_UTTERANCES_LABEL}:\n"
+        content = f"{self._INTENT_NAME_LABEL}: {intent_data.name}\n\n{self._EXAMPLE_UTTERANCES_LABEL}:\n"
 
         if sample_utterances:
-            numbered_utterances = "\n".join(f"{i+1}. {utt}" for i, utt in enumerate(sample_utterances))
+            numbered_utterances = "\n".join(f"{i + 1}. {utt}" for i, utt in enumerate(sample_utterances))
             content += numbered_utterances + "\n\n"
 
         content += self._GENERATE_INSTRUCTION.format(n_examples=n_examples)

autointent/generation/utterances/basic/chat_templates/_synthesizer_en.py

@@ -70,7 +70,7 @@ class EnglishSynthesizerTemplate(BaseSynthesizerTemplate):
         ),
         Message(
             role=Role.ASSISTANT,
-            content=("1. Can you tell me the forecast for tomorrow?\n" "2. Is it going to rain this weekend?"),
+            content="1. Can you tell me the forecast for tomorrow?\n2. Is it going to rain this weekend?",
         ),
         Message(
             role=Role.USER,

autointent/generation/utterances/basic/chat_templates/_synthesizer_ru.py

@@ -53,7 +53,7 @@ class RussianSynthesizerTemplate(BaseSynthesizerTemplate):
         ),
         Message(
             role=Role.ASSISTANT,
-            content=("1. Забронируйте люкс в Санкт-Петербурге на выходные\n" "2. Ищу номер с видом на море в Сочи"),
+            content=("1. Забронируйте люкс в Санкт-Петербурге на выходные\n2. Ищу номер с видом на море в Сочи"),
         ),
         Message(
             role=Role.USER,
@@ -66,7 +66,7 @@ class RussianSynthesizerTemplate(BaseSynthesizerTemplate):
         ),
         Message(
             role=Role.ASSISTANT,
-            content=("1. Какой прогноз на завтра?\n" "2. Будет ли дождь в субботу?"),
+            content=("1. Какой прогноз на завтра?\n2. Будет ли дождь в субботу?"),
         ),
         Message(
             role=Role.USER,

autointent/generation/utterances/basic/utterance_generator.py

@@ -1,15 +1,14 @@
 """Basic generation of new utterances from existing ones."""
 
 import asyncio
-from collections.abc import Callable
 
 from datasets import Dataset as HFDataset
 from datasets import concatenate_datasets
 
 from autointent import Dataset
 from autointent.custom_types import Split
+from autointent.generation.utterances.basic.chat_templates import BaseSynthesizerTemplate
 from autointent.generation.utterances.generator import Generator
-from autointent.generation.utterances.schemas import Message
 from autointent.schemas import Intent, Sample
 
 
@@ -22,9 +21,7 @@ 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: BaseSynthesizerTemplate, async_mode: bool = False) -> None:
         """Initialize."""
         self.generator = generator
         self.prompt_maker = prompt_maker

autointent/generation/utterances/evolution/chat_templates/__init__.py

@@ -7,7 +7,13 @@
 from .informal import InformalEvolution
 from .reasoning import ReasoningEvolution
 
+EVOLUTION_NAMES = [evolution.name for evolution in EvolutionChatTemplate.__subclasses__()]
+
+EVOLUTION_MAPPING = {evolution.name: evolution() for evolution in EvolutionChatTemplate.__subclasses__()}
+
 __all__ = [
+    "EVOLUTION_MAPPING",
+    "EVOLUTION_NAMES",
     "AbstractEvolution",
     "ConcreteEvolution",
     "EvolutionChatTemplate",

autointent/generation/utterances/evolution/chat_templates/abstract.py

@@ -3,14 +3,14 @@
 from typing import ClassVar
 
 from autointent.generation.utterances.schemas import Message, Role
-from autointent.schemas import Intent
 
 from .base import EvolutionChatTemplate
 
 
 class AbstractEvolution(EvolutionChatTemplate):
     """Chat template for evolution augmentation via abstraction."""
 
+    name = "abstract"
     _messages: ClassVar[list[Message]] = [
         Message(
             role=Role.USER,
@@ -36,10 +36,3 @@ class AbstractEvolution(EvolutionChatTemplate):
         ),
         Message(role=Role.ASSISTANT, content="I'm having trouble with my laptop."),
     ]
-
-    def __call__(self, utterance: str, intent_data: Intent) -> list[Message]:
-        """Make chat to complete."""
-        return [
-            *self._messages,
-            Message(role=Role.USER, content=f"Intent name: {intent_data.name or ''}\nUtterance: {utterance}"),
-        ]

autointent/generation/utterances/evolution/chat_templates/base.py

@@ -1,14 +1,21 @@
 """Base class for chat templates for evolution augmentation."""
 
-from abc import ABC, abstractmethod
+from typing import ClassVar
 
-from autointent.generation.utterances.schemas import Message
+from autointent.generation.utterances.schemas import Message, Role
 from autointent.schemas import Intent
 
 
-class EvolutionChatTemplate(ABC):
+class EvolutionChatTemplate:
     """Base class for chat templates for evolution augmentation."""
 
-    @abstractmethod
+    _messages: ClassVar[list[Message]]
+    name: str
+
     def __call__(self, utterance: str, intent_data: Intent) -> list[Message]:
         """Make a chat to complete by LLM."""
+        invoke_message = Message(
+            role=Role.USER,
+            content=f"Intent name: {intent_data.name or ''}\nUtterance: {utterance}",
+        )
+        return [*self._messages, invoke_message]

autointent/generation/utterances/evolution/chat_templates/concrete.py

@@ -3,14 +3,15 @@
 from typing import ClassVar
 
 from autointent.generation.utterances.schemas import Message, Role
-from autointent.schemas import Intent
 
 from .base import EvolutionChatTemplate
 
 
 class ConcreteEvolution(EvolutionChatTemplate):
     """Chat template for evolution augmentation via concretizing."""
 
+    name = "concrete"
+
     _messages: ClassVar[list[Message]] = [
         Message(
             role=Role.USER,
@@ -29,14 +30,7 @@ class ConcreteEvolution(EvolutionChatTemplate):
         Message(role=Role.ASSISTANT, content="I want to reserve a table for 4 persons at 9 pm."),
         Message(
             role=Role.USER,
-            content=("Intent name: requesting technical support\n" "Utterance: I'm having trouble with my laptop."),
+            content="Intent name: requesting technical support\nUtterance: I'm having trouble with my laptop.",
         ),
         Message(role=Role.ASSISTANT, content="My laptop is constantly rebooting and overheating."),
     ]
-
-    def __call__(self, utterance: str, intent_data: Intent) -> list[Message]:
-        """Make chat to complete."""
-        return [
-            *self._messages,
-            Message(role=Role.USER, content=f"Intent name: {intent_data.name or ''}\nUtterance: {utterance}"),
-        ]