Repository: huggingface/disaggregators (language: python)

Issue:
Add a basic suite of text disaggregators
Characteristic to consider include:

- [x] Age (#13)
- [x] ~~Ability Status~~
- [x] Gender (NOT pronouns) (#14)
- [x] ~~Ethnicity~~
- [x] Creed/religion (#22)

It's important for these to be implemented following best-practices from existing literature. Note that some of these may not lend themselves well to text data, in which case they might be implemented for image data instead (#6). There may also be some cases where both an image and text version for a particular characteristic will be good to have.

Would be neat to have options for rule-based disaggregators (e.g. word presence for pronouns) and ML-powered ones (e.g. using NLI to find pronouns). Maybe for overlapping ones they could exist as configurations within the same disaggregator module? It's important though that any ML-powered ones use _optional_ extra dependencies.


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 8653d0f..da7e91c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -51,3 +51,5 @@ venv.bak/
 # Sphinx documentation
 docs/_build/
 docs/source/_build/
+
+notebooks/
\ No newline at end of file
diff --git a/Makefile b/Makefile
index abbb2e3..12acb97 100644
--- a/Makefile
+++ b/Makefile
@@ -17,3 +17,6 @@ style:
 
 test:
 	python -m pytest -n auto --dist=loadfile -s -v ./tests/
+
+test-fast:
+	python -m pytest -m "not slow" -n auto --dist=loadfile -s -v ./tests/
\ No newline at end of file
diff --git a/README.md b/README.md
index 650b352..f95e02f 100644
--- a/README.md
+++ b/README.md
@@ -10,6 +10,8 @@ The `disaggregators` library allows you to easily add new features to your datas
 
 `pip install disaggregators` will fetch the latest release from PyPI.
 
+Note that some disaggregation modules require extra dependencies such as SpaCy modules, which may need to be installed manually. If these dependencies aren't installed, `disaggregators` will inform you about how to install them.
+
 To install directly from this GitHub repo, use the following command:
 ```shell
 pip install git+https://github.com/huggingface/disaggregators.git
@@ -23,16 +25,16 @@ You will likely want to use 🤗 Datasets with `disaggregators`.
 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.
+The snippet below loads the IMDB dataset from the Hugging Face Hub, and initializes a disaggregator for "pronoun" 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(["pronoun", "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")
+disaggregator = Disaggregator("pronoun", column="text")
 
-ds = dataset.map(disaggregator.get_function())  # New boolean columns are added for she/her, he/him, and they/them
+ds = dataset.map(disaggregator)  # 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.
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..dff6b32
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,4 @@
+[tool.pytest.ini_options]
+markers = [
+    "slow: marks tests as slow (deselect with '-m \"not slow\"')",
+]
diff --git a/setup.py b/setup.py
index 1edf959..bdc9829 100644
--- a/setup.py
+++ b/setup.py
@@ -58,6 +58,7 @@
 REQUIRED_PKGS = [
     # Utilities from PyPA to e.g., compare versions
     "packaging",
+    "spacy",
 ]
 
 TESTS_REQUIRE = [
diff --git a/src/disaggregators/disaggregation_modules/__init__.py b/src/disaggregators/disaggregation_modules/__init__.py
index aaedf03..2290caf 100644
--- a/src/disaggregators/disaggregation_modules/__init__.py
+++ b/src/disaggregators/disaggregation_modules/__init__.py
@@ -1,7 +1,8 @@
+from .age import Age
 from .disaggregation_module import DisaggregationModule, DisaggregationModuleFactory, DisaggregationModuleLabels
-from .pronouns import Pronouns
+from .pronoun import Pronoun
 
 
-AVAILABLE_MODULES = {"pronouns": Pronouns}
+AVAILABLE_MODULES = {"pronoun": Pronoun, "age": Age}
 
 __all__ = ["DisaggregationModule", "DisaggregationModuleFactory", "DisaggregationModuleLabels"]
diff --git a/src/disaggregators/disaggregation_modules/age/__init__.py b/src/disaggregators/disaggregation_modules/age/__init__.py
new file mode 100644
index 0000000..1a91ac4
--- /dev/null
+++ b/src/disaggregators/disaggregation_modules/age/__init__.py
@@ -0,0 +1,49 @@
+import re
+from bisect import bisect
+
+import spacy
+
+from ..disaggregation_module import DisaggregationModule, DisaggregationModuleLabels
+
+
+class AgeLabels(DisaggregationModuleLabels):
+    CHILD = "child"
+    YOUTH = "youth"
+    ADULT = "adult"
+    SENIOR = "senior"
+
+
+class Age(DisaggregationModule):
+    def __init__(self, *args, **kwargs):
+        super().__init__(module_id="age", labels=AgeLabels, *args, **kwargs)
+        self.AGES = [AgeLabels.CHILD, AgeLabels.YOUTH, AgeLabels.ADULT, AgeLabels.SENIOR]
+        self.AGE_BREAKPOINTS = [0, 12, 20, 65]
+        spacy_model = "en_core_web_lg"
+        try:
+            self.nlp = spacy.load(spacy_model, enable="ner")
+        except OSError:
+            raise ValueError(
+                f"This disaggregation module depends on the {spacy_model} model from spaCy.\n"
+                "You can install it by running: python -m spacy download en_core_web_sm"
+            )
+
+    def __call__(self, row, *args, **kwargs):
+        return_ages = {age: False for age in self.AGES}
+        text = row[self.column]
+        doc = self.nlp(text)
+        date_entities = [y for y in doc.ents if y.label_ == "DATE"]
+
+        for date_entity in date_entities:
+            value = re.search(r"\d+", date_entity.text)
+
+            if value is None:
+                continue
+
+            value = int(value[0])
+
+            if value <= 120:
+                age_bucket = bisect(self.AGE_BREAKPOINTS, value)
+                if age_bucket >= 1:
+                    return_ages.update({self.AGES[age_bucket - 1]: True})
+
+        return return_ages
diff --git a/src/disaggregators/disaggregation_modules/pronouns/__init__.py b/src/disaggregators/disaggregation_modules/pronoun/__init__.py
similarity index 55%
rename from src/disaggregators/disaggregation_modules/pronouns/__init__.py
rename to src/disaggregators/disaggregation_modules/pronoun/__init__.py
index 0a2a085..cdded98 100644
--- a/src/disaggregators/disaggregation_modules/pronouns/__init__.py
+++ b/src/disaggregators/disaggregation_modules/pronoun/__init__.py
@@ -1,19 +1,19 @@
 from ..disaggregation_module import DisaggregationModule, DisaggregationModuleLabels
 
 
-class PronounsLabels(DisaggregationModuleLabels):
+class PronounLabels(DisaggregationModuleLabels):
     SHE_HER = "she_her"
     HE_HIM = "he_him"
     THEY_THEM = "they_them"
 
 
-class Pronouns(DisaggregationModule):
+class Pronoun(DisaggregationModule):
     def __init__(self, *args, **kwargs):
-        super().__init__(module_id="pronouns", labels=PronounsLabels, *args, **kwargs)
+        super().__init__(module_id="pronoun", labels=PronounLabels, *args, **kwargs)
         self.AVAILABLE_PRONOUNS = {
-            PronounsLabels.SHE_HER: {"she", "her", "hers", "herself"},
-            PronounsLabels.HE_HIM: {"he", "him", "his", "himself"},
-            PronounsLabels.THEY_THEM: {"they", "them", "their", "theirs", "themself", "themselves"},
+            PronounLabels.SHE_HER: {"she", "her", "hers", "herself"},
+            PronounLabels.HE_HIM: {"he", "him", "his", "himself"},
+            PronounLabels.THEY_THEM: {"they", "them", "their", "theirs", "themself", "themselves"},
         }
 
     def __call__(self, row, *args, **kwargs):


Test patch:
diff --git a/tests/disaggregation_modules/test_age.py b/tests/disaggregation_modules/test_age.py
new file mode 100644
index 0000000..ffedd04
--- /dev/null
+++ b/tests/disaggregation_modules/test_age.py
@@ -0,0 +1,30 @@
+import pytest
+
+from disaggregators.disaggregation_modules.age import Age, AgeLabels
+
+
+def test_initialize():
+    disagg_module = Age(column=None)
+    assert disagg_module.name == "age"
+    assert disagg_module.labels == {AgeLabels.CHILD, AgeLabels.YOUTH, AgeLabels.ADULT, AgeLabels.SENIOR}
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize(
+    "text,expected",
+    [
+        ("The man went to the park.", []),
+        ("The 40-year-old man went to the park.", [AgeLabels.ADULT]),
+        ("The 10 year old child ate a hot dog.", [AgeLabels.CHILD]),
+        ("Clara is 18 years old.", [AgeLabels.YOUTH]),
+        ("Hamed's grandpa is an 86 yr old smoker.", [AgeLabels.SENIOR]),
+        ("The 40-year-old man went to the park with his 10 year old daughter.", [AgeLabels.ADULT, AgeLabels.CHILD]),
+        ("Farzaneh is 18 years old and her grandmother is 86 yrs old", [AgeLabels.YOUTH, AgeLabels.SENIOR]),
+    ],
+)
+def test_call_default_pronouns(text, expected):
+    base_labels = {age: False for age in AgeLabels}
+    data = {"text": text}
+    disagg_module = Age(column="text")
+    results = disagg_module(data)
+    assert results == {**base_labels, **{label: True for label in expected}}
diff --git a/tests/disaggregation_modules/test_pronoun.py b/tests/disaggregation_modules/test_pronoun.py
new file mode 100644
index 0000000..84c5460
--- /dev/null
+++ b/tests/disaggregation_modules/test_pronoun.py
@@ -0,0 +1,14 @@
+from disaggregators.disaggregation_modules.pronoun import Pronoun, PronounLabels
+
+
+def test_initialize():
+    disagg_module = Pronoun(column=None)
+    assert disagg_module.name == "pronoun"
+    assert disagg_module.labels == {PronounLabels.HE_HIM, PronounLabels.SHE_HER, PronounLabels.THEY_THEM}
+
+
+def test_call_default_pronouns():
+    data = {"text": "He went to the park."}
+    disagg_module = Pronoun(column="text")
+    results = disagg_module(data)
+    assert results == {PronounLabels.HE_HIM: True, PronounLabels.SHE_HER: False, PronounLabels.THEY_THEM: False}
diff --git a/tests/disaggregation_modules/test_pronouns.py b/tests/disaggregation_modules/test_pronouns.py
deleted file mode 100644
index 592deab..0000000
--- a/tests/disaggregation_modules/test_pronouns.py
+++ /dev/null
@@ -1,14 +0,0 @@
-from disaggregators.disaggregation_modules.pronouns import Pronouns, PronounsLabels
-
-
-def test_initialize():
-    disagg_module = Pronouns(column=None)
-    assert disagg_module.name == "pronouns"
-    assert disagg_module.labels == {PronounsLabels.HE_HIM, PronounsLabels.SHE_HER, PronounsLabels.THEY_THEM}
-
-
-def test_call_default_pronouns():
-    data = {"text": "He went to the park."}
-    disagg_module = Pronouns(column="text")
-    results = disagg_module(data)
-    assert results == {PronounsLabels.HE_HIM: True, PronounsLabels.SHE_HER: False, PronounsLabels.THEY_THEM: False}
diff --git a/tests/test_disaggregator.py b/tests/test_disaggregator.py
index 5fec374..abba136 100644
--- a/tests/test_disaggregator.py
+++ b/tests/test_disaggregator.py
@@ -22,16 +22,16 @@ def test_create_empty_disaggregator(self):
         assert len(disagg.modules) == 0
 
     def test_create_disaggregator_with_single_module(self):
-        disagg = Disaggregator("pronouns")
+        disagg = Disaggregator("pronoun")
         assert len(disagg.modules) == 1
         assert isinstance(disagg.modules[0], DisaggregationModule)
-        assert disagg.modules[0].name == "pronouns"
+        assert disagg.modules[0].name == "pronoun"
 
     def test_create_disaggregator_with_multiple_modules(self):
-        disagg = Disaggregator(["pronouns", "spelling"])
+        disagg = Disaggregator(["pronoun", "spelling"])
         assert len(disagg.modules) == 2
         assert all([isinstance(module, DisaggregationModule) for module in disagg.modules])
-        assert disagg.modules[0].name == "pronouns"
+        assert disagg.modules[0].name == "pronoun"
         assert disagg.modules[1].name == "spelling"
 
     def test_get_disaggregator_function_single_aggregation_module(self):