Feat/augmentation (#107)

Darinochka · web-flow · commit 15016a57a515 · 2025-02-04T21:46:20.000+03:00
* fix: change location of prompts

* feat: update evolution class

* feat: update chat templates

* feat: added new evolutions

* feat: change test structure &amp; add test_generator

* fix: fixex ruff

* fix: fixed import

* fix: fixed chat_templates
diff --git a/autointent/generation/intents/prompt_scheme.py b/autointent/generation/intents/prompt_scheme.py
@@ -2,7 +2,7 @@
 
 from pydantic import BaseModel, field_validator
 
-from autointent.generation.utterances.prompts import PROMPT_DESCRIPTION
+from autointent.generation.intents.prompts import PROMPT_DESCRIPTION
 
 
 class PromptDescription(BaseModel):
diff --git a/autointent/generation/intents/prompts.py b/autointent/generation/intents/prompts.py
diff --git a/autointent/generation/utterances/__init__.py b/autointent/generation/utterances/__init__.py
@@ -1,14 +1,28 @@
 from .basic import SynthesizerChatTemplate, UtteranceGenerator
-from .evolution import AbstractEvolution, ConcreteEvolution, EvolutionChatTemplate, ReasoningEvolution, UtteranceEvolver
+from .evolution import (
+                        AbstractEvolution,
+                        ConcreteEvolution,
+                        EvolutionChatTemplate,
+                        FormalEvolution,
+                        FunnyEvolution,
+                        GoofyEvolution,
+                        InformalEvolution,
+                        ReasoningEvolution,
+                        UtteranceEvolver,
+)
 from .generator import Generator
 
 __all__ = [
-    "AbstractEvolution",
-    "ConcreteEvolution",
-    "EvolutionChatTemplate",
-    "Generator",
-    "ReasoningEvolution",
-    "SynthesizerChatTemplate",
-    "UtteranceEvolver",
-    "UtteranceGenerator",
+                        "AbstractEvolution",
+                        "ConcreteEvolution",
+                        "EvolutionChatTemplate",
+                        "FormalEvolution",
+                        "FunnyEvolution",
+                        "Generator",
+                        "GoofyEvolution",
+                        "InformalEvolution",
+                        "ReasoningEvolution",
+                        "SynthesizerChatTemplate",
+                        "UtteranceEvolver",
+                        "UtteranceGenerator",
 ]
diff --git a/autointent/generation/utterances/evolution/__init__.py b/autointent/generation/utterances/evolution/__init__.py
@@ -1,4 +1,23 @@
-from .chat_templates import AbstractEvolution, ConcreteEvolution, EvolutionChatTemplate, ReasoningEvolution
+from .chat_templates import (
+                             AbstractEvolution,
+                             ConcreteEvolution,
+                             EvolutionChatTemplate,
+                             FormalEvolution,
+                             FunnyEvolution,
+                             GoofyEvolution,
+                             InformalEvolution,
+                             ReasoningEvolution,
+)
 from .evolver import UtteranceEvolver
 
-__all__ = ["AbstractEvolution", "ConcreteEvolution", "EvolutionChatTemplate", "ReasoningEvolution", "UtteranceEvolver"]
+__all__ = [
+                             "AbstractEvolution",
+                             "ConcreteEvolution",
+                             "EvolutionChatTemplate",
+                             "FormalEvolution",
+                             "FunnyEvolution",
+                             "GoofyEvolution",
+                             "InformalEvolution",
+                             "ReasoningEvolution",
+                             "UtteranceEvolver",
+]
diff --git a/autointent/generation/utterances/evolution/chat_templates/__init__.py b/autointent/generation/utterances/evolution/chat_templates/__init__.py
@@ -1,6 +1,19 @@
 from .abstract import AbstractEvolution
 from .base import EvolutionChatTemplate
 from .concrete import ConcreteEvolution
+from .formal import FormalEvolution
+from .funny import FunnyEvolution
+from .goofy import GoofyEvolution
+from .informal import InformalEvolution
 from .reasoning import ReasoningEvolution
 
-__all__ = ["AbstractEvolution", "ConcreteEvolution", "EvolutionChatTemplate", "ReasoningEvolution"]
+__all__ = [
+    "AbstractEvolution",
+    "ConcreteEvolution",
+    "EvolutionChatTemplate",
+    "FormalEvolution",
+    "FunnyEvolution",
+    "GoofyEvolution",
+    "InformalEvolution",
+    "ReasoningEvolution",
+]
diff --git a/autointent/generation/utterances/evolution/chat_templates/abstract.py b/autointent/generation/utterances/evolution/chat_templates/abstract.py
@@ -28,9 +28,9 @@ class AbstractEvolution(EvolutionChatTemplate):
         ),
         Message(role=Role.ASSISTANT, content="Please, reserve a table for me."),
         Message(
-            role=Role.ASSISTANT,
+            role=Role.USER,
             content=(
-                "Intent name: requesting technical support"
+                "Intent name: requesting technical support\n"
                 "Utterance: My Lenovo laptop is constantly rebooting and overheating."
             ),
         ),
diff --git a/autointent/generation/utterances/evolution/chat_templates/concrete.py b/autointent/generation/utterances/evolution/chat_templates/concrete.py
@@ -29,7 +29,10 @@ 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\n"
+                "Utterance: I'm having trouble with my laptop."
+            ),
         ),
         Message(role=Role.ASSISTANT, content="My laptop is constantly rebooting and overheating."),
     ]
diff --git a/autointent/generation/utterances/evolution/chat_templates/formal.py b/autointent/generation/utterances/evolution/chat_templates/formal.py
@@ -0,0 +1,48 @@
+"""Chat template for formal tone augmentation."""
+
+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."""
+
+    _messages: ClassVar[list[Message]] = [
+        Message(
+            role=Role.USER,
+            content=(
+                "I want you to act as a rewriter. "
+                "You will be provided with an utterance and the topic (name of intent class) of the utterance. "
+                "You need to rewrite the utterance in a more formal tone using the following method:\n"
+                "1. Rewrite the utterance in a more formal tone.\n"
+                "2. Use polite and professional language while maintaining clarity.\n"
+                "3. The rewritten utterance should be grammatically correct and complete.\n"
+                "4. Keep the rewritten utterance within 15 words.\n\n"
+                "Intent name: Reserve Restaurant"
+                "Utterance: I want to reserve a table for 4 persons at 9 pm."
+            ),
+        ),
+        Message(role=Role.ASSISTANT, content="I would like to make a reservation for four guests at 9 pm."),
+        Message(
+            role=Role.ASSISTANT,
+            content=(
+                "Intent name: requesting technical support\n"
+                "Utterance: My Lenovo laptop is constantly rebooting and overheating."
+            ),
+        ),
+        Message(
+            role=Role.ASSISTANT,
+            content="My Lenovo laptop frequently restarts and experiences overheating issues. Kindly assist.",
+        ),
+    ]
+
+    def __call__(self, utterance: str, intent_data: Intent) -> list[Message]:
+        """Generate chat for formal tone adaptation."""
+        return [
+            *self._messages,
+            Message(role=Role.USER, content=f"Intent name: {intent_data.name or ''}\nUtterance: {utterance}"),
+        ]
diff --git a/autointent/generation/utterances/evolution/chat_templates/funny.py b/autointent/generation/utterances/evolution/chat_templates/funny.py
@@ -0,0 +1,46 @@
+"""Chat template for humorous tone augmentation."""
+
+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."""
+
+    _messages: ClassVar[list[Message]] = [
+        Message(
+            role=Role.USER,
+            content=(
+                "I want you to act as a rewriter. "
+                "You will be provided with an utterance and the topic (name of intent class) of the utterance. "
+                "You need to rewrite the utterance in a humorous way while maintaining its original meaning using "
+                "the following method:\n"
+                "1. Rewrite the utterance in a humorous way while maintaining its original meaning.\n"
+                "2. Use wordplay, exaggeration, or lighthearted phrasing.\n"
+                "3. The rewritten utterance should still be understandable and relevant.\n"
+                "4. Keep it within 15 words.\n\n"
+                "Intent name: Reserve Restaurant"
+                "Utterance: I want to reserve a table for 4 persons at 9 pm."
+            ),
+        ),
+        Message(role=Role.ASSISTANT, content="Gotta feed my squad at 9 pm. Got a table for us?"),
+        Message(
+            role=Role.USER,
+            content=(
+                "Intent name: requesting technical support\n"
+                "Utterance: My Lenovo laptop is constantly rebooting and overheating."
+            ),
+        ),
+        Message(role=Role.ASSISTANT, content="My Lenovo thinks it's a phoenix—keeps dying and rising in flames."),
+    ]
+
+    def __call__(self, utterance: str, intent_data: Intent) -> list[Message]:
+        """Generate chat for humorous tone adaptation."""
+        return [
+            *self._messages,
+            Message(role=Role.USER, content=f"Intent name: {intent_data.name or ''}\nUtterance: {utterance}"),
+        ]
diff --git a/autointent/generation/utterances/evolution/chat_templates/goofy.py b/autointent/generation/utterances/evolution/chat_templates/goofy.py
@@ -0,0 +1,49 @@
+"""Chat template for goofy tone augmentation."""
+
+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."""
+
+    _messages: ClassVar[list[Message]] = [
+        Message(
+            role=Role.USER,
+            content=(
+                "I want you to act as a rewriter. "
+                "You will be provided with an utterance and the topic (name of intent class) of the utterance. "
+                "You need to rewrite the utterance in a silly, exaggerated, or nonsensical way while keeping "
+                "the intent clear using the following method:\n"
+                "1. Rewrite the utterance in a silly, exaggerated, or nonsensical way while keeping the intent clear.\n"
+                "2. Use playful words, randomness, or exaggeration.\n"
+                "3. The rewritten utterance should still be answerable.\n"
+                "4. Keep it under 15 words.\n\n"
+                "Intent name: Reserve Restaurant"
+                "Utterance: I want to reserve a table for 4 persons at 9 pm."
+            ),
+        ),
+        Message(role=Role.ASSISTANT, content="Need a feast for my hungry goblins at 9. Got room?"),
+        Message(
+            role=Role.USER,
+            content=(
+                "Intent name: requesting technical support\n"
+                "Utterance: My Lenovo laptop is constantly rebooting and overheating."
+            ),
+        ),
+        Message(
+            role=Role.ASSISTANT,
+            content="My laptop's having an existential crisis—keeps rebooting and melting. Help!"
+        ),
+    ]
+
+    def __call__(self, utterance: str, intent_data: Intent) -> list[Message]:
+        """Generate chat for goofy tone adaptation."""
+        return [
+            *self._messages,
+            Message(role=Role.USER, content=f"Intent name: {intent_data.name or ''}\nUtterance: {utterance}"),
+        ]
diff --git a/autointent/generation/utterances/evolution/chat_templates/informal.py b/autointent/generation/utterances/evolution/chat_templates/informal.py
@@ -0,0 +1,45 @@
+"""Chat template for informal tone augmentation."""
+
+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."""
+
+    _messages: ClassVar[list[Message]] = [
+        Message(
+            role=Role.USER,
+            content=(
+                "I want you to act as a rewriter. "
+                "You will be provided with an utterance and the topic (name of intent class) of the utterance. "
+                "You need to rewrite the utterance in a more casual and relaxed tone using the following method:\n"
+                "1. Rewrite the utterance in a more casual and relaxed tone.\n"
+                "2. Use contractions, friendly language, and a conversational style.\n"
+                "3. The rewritten utterance should feel natural in an informal conversation.\n"
+                "4. Keep it under 15 words.\n\n"
+                "Intent name: Reserve Restaurant"
+                "Utterance: I want to reserve a table for 4 persons at 9 pm."
+            ),
+        ),
+        Message(role=Role.ASSISTANT, content="Hey, can I book a table for 4 at 9?"),
+        Message(
+            role=Role.USER,
+            content=(
+                "Intent name: requesting technical support\n"
+                "Utterance: My Lenovo laptop is constantly rebooting and overheating."
+            ),
+        ),
+        Message(role=Role.ASSISTANT, content="My Lenovo keeps crashing and getting super hot. Any ideas?"),
+    ]
+
+    def __call__(self, utterance: str, intent_data: Intent) -> list[Message]:
+        """Generate chat for informal tone adaptation."""
+        return [
+            *self._messages,
+            Message(role=Role.USER, content=f"Intent name: {intent_data.name or ''}\nUtterance: {utterance}"),
+        ]
diff --git a/autointent/generation/utterances/evolution/cli.py b/autointent/generation/utterances/evolution/cli.py
@@ -7,7 +7,16 @@
 from autointent.generation.utterances.evolution.evolver import UtteranceEvolver
 from autointent.generation.utterances.generator import Generator
 
-from .chat_templates import AbstractEvolution, ConcreteEvolution, EvolutionChatTemplate, ReasoningEvolution
+from .chat_templates import (
+    AbstractEvolution,
+    ConcreteEvolution,
+    EvolutionChatTemplate,
+    FormalEvolution,
+    FunnyEvolution,
+    GoofyEvolution,
+    InformalEvolution,
+    ReasoningEvolution,
+)
 
 logging.basicConfig(level="INFO")
 logger = logging.getLogger(__name__)
@@ -40,7 +49,12 @@ def main() -> None:
     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("--seed", type=int, default=0)
+
     args = parser.parse_args()
 
     evolutions: list[EvolutionChatTemplate] = []
@@ -50,6 +64,18 @@ def main() -> None:
         evolutions.append(ConcreteEvolution())
     if args.abstract:
         evolutions.append(AbstractEvolution())
+    if args.formal:
+        evolutions.append(FormalEvolution())
+    if args.funny:
+        evolutions.append(FunnyEvolution())
+    if args.goofy:
+        evolutions.append(GoofyEvolution())
+    if args.informal:
+        evolutions.append(InformalEvolution())
+
+    if not evolutions:
+        logger.warning("No evolutions selected. Exiting.")
+        return
 
     dataset = load_dataset(args.input_path)
     n_before = len(dataset[args.split])
diff --git a/tests/generation/intents/__init__.py b/tests/generation/intents/__init__.py
diff --git a/tests/generation/intents/test_description_generation.py b/tests/generation/intents/test_description_generation.py
diff --git a/tests/generation/utterances/__init__.py b/tests/generation/utterances/__init__.py
diff --git a/tests/generation/utterances/test_basic_synthesizer.py b/tests/generation/utterances/test_basic_synthesizer.py
diff --git a/tests/generation/utterances/test_evolver.py b/tests/generation/utterances/test_evolver.py
diff --git a/tests/generation/utterances/test_generator.py b/tests/generation/utterances/test_generator.py
@@ -0,0 +1,38 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from autointent.generation.utterances.generator import Generator
+from autointent.generation.utterances.schemas import Message
+
+
+@pytest.fixture(autouse=True)
+def set_env_vars(monkeypatch):
+    monkeypatch.setenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
+    monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
+    monkeypatch.setenv("OPENAI_MODEL_NAME", "gpt-3.5-turbo")
+
+
+@pytest.fixture
+def mock_openai_client():
+    with patch("openai.OpenAI") as mock_client:
+        yield mock_client
+
+
+def test_generator_initialization(mock_openai_client):
+    generator = Generator()
+    assert generator.client == mock_openai_client.return_value
+    assert generator.model_name == "gpt-3.5-turbo"
+
+
+def test_get_chat_completion(mock_openai_client):
+    mock_response = MagicMock()
+    mock_response.choices = [MagicMock(message=MagicMock(content="Test response"))]
+    mock_openai_client.return_value.chat.completions.create.return_value = mock_response
+
+    generator = Generator()
+    messages = [Message(role="user", content="Hello")]
+    response = generator.get_chat_completion(messages)
+
+    assert response == "Test response"
+    mock_openai_client.return_value.chat.completions.create.assert_called_once()
