fix few shot split (#219)

* fix few shot split

* lint

Author: Samoed

autointent/context/data_handler/_stratification.py

@@ -72,7 +72,16 @@ def __call__(
             ValueError: If OOS samples are present but allow_oos_in_train is not specified.
         """
         if not self._has_oos_samples(dataset):
-            return self._split_without_oos(dataset, multilabel, self.test_size)
+            train, test = self._split_without_oos(dataset, multilabel, self.test_size)
+            if self.is_few_shot:
+                train, test = create_few_shot_split(
+                    train,
+                    test,
+                    multilabel=multilabel,
+                    label_column=self.label_feature,
+                    examples_per_label=self.examples_per_label,
+                )
+            return train, test
         if allow_oos_in_train is None:
             msg = (
                 "Error while splitting dataset. It contains OOS samples, "

tests/data/test_stratificaiton.py

@@ -63,20 +63,29 @@ def test_multilabel_train_test_split_few_shot(dataset_unsplitted):
     assert dataset.get_n_classes(Split.TRAIN) == dataset.get_n_classes(Split.TEST)
 
 
-def test_multiclass_train_test_split_few_shot(dataset_unsplitted):
+@pytest.mark.parametrize("allow_oos_in_train", [True, False])
+def test_multiclass_train_test_split_few_shot(dataset_unsplitted, allow_oos_in_train):
+    train_num_rows = 10 if allow_oos_in_train else 8
+    test_num_rows = 26 if allow_oos_in_train else 28
+    examples_per_intent = 2
+
     dataset = dataset_unsplitted
     dataset[Split.TRAIN], dataset[Split.TEST] = split_dataset(
         dataset,
         split=Split.TRAIN,
         test_size=0.5,
         random_seed=42,
-        allow_oos_in_train=False,
+        allow_oos_in_train=allow_oos_in_train,
         is_few_shot=True,
-        examples_per_intent=2,
+        examples_per_intent=examples_per_intent,
     )
 
     assert Split.TRAIN in dataset
     assert Split.TEST in dataset
-    assert dataset[Split.TRAIN].num_rows == 8
-    assert dataset[Split.TEST].num_rows == 28
+    assert dataset[Split.TRAIN].num_rows == train_num_rows
+    assert dataset[Split.TEST].num_rows == test_num_rows
     assert dataset.get_n_classes(Split.TRAIN) == dataset.get_n_classes(Split.TEST)
+
+    for class_id in range(dataset.get_n_classes(Split.TRAIN)):
+        class_ds = dataset[Split.TRAIN].filter(lambda x: x["label"] == class_id)  # noqa: B023
+        assert len(class_ds) <= examples_per_intent