[DRAFT] added agumentation yielding evenly distributed classes

textattack/augmentation/augmenter.py

@@ -3,6 +3,7 @@
 ===================
 """
 import random
+from collections import Counter
 
 import tqdm
 
@@ -230,6 +231,98 @@ def augment_text_with_ids(self, text_list, id_list, show_progress=True):
             all_id_list.extend([_id] * (1 + len(augmented_texts)))
         return all_text_list, all_id_list
 
+    def augment_text_with_ids_evenly(
+        self,
+        text_list,
+        id_list,
+        additional_examples=0,
+        perfectly_even=True,
+        show_progress=True,
+    ):
+        """Supplements a list of text with more text data so that there is approximately
+        the same number of sentences for each label.
+        Each ID from `id_list` will be represented the same number of times
+        as the most frequent ID plus `additional_examples`.
+        If `perfectly_even` is set to `True`, every ID will be occurring exactly the same number of times (recommended,
+        but slightly slower).
+
+        Returns the augmented text along with the corresponding IDs for
+        each augmented example.
+        """
+        if len(text_list) != len(id_list):
+            raise ValueError("List of text must be same length as list of IDs")
+        if additional_examples < 0:
+            raise ValueError("Additional examples must be non-negative")
+        all_text_list = []
+        all_id_list = []
+        examples_per_id = Counter(id_list)
+        max_examples = max(examples_per_id.values()) + additional_examples
+        diff_per_example = {k: max_examples - v for k, v in examples_per_id.items()}
+        original_transformations_per_example = self.transformations_per_example
+        remainders = {}
+        if show_progress:
+            text_list = tqdm.tqdm(text_list, desc="Augmenting data...")
+        for text, _id in zip(text_list, id_list):
+            # distribute augmentation of the original documents evenly
+            self.transformations_per_example = (
+                diff_per_example[_id] // examples_per_id[_id]
+            )
+            remainders[_id] = diff_per_example[_id] % examples_per_id[_id]
+            all_text_list.append(text)
+            all_id_list.append(_id)
+            if self.transformations_per_example > 0:
+                augmented_texts = []
+                while len(augmented_texts) < self.transformations_per_example:
+                    augmented_texts.extend(self.augment(text))
+                all_text_list.extend(augmented_texts)
+                all_id_list.extend([_id] * len(augmented_texts))
+
+        if perfectly_even:
+            self.transformations_per_example = 1
+            # (1) add missing examples:
+            if show_progress:
+                added = tqdm.tqdm(
+                    desc="Adding additional examples...", total=sum(remainders.values())
+                )
+            while any(remainders.values()):
+                for text, _id in zip(text_list, id_list):
+                    if remainders[_id] > 0:
+                        # add missing elements one-by-one
+                        remainders[_id] -= 1
+                        if show_progress:
+                            added.update(1)
+                        augmented_texts = self.augment(text)
+                        all_text_list.extend(augmented_texts)
+                        all_id_list.append(_id)
+            if show_progress:
+                added.close()
+            # (2) remove excess:
+            excess = {k: v - max_examples for k, v in Counter(all_id_list).items()}
+            new_id_list = []
+            new_text_list = []
+            if show_progress:
+                to_be_removed = int(sum([e > 0 for e in excess.values()]))
+                removed = tqdm.tqdm(
+                    desc="Removing abundant examples...", total=to_be_removed
+                )
+            # count backwards so that the newer elements (most probably being augmented) are deleted first
+            for i in range(len(all_id_list) - 1, -1, -1):
+                if excess[all_id_list[i]] <= 0:
+                    new_id_list.append(all_id_list[i])
+                    new_text_list.append(all_text_list[i])
+                else:
+                    # skip entry for new id and text list
+                    excess[all_id_list[i]] -= 1
+                    if show_progress:
+                        removed.update(1)
+            all_id_list = new_id_list
+            all_text_list = new_text_list
+            if show_progress:
+                removed.close()
+
+        self.transformations_per_example = original_transformations_per_example
+        return all_text_list, all_id_list
+
     def __repr__(self):
         main_str = "Augmenter" + "("
         lines = []