тесты для адверсариал аугментации

Author: Tetragrammaton123

src/autointent/generation/utterances/__init__.py

@@ -2,10 +2,13 @@
 
 from ._basic import DatasetBalancer, UtteranceGenerator
 from ._evolution import IncrementalUtteranceEvolver, UtteranceEvolver
+from ._adversarial import HumanUtteranceGenerator, CriticHumanLike
 
 __all__ = [
     "DatasetBalancer",
     "IncrementalUtteranceEvolver",
     "UtteranceEvolver",
     "UtteranceGenerator",
+    "HumanUtteranceGenerator",
+    "CriticHumanLike"
 ]

src/autointent/generation/utterances/_adversarial/__init__.py

@@ -1,4 +1,4 @@
 from .critic_human_like import CriticHumanLike
 from .human_utterance_generator import HumanUtteranceGenerator
 
-__all__ = ["HumanUtteranceGenerator"]
+__all__ = ["HumanUtteranceGenerator", "CriticHumanLike"]

src/autointent/generation/utterances/_adversarial/human_utterance_generator.py

@@ -155,7 +155,8 @@ async def generate_one(intent_id: str, intent_name: str) -> list[dict[str, str]]
 
         for result in results:
             new_samples.extend(result)
-
+        for s in new_samples:
+            s['label'] = int(s['label'])
         if update_split:
             generated_split = HFDataset.from_list(new_samples)
             dataset[split_name] = concatenate_datasets([original_split, generated_split])

tests/generation/utterances/test_adversarial.py

@@ -0,0 +1,59 @@
+from unittest.mock import Mock, AsyncMock
+import pytest
+from autointent.generation import Generator
+
+from autointent.generation.utterances import HumanUtteranceGenerator, CriticHumanLike
+from autointent import Dataset, Sample
+
+
+def test_human_utterance_generator_sync(dataset):
+    mock_llm = Mock()
+    mock_llm.get_chat_completion.return_value = "Human-like utterance"
+    
+    mock_critic = Mock(spec=CriticHumanLike)
+    mock_critic.is_human.return_value = True
+
+    generator = HumanUtteranceGenerator(mock_llm, mock_critic, async_mode=False)
+    
+    n_before = len(dataset["train_0"])
+    new_samples = generator.augment(dataset, split_name="train_0", update_split=False, n_final_per_class=2)
+    n_after = len(dataset["train_0"])
+
+    assert n_before == n_after
+    assert len(new_samples) > 0
+    assert all(isinstance(sample, Sample) for sample in new_samples)
+    assert all("utterance" in sample.dict() for sample in new_samples)
+    assert all("label" in sample.dict() for sample in new_samples)
+
+
+def test_human_utterance_generator_async(dataset):
+    mock_llm = AsyncMock()
+    mock_llm.get_chat_completion_async.return_value = "Human-like utterance"
+    
+    mock_critic = AsyncMock(spec=CriticHumanLike)
+    mock_critic.is_human_async.return_value = True
+    generator = HumanUtteranceGenerator(mock_llm, mock_critic, async_mode=True)
+
+    n_before = len(dataset["train_0"])
+    new_samples = generator.augment(dataset, split_name="train_0", update_split=False, n_final_per_class=2)
+    n_after = len(dataset["train_0"])
+
+    assert n_before == n_after
+    assert len(new_samples) > 0
+    assert all(isinstance(sample, Sample) for sample in new_samples)
+    assert all("utterance" in sample.dict() for sample in new_samples)
+    assert all("label" in sample.dict() for sample in new_samples)
+
+
+def test_human_utterance_generator_respects_critic(dataset):
+    mock_llm = Mock()
+    mock_llm.get_chat_completion.return_value = "Generated utterance"
+
+    mock_critic = Mock(spec=CriticHumanLike)
+    mock_critic.is_human.side_effect = [False, True]
+
+    generator = HumanUtteranceGenerator(mock_llm, mock_critic, async_mode=False)
+
+    new_samples = generator.augment(dataset, split_name="train_0", update_split=False, n_final_per_class=1)
+    assert len(new_samples) > 0
+    assert all(mock_critic.is_human.call_count >= 1 for _ in range(len(new_samples)))