diff --git a/autointent/generation/intents/prompt_scheme.py b/autointent/generation/intents/prompt_scheme.py index b0883cff1..026e3ec83 100644 --- 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/utterances/prompts.py b/autointent/generation/intents/prompts.py similarity index 100% rename from autointent/generation/utterances/prompts.py rename to autointent/generation/intents/prompts.py diff --git a/autointent/generation/utterances/__init__.py b/autointent/generation/utterances/__init__.py index db54310ee..7534eb1a0 100644 --- 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 index 27007e7f1..7e352bd86 100644 --- 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 index 8417ff1a6..f6c4e2902 100644 --- 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 index 4ab88c359..20c2da5b2 100644 --- 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 index dcca78bac..4a7ab52f2 100644 --- 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 new file mode 100644 index 000000000..a527826c9 --- /dev/null +++ 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 new file mode 100644 index 000000000..0f401b464 --- /dev/null +++ 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 new file mode 100644 index 000000000..15a6fcb17 --- /dev/null +++ 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 new file mode 100644 index 000000000..7d9243f20 --- /dev/null +++ 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 index f05991ed1..fc271b45a 100644 --- 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 new file mode 100644 index 000000000..e69de29bb diff --git a/tests/generation/test_description_generation.py b/tests/generation/intents/test_description_generation.py similarity index 100% rename from tests/generation/test_description_generation.py rename to tests/generation/intents/test_description_generation.py diff --git a/tests/generation/utterances/__init__.py b/tests/generation/utterances/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/generation/test_basic_synthesizer.py b/tests/generation/utterances/test_basic_synthesizer.py similarity index 100% rename from tests/generation/test_basic_synthesizer.py rename to tests/generation/utterances/test_basic_synthesizer.py diff --git a/tests/generation/test_evolver.py b/tests/generation/utterances/test_evolver.py similarity index 100% rename from tests/generation/test_evolver.py rename to tests/generation/utterances/test_evolver.py diff --git a/tests/generation/utterances/test_generator.py b/tests/generation/utterances/test_generator.py new file mode 100644 index 000000000..667f3dc9c --- /dev/null +++ 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()