Merge pull request #49 from DataFog/feature/add-char-length-param

_chunk_text + tests

Author: sidmohan0

datafog/services/text_service.py

@@ -1,28 +1,61 @@
 import asyncio
+from typing import Dict, List
 
 from datafog.processing.text_processing.spacy_pii_annotator import SpacyPIIAnnotator
 
 
 class TextService:
-    def __init__(self):
+    def __init__(self, text_chunk_length: int = 1000):
         self.annotator = SpacyPIIAnnotator.create()
+        self.text_chunk_length = text_chunk_length
 
-    def annotate_text_sync(self, text):
-        """Synchronously Annotate a text string."""
-        res = self.annotator.annotate(text)
-        return res
+    def _chunk_text(self, text: str) -> List[str]:
+        """Split the text into chunks of specified length."""
+        return [
+            text[i : i + self.text_chunk_length]
+            for i in range(0, len(text), self.text_chunk_length)
+        ]
 
-    def batch_annotate_text_sync(self, texts: list):
+    def _combine_annotations(self, annotations: List[Dict]) -> Dict:
+        """Combine annotations from multiple chunks."""
+        combined = {}
+        for annotation in annotations:
+            for key, value in annotation.items():
+                if key not in combined:
+                    combined[key] = []
+                combined[key].extend(value)
+        return combined
+
+    def annotate_text_sync(self, text: str) -> Dict:
+        """Synchronously annotate a text string."""
+        if not text:
+            return {}
+        print(f"Starting on {text.split()[0]}")
+        chunks = self._chunk_text(text)
+        annotations = []
+        for chunk in chunks:
+            res = self.annotator.annotate(chunk)
+            annotations.append(res)
+        combined = self._combine_annotations(annotations)
+        print(f"Done processing {text.split()[0]}")
+        return combined
+
+    def batch_annotate_text_sync(self, texts: List[str]) -> Dict[str, Dict]:
         """Synchronously annotate a list of text input."""
         results = [self.annotate_text_sync(text) for text in texts]
         return dict(zip(texts, results, strict=True))
 
-    async def annotate_text_async(self, text):
+    async def annotate_text_async(self, text: str) -> Dict:
         """Asynchronously annotate a text string."""
-        return await asyncio.to_thread(self.annotator.annotate, text)
+        if not text:
+            return {}
+        chunks = self._chunk_text(text)
+        tasks = [asyncio.to_thread(self.annotator.annotate, chunk) for chunk in chunks]
+        annotations = await asyncio.gather(*tasks)
+        return self._combine_annotations(annotations)
 
-    async def batch_annotate_text_async(self, text: list):
+    async def batch_annotate_text_async(self, texts: List[str]) -> Dict[str, Dict]:
         """Asynchronously annotate a list of text input."""
-        tasks = [self.annotate_text_async(txt) for txt in text]
+        tasks = [self.annotate_text_async(txt) for txt in texts]
         results = await asyncio.gather(*tasks)
-        return dict(zip(text, results, strict=True))
+        return dict(zip(texts, results, strict=True))

error_log.txt

@@ -0,0 +1,58 @@
+============================= test session starts ==============================
+platform darwin -- Python 3.10.1, pytest-8.2.2, pluggy-1.5.0
+rootdir: /Users/sidmohan/Desktop/datafog_local/datafog-python/datafog-python
+plugins: cov-5.0.0, asyncio-0.23.7, anyio-4.4.0
+asyncio: mode=strict
+collected 12 items
+
+tests/test_text_service.py .F..........                                  [100%]
+
+=================================== FAILURES ===================================
+_______________________________ test_chunk_text ________________________________
+
+text_service = <datafog.services.text_service.TextService object at 0x16fa68a60>
+
+    def test_chunk_text(text_service):
+        text = "This is a test sentence for chunking."
+        chunks = text_service._chunk_text(text)
+        assert len(chunks) == 4
+>       assert chunks == ["This is a", "test", "sentence f", "or chunki"]
+E       AssertionError: assert ['This is a '...h', 'unking.'] == ['This is a',..., 'or chunki']
+E         
+E         At index 0 diff: 'This is a ' != 'This is a'
+E         Use -v to get more diff
+
+tests/test_text_service.py:23: AssertionError
+=============================== warnings summary ===============================
+.venv/lib/python3.10/site-packages/spacy/cli/info.py:3
+  /Users/sidmohan/Desktop/datafog_local/datafog-python/datafog-python/.venv/lib/python3.10/site-packages/spacy/cli/info.py:3: DeprecationWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html
+    import pkg_resources
+
+-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
+
+---------- coverage: platform darwin, python 3.10.1-final-0 ----------
+Name                                                           Stmts   Miss  Cover
+----------------------------------------------------------------------------------
+datafog/__about__.py                                               1      0   100%
+datafog/__init__.py                                               11      0   100%
+datafog/config.py                                                  6      0   100%
+datafog/main.py                                                   74     55    26%
+datafog/processing/__init__.py                                     0      0   100%
+datafog/processing/image_processing/__init__.py                    0      0   100%
+datafog/processing/image_processing/donut_processor.py            44     29    34%
+datafog/processing/image_processing/image_downloader.py           17      7    59%
+datafog/processing/image_processing/pytesseract_processor.py       8      1    88%
+datafog/processing/spark_processing/__init__.py                    3      3     0%
+datafog/processing/spark_processing/pyspark_udfs.py               39     39     0%
+datafog/processing/text_processing/__init__.py                     1      0   100%
+datafog/processing/text_processing/spacy_pii_annotator.py         36     20    44%
+datafog/services/__init__.py                                       3      0   100%
+datafog/services/image_service.py                                 26     11    58%
+datafog/services/spark_service.py                                 26     16    38%
+datafog/services/text_service.py                                  43      1    98%
+----------------------------------------------------------------------------------
+TOTAL                                                            338    182    46%
+
+=========================== short test summary info ============================
+FAILED tests/test_text_service.py::test_chunk_text - AssertionError: assert [...
+=================== 1 failed, 11 passed, 1 warning in 2.06s ====================

tests/beta_pypi_test.ipynb

@@ -0,0 +1,117 @@
+from unittest.mock import Mock, patch
+
+import pytest
+
+from datafog.services.text_service import TextService
+
+
+@pytest.fixture
+def mock_annotator():
+    mock = Mock()
+    mock.annotate.return_value = {"PER": ["John Doe"], "ORG": ["Acme Inc"]}
+    return mock
+
+
+@pytest.fixture
+def text_service(mock_annotator):
+    with patch(
+        "datafog.services.text_service.SpacyPIIAnnotator.create",
+        return_value=mock_annotator,
+    ):
+        return TextService(text_chunk_length=10)
+
+
+def test_init(text_service):
+    assert text_service.text_chunk_length == 10
+
+
+def test_chunk_text(text_service):
+    text = "This is a test sentence for chunking."
+    chunks = text_service._chunk_text(text)
+    assert len(chunks) == 4
+    assert chunks == ["This is a ", "test sente", "nce for ch", "unking."]
+
+
+def test_combine_annotations(text_service):
+    annotations = [
+        {"PER": ["John"], "ORG": ["Acme"]},
+        {"PER": ["Doe"], "LOC": ["New York"]},
+    ]
+    combined = text_service._combine_annotations(annotations)
+    assert combined == {"PER": ["John", "Doe"], "ORG": ["Acme"], "LOC": ["New York"]}
+
+
+def test_annotate_text_sync(text_service):
+    result = text_service.annotate_text_sync("John Doe works at Acme Inc")
+    assert result == {
+        "PER": ["John Doe", "John Doe", "John Doe"],
+        "ORG": ["Acme Inc", "Acme Inc", "Acme Inc"],
+    }
+
+
+def test_batch_annotate_text_sync(text_service):
+    texts = ["John Doe", "Acme Inc"]
+    result = text_service.batch_annotate_text_sync(texts)
+    assert result == {
+        "John Doe": {"PER": ["John Doe"], "ORG": ["Acme Inc"]},
+        "Acme Inc": {"PER": ["John Doe"], "ORG": ["Acme Inc"]},
+    }
+
+
+@pytest.mark.asyncio
+async def test_annotate_text_async(text_service):
+    result = await text_service.annotate_text_async("John Doe works at Acme Inc")
+    assert result == {
+        "PER": ["John Doe", "John Doe", "John Doe"],
+        "ORG": ["Acme Inc", "Acme Inc", "Acme Inc"],
+    }
+
+
+@pytest.mark.asyncio
+async def test_batch_annotate_text_async(text_service):
+    texts = ["John Doe", "Acme Inc"]
+    result = await text_service.batch_annotate_text_async(texts)
+    assert result == {
+        "John Doe": {"PER": ["John Doe"], "ORG": ["Acme Inc"]},
+        "Acme Inc": {"PER": ["John Doe"], "ORG": ["Acme Inc"]},
+    }
+
+
+def test_long_text_chunking(text_service):
+    long_text = "John Doe works at Acme Inc. Jane Smith is from New York City."
+    result = text_service.annotate_text_sync(long_text)
+    expected_count = len(text_service._chunk_text(long_text))
+    assert result == {
+        "PER": ["John Doe"] * expected_count,
+        "ORG": ["Acme Inc"] * expected_count,
+    }
+
+
+@pytest.mark.asyncio
+async def test_long_text_chunking_async(text_service):
+    long_text = "John Doe works at Acme Inc. Jane Smith is from New York City."
+    result = await text_service.annotate_text_async(long_text)
+    expected_count = len(text_service._chunk_text(long_text))
+    assert result == {
+        "PER": ["John Doe"] * expected_count,
+        "ORG": ["Acme Inc"] * expected_count,
+    }
+
+
+def test_empty_string(text_service):
+    result = text_service.annotate_text_sync("")
+    assert result == {}
+
+
+def test_short_string(text_service):
+    result = text_service.annotate_text_sync("Short")
+    assert result == {"PER": ["John Doe"], "ORG": ["Acme Inc"]}
+
+
+def test_special_characters(text_service):
+    result = text_service.annotate_text_sync("John@Doe.com works at Acme-Inc!!!")
+    expected_count = len(text_service._chunk_text("John@Doe.com works at Acme-Inc!!!"))
+    assert result == {
+        "PER": ["John Doe"] * expected_count,
+        "ORG": ["Acme Inc"] * expected_count,
+    }