Repository: huggingface/disaggregators (language: python)

Issue:
Initial library setup
Expose a `Disaggregator`, with a sample `DisaggregationModule`.


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 718713f..650b352 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,42 @@
 # Disaggregators
 
-TO DO.
+The `disaggregators` library allows you to easily add new features to your datasets to enable disaggregated data exploration and disaggregated model evaluation. `disaggregators` is preloaded with disaggregation modules intended for text and image data.
 
-Contact: nima@huggingface.co
\ No newline at end of file
+`disaggregators` is intended to be used with 🤗 Datasets, but should work with any other "mappable" interface to a dataset. 
+
+# Requirements and Installation
+
+`disaggregators` has been tested on Python 3.10.8.
+
+`pip install disaggregators` will fetch the latest release from PyPI.
+
+To install directly from this GitHub repo, use the following command:
+```shell
+pip install git+https://github.com/huggingface/disaggregators.git
+```
+
+# Usage
+
+You will likely want to use 🤗 Datasets with `disaggregators`.
+
+```shell
+pip install datasets
+```
+
+The snippet below loads the IMDB dataset from the Hugging Face Hub, and initializes a disaggregator for "pronouns" that will run on the IMDB dataset's "text" column. Note that if you would like to run multiple disaggregations, you can pass a list to the `Disaggregator` constructor (e.g. `Disaggregator(["pronouns", "sentiment"], column="text")`). We then use the 🤗 Datasets `map` method to apply the disaggregation to the dataset.
+
+```python
+from disaggregators import Disaggregator
+from datasets import load_dataset
+
+dataset = load_dataset("imdb", split="train")
+disaggregator = Disaggregator("pronouns", column="text")
+
+ds = dataset.map(disaggregator.get_function())  # New boolean columns are added for she/her, he/him, and they/them
+```
+
+The resulting dataset can now be used for data exploration and disaggregated model evaluation.
+
+# Contact
+
+Nima Boscarino – `nima <at> huggingface <dot> co`
\ No newline at end of file
diff --git a/src/disaggregators/__init__.py b/src/disaggregators/__init__.py
index 189209c..a36c0fa 100644
--- a/src/disaggregators/__init__.py
+++ b/src/disaggregators/__init__.py
@@ -21,6 +21,10 @@
 
 from packaging import version
 
+from disaggregators.disaggregation_modules import DisaggregationModule, DisaggregationModuleFactory
+
+from .disaggregator import Disaggregator
+
 
 SCRIPTS_VERSION = "main" if version.parse(__version__).is_devrelease else __version__
 
diff --git a/src/disaggregators/disaggregation_modules/__init__.py b/src/disaggregators/disaggregation_modules/__init__.py
new file mode 100644
index 0000000..27f0d69
--- /dev/null
+++ b/src/disaggregators/disaggregation_modules/__init__.py
@@ -0,0 +1,7 @@
+from .disaggregation_module import DisaggregationModule, DisaggregationModuleFactory
+from .pronouns import Pronouns
+
+
+AVAILABLE_MODULES = {"pronouns": Pronouns}
+
+__all__ = ["DisaggregationModule", "DisaggregationModuleFactory"]
diff --git a/src/disaggregators/disaggregation_modules/disaggregation_module.py b/src/disaggregators/disaggregation_modules/disaggregation_module.py
new file mode 100644
index 0000000..c579e9f
--- /dev/null
+++ b/src/disaggregators/disaggregation_modules/disaggregation_module.py
@@ -0,0 +1,22 @@
+from abc import ABC, abstractmethod
+
+from disaggregators import disaggregation_modules
+
+
+class DisaggregationModule(ABC):
+    def __init__(self, module_id: str, column: str):
+        self.name = module_id
+        self.column = column
+
+    @abstractmethod
+    def __call__(self, row, *args, **kwargs):
+        raise NotImplementedError()
+
+
+class DisaggregationModuleFactory:
+    @staticmethod
+    def create_from_id(module_id: str, column: str) -> DisaggregationModule:
+        if module_id not in disaggregation_modules.AVAILABLE_MODULES:
+            raise ValueError("Invalid module_id received.")
+
+        return disaggregation_modules.AVAILABLE_MODULES[module_id](column=column)
diff --git a/src/disaggregators/disaggregation_modules/pronouns/__init__.py b/src/disaggregators/disaggregation_modules/pronouns/__init__.py
new file mode 100644
index 0000000..0121dd6
--- /dev/null
+++ b/src/disaggregators/disaggregation_modules/pronouns/__init__.py
@@ -0,0 +1,19 @@
+from ..disaggregation_module import DisaggregationModule
+
+
+class Pronouns(DisaggregationModule):
+    def __init__(self, *args, **kwargs):
+        super().__init__(module_id="pronouns", *args, **kwargs)
+        self.AVAILABLE_PRONOUNS = {
+            "she/her": {"she", "her", "hers", "herself"},
+            "he/him": {"he", "him", "his", "himself"},
+            "they/them": {"they", "them", "their", "theirs", "themself", "themselves"},
+        }
+
+    def __call__(self, row, *args, **kwargs):
+        text = row[self.column]
+        pronoun_flag = {
+            av_p: any(p in text.lower() for p in self.AVAILABLE_PRONOUNS[av_p]) for av_p in self.AVAILABLE_PRONOUNS
+        }
+
+        return pronoun_flag
diff --git a/src/disaggregators/disaggregator.py b/src/disaggregators/disaggregator.py
new file mode 100644
index 0000000..3a80800
--- /dev/null
+++ b/src/disaggregators/disaggregator.py
@@ -0,0 +1,22 @@
+from typing import Callable, Optional, Union
+
+from disaggregators.disaggregation_modules import DisaggregationModuleFactory
+
+
+class Disaggregator:
+    def __init__(self, module_ids: Optional[Union[str, list[str]]] = None, column: str = None):
+        if module_ids is None:
+            module_ids = []
+
+        if isinstance(module_ids, str):
+            self.modules = [DisaggregationModuleFactory.create_from_id(module_ids, column=column)]
+        elif isinstance(module_ids, list):
+            self.modules = [
+                DisaggregationModuleFactory.create_from_id(module_id, column=column) for module_id in module_ids
+            ]
+        else:
+            raise ValueError("Invalid argument passed for module")
+
+    def get_function(self) -> Callable:
+        # Merge dicts - https://stackoverflow.com/a/3495395
+        return lambda x: {k: v for d in [module(x) for module in self.modules] for k, v in d.items()}


Test patch:
diff --git a/tests/disaggregation_modules/test_pronouns.py b/tests/disaggregation_modules/test_pronouns.py
new file mode 100644
index 0000000..07aac78
--- /dev/null
+++ b/tests/disaggregation_modules/test_pronouns.py
@@ -0,0 +1,19 @@
+import unittest
+
+from disaggregators.disaggregation_modules.pronouns import Pronouns
+
+
+class TestPronouns(unittest.TestCase):
+    def test_initialize(self):
+        disagg_module = Pronouns(column=None)
+        self.assertEqual(disagg_module.name, "pronouns")
+
+    def test_call_default_pronouns(self):
+        data = {"text": "He went to the park."}
+        disagg_module = Pronouns(column="text")
+        results = disagg_module(data)
+        self.assertEqual(results, {"he/him": True, "she/her": False, "they/them": False})
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/sample_test.py b/tests/sample_test.py
deleted file mode 100644
index b475fb4..0000000
--- a/tests/sample_test.py
+++ /dev/null
@@ -1,6 +0,0 @@
-import unittest
-
-
-class TestSample(unittest.TestCase):
-    def test_sample(self):
-        self.assertTrue(True)
diff --git a/tests/test_disaggregation_module.py b/tests/test_disaggregation_module.py
new file mode 100644
index 0000000..3ee658d
--- /dev/null
+++ b/tests/test_disaggregation_module.py
@@ -0,0 +1,37 @@
+import unittest
+from unittest.mock import patch
+
+from disaggregators import DisaggregationModule, DisaggregationModuleFactory
+
+
+class DummyModule(DisaggregationModule):
+    def __init__(self, module_id="dummy-module", column=None):
+        super().__init__(module_id=module_id, column=column)
+
+    def __call__(self, row, *args, **kwargs):
+        return {self.name: True}
+
+
+class TestModule(unittest.TestCase):
+    def test_create_subclassed_module(self):
+        custom_module = DummyModule()
+        self.assertEqual(custom_module.name, "dummy-module")
+
+
+@patch("disaggregators.disaggregation_modules.disaggregation_module.disaggregation_modules")
+class TestModuleCreator(unittest.TestCase):
+    def test_load_module_from_string_id(self, mock_disaggregation_modules):
+        mock_disaggregation_modules.AVAILABLE_MODULES = {"dummy-module": DummyModule}
+        loaded_module = DisaggregationModuleFactory.create_from_id(module_id="dummy-module", column="")
+        self.assertTrue(issubclass(type(loaded_module), DisaggregationModule))
+        self.assertEqual(loaded_module.name, "dummy-module")
+
+    def test_load_module_with_bad_string_id(self, mock_disaggregation_modules):
+        mock_disaggregation_modules.AVAILABLE_MODULES = {}
+
+        with self.assertRaises(ValueError):
+            DisaggregationModuleFactory.create_from_id(module_id="bad-module", column="")
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_disaggregator.py b/tests/test_disaggregator.py
new file mode 100644
index 0000000..f6f542c
--- /dev/null
+++ b/tests/test_disaggregator.py
@@ -0,0 +1,49 @@
+import unittest
+from typing import Callable
+from unittest.mock import MagicMock, patch
+
+from disaggregators import DisaggregationModule, Disaggregator
+from tests.test_disaggregation_module import DummyModule
+
+
+MockDisaggregationModuleFactory = MagicMock()
+MockDisaggregationModuleFactory.create_from_id.side_effect = lambda module_id=None, column=None: DummyModule(
+    module_id=module_id, column=column
+)
+
+
+@patch("disaggregators.disaggregator.DisaggregationModuleFactory", MockDisaggregationModuleFactory)
+class TestDisaggregator(unittest.TestCase):
+    def test_create_empty_disaggregator(self):
+        disagg = Disaggregator()
+        self.assertIsInstance(disagg.modules, list)
+        self.assertEqual(len(disagg.modules), 0)
+
+    def test_create_disaggregator_with_single_module(self):
+        disagg = Disaggregator("pronouns")
+        self.assertEqual(len(disagg.modules), 1)
+        self.assertIsInstance(disagg.modules[0], DisaggregationModule)
+        self.assertEqual(disagg.modules[0].name, "pronouns")
+
+    def test_create_disaggregator_with_multiple_modules(self):
+        disagg = Disaggregator(["pronouns", "spelling"])
+        self.assertEqual(len(disagg.modules), 2)
+        self.assertTrue(all([isinstance(module, DisaggregationModule) for module in disagg.modules]))
+        self.assertEqual(disagg.modules[0].name, "pronouns")
+        self.assertEqual(disagg.modules[1].name, "spelling")
+
+    def test_get_disaggregator_function_single_aggregation_module(self):
+        disagg = Disaggregator("dummy-module")
+        disagg_func = disagg.get_function()
+        self.assertIsInstance(disagg_func, Callable)
+        self.assertEqual(disagg_func({"a": 1, "b": 2}), {"dummy-module": True})
+
+    def test_get_disaggregator_function_multiple_aggregation_modules(self):
+        disagg = Disaggregator(["dummy-one", "dummy-two"])
+        disagg_func = disagg.get_function()
+        self.assertIsInstance(disagg_func, Callable)
+        self.assertEqual(disagg_func({"a": 1, "b": 2}), {"dummy-one": True, "dummy-two": True})
+
+
+if __name__ == "__main__":
+    unittest.main()