update chat template structure

Author: Samoed

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,
@@ -35,4 +35,4 @@ class AbstractEvolution(EvolutionChatTemplate):
             ),
         ),
         Message(role=Role.ASSISTANT, content="I'm having trouble with my laptop."),
-    ]
+    ]

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

@@ -1,6 +1,5 @@
 """Base class for chat templates for evolution augmentation."""
 
-from abc import ABC, abstractmethod
 from typing import ClassVar
 
 from autointent.generation.utterances.schemas import Message, Role
@@ -10,13 +9,13 @@
 class EvolutionChatTemplate:
     """Base class for chat templates for evolution augmentation."""
 
-    invoke_message: Message
     _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]
+        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,
@@ -33,4 +34,3 @@ class ConcreteEvolution(EvolutionChatTemplate):
         ),
         Message(role=Role.ASSISTANT, content="My laptop is constantly rebooting and overheating."),
     ]
-

autointent/generation/utterances/evolution/chat_templates/formal.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 FormalEvolution(EvolutionChatTemplate):
     """Chat template for formal tone augmentation."""
 
+    name: str = "formal"
+
     _messages: ClassVar[list[Message]] = [
         Message(
             role=Role.USER,
@@ -39,4 +40,3 @@ class FormalEvolution(EvolutionChatTemplate):
             content="My Lenovo laptop frequently restarts and experiences overheating issues. Kindly assist.",
         ),
     ]
-

autointent/generation/utterances/evolution/chat_templates/funny.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 FunnyEvolution(EvolutionChatTemplate):
     """Chat template for humorous tone augmentation."""
 
+    name: str = "funny"
     _messages: ClassVar[list[Message]] = [
         Message(
             role=Role.USER,

autointent/generation/utterances/evolution/chat_templates/goofy.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 GoofyEvolution(EvolutionChatTemplate):
     """Chat template for goofy tone augmentation."""
 
+    name: str = "goofy"
     _messages: ClassVar[list[Message]] = [
         Message(
             role=Role.USER,

autointent/generation/utterances/evolution/chat_templates/informal.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 InformalEvolution(EvolutionChatTemplate):
     """Chat template for informal tone augmentation."""
 
+    name: str = "informal"
     _messages: ClassVar[list[Message]] = [
         Message(
             role=Role.USER,
@@ -35,4 +35,4 @@ class InformalEvolution(EvolutionChatTemplate):
             ),
         ),
         Message(role=Role.ASSISTANT, content="My Lenovo keeps crashing and getting super hot. Any ideas?"),
-    ]
+    ]

autointent/generation/utterances/evolution/chat_templates/reasoning.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 ReasoningEvolution(EvolutionChatTemplate):
     """Chat template for evolution augmentation via reasoning."""
 
+    name = "reasoning"
+
     _messages: ClassVar[list[Message]] = [
         Message(
             role=Role.USER,

autointent/generation/utterances/evolution/cli.py

@@ -8,13 +8,8 @@
 from autointent.generation.utterances.generator import Generator
 
 from .chat_templates import (
-    AbstractEvolution,
-    ConcreteEvolution,
-    FormalEvolution,
-    FunnyEvolution,
-    GoofyEvolution,
-    InformalEvolution,
-    ReasoningEvolution,
+    EVOLUTION_MAPPING,
+    EVOLUTION_NAMES,
 )
 
 logging.basicConfig(level="INFO")
@@ -44,14 +39,8 @@ def _parse_args() -> Namespace:
     )
     parser.add_argument("--private", action="store_true", help="Publish privately if --output-repo option is used")
     parser.add_argument("--n-evolutions", type=int, default=1, help="Number of utterances to generate for each intent")
-    parser.add_argument("--decide-for-me", action="store_true")
-    parser.add_argument("--reasoning", action="store_true", help="Whether to use `Reasoning` evolution")
-    parser.add_argument("--concretizing", action="store_true", help="Whether to use `Concretizing` evolution")
-    parser.add_argument("--abstract", action="store_true", help="Whether to use `Abstract` evolution")
-    parser.add_argument("--formal", action="store_true", help="Whether to use `Formal` evolution")
-    parser.add_argument("--funny", action="store_true", help="Whether to use `Funny` evolution")
-    parser.add_argument("--goofy", action="store_true", help="Whether to use `Goofy` evolution")
-    parser.add_argument("--informal", action="store_true", help="Whether to use `Informal` evolution")
+    parser.add_argument("--decide-for-me", action="store_true", help="Enable incremental evolution")
+    parser.add_argument("--template", type=str, choices=EVOLUTION_NAMES, help="Template to use", nargs="+")
     parser.add_argument("--async-mode", action="store_true", help="Enable asynchronous generation")
     parser.add_argument("--seed", type=int, default=0)
     parser.add_argument("--batch-size", type=int, default=4)
@@ -62,25 +51,8 @@ def _parse_args() -> Namespace:
 
 def main() -> None:
     """CLI endpoint."""
-    mapping = {
-        "reasoning": ReasoningEvolution,
-        "concretizing": ConcreteEvolution,
-        "abstract": AbstractEvolution,
-        "formal": FormalEvolution,
-        "funny": FunnyEvolution,
-        "goofy": GoofyEvolution,
-        "informal": InformalEvolution,
-    }
     args = _parse_args()
-    evolutions = []
-
-    for arg_name, evolution_cls in mapping.items():
-        if getattr(args, arg_name):
-            evolutions.append(evolution_cls())  # type: ignore[abstract]
-
-    if not evolutions:
-        logger.warning("No evolutions selected. Exiting.")
-        return
+    evolutions = [EVOLUTION_MAPPING[template_name] for template_name in args.template]
 
     utterance_evolver: UtteranceEvolver
     if args.decide_for_me:
@@ -103,7 +75,7 @@ def main() -> None:
     dataset.to_json(args.output_path)
 
     if args.output_repo is not None:
-        dataset.push_to_hub(args.output_repo)
+        dataset.push_to_hub(args.output_repo, args.private)
 
 
 if __name__ == "__main__":