supports custom proportions for augument

nlptest/augmentation/__init__.py

@@ -66,7 +66,7 @@ class AugmentRobustness(BaseAugmentaion):
 
     """
 
-    def __init__(self, task, h_report, config, max_prop=0.5) -> None:
+    def __init__(self, task, h_report, config, custom_proportions=None, max_prop=0.5) -> None:
         """
         Initializes an instance of MyClass with the specified parameters.
 
@@ -87,6 +87,7 @@ def __init__(self, task, h_report, config, max_prop=0.5) -> None:
         self.config = config
         self.h_report = h_report
         self.max_prop = max_prop
+        self.custom_proportions = custom_proportions
 
         if isinstance(self.config, str):
             with open(self.config) as fread:
@@ -168,6 +169,7 @@ def fix(
     def suggestions(self, report):
         """
         Calculates suggestions for improving test performance based on a given report.
+        Supports for custom proportion values passes in dict or list.
 
         Args:
             report (pandas.DataFrame): A DataFrame containing test results by category and test type,
@@ -183,10 +185,21 @@ def suggestions(self, report):
 
         """
         report['ratio'] = report['pass_rate'] / report['minimum_pass_rate']
-        report['proportion_increase'] = report['ratio'].apply(
-            lambda x: self._proportion_values(x)
-        )
-        return report[~report['proportion_increase'].isna()][['category', 'test_type', 'ratio', 'proportion_increase']]
+
+        if self.custom_proportions and isinstance(self.custom_proportions, dict):
+            report['proportion_increase'] = report['test_type'].map(
+                self.custom_proportions)
+        elif self.custom_proportions and isinstance(self.custom_proportions, list):
+            report['proportion_increase'] = report['ratio'].apply(
+                self._proportion_values)
+            report = report[report['test_type'].isin(self.custom_proportions)]
+        else:
+            report['proportion_increase'] = report['ratio'].apply(
+                self._proportion_values)
+
+        report = report.dropna(subset=['proportion_increase'])[
+            ['category', 'test_type', 'ratio', 'proportion_increase']]
+        return report
 
     def _proportion_values(self, x):
         """

nlptest/nlptest.py

@@ -2,7 +2,7 @@
 import os
 import pickle
 from collections import defaultdict
-from typing import Optional, Union, Any
+from typing import Dict, List, Optional, Union, Any
 import langchain
 
 import pandas as pd
@@ -312,7 +312,7 @@ def generated_results(self) -> Optional[pd.DataFrame]:
 
         return generated_results_df.fillna("-")
 
-    def augment(self, input_path: str, output_path: str, inplace: bool = False) -> "Harness":
+    def augment(self, input_path: str, output_path: str, custom_proportions: Union[Dict, List] = None, inplace: bool = False) -> "Harness":
         """
         Augments the data in the input file located at `input_path` and saves the result to `output_path`.
 
@@ -346,7 +346,8 @@ def augment(self, input_path: str, output_path: str, inplace: bool = False) -> "
         _ = AugmentRobustness(
             task=self.task,
             config=self._config,
-            h_report=self.df_report
+            h_report=self.df_report,
+            custom_proportions=custom_proportions
         ).fix(
             input_path=input_path,
             output_path=output_path,

tests/test_augmentation.py

@@ -74,7 +74,6 @@ def test_augment_robustness(self):
             'pass': [True, True, False, False]
         })
 
-
         augment = AugmentRobustness(
             task='ner',
             h_report=temp_df,
@@ -113,3 +112,20 @@ def test_spacy_ner_augmentation(self):
         is_file_exist = pl.Path(
             'tests/fixtures/augmentated_train.conll').is_file()
         self.assertTrue(is_file_exist)
+
+    def test_custom_proportions_augment_harness(self):
+        """"""
+        harness = Harness(**self.params['huggingface_ner'])
+        self.assertIsInstance(harness, Harness)
+        report = harness.generate().run().report()
+        self.assertIsInstance(report, pd.DataFrame)
+
+        proportions = {'uppercase': 0.5, 'lowercase': 0.5}
+
+        harness.augment('tests/fixtures/train.conll',
+                        'tests/fixtures/augmentated_train.conll',
+                        custom_proportions=proportions, inplace=True)
+
+        is_file_exist = pl.Path(
+            'tests/fixtures/augmentated_train.conll').is_file()
+        self.assertTrue(is_file_exist)