feat: Splitter tracks token and char length, embedding step logs the length statistics (#167)

Author: malteos

tests/datacontract/md_test.py

@@ -50,6 +50,9 @@ def test_manual_step_md_parsing(tmp_path, md, url, bread):
     else:
         assert s.md == "Text"
 
+    # check metadata field
+    assert s.metadata is None
+
 
 class MDCChild(MarkdownDataContract):
     pass
@@ -200,6 +203,40 @@ def test_topics_deprecation_warning(tmp_path):
         assert s.md.startswith("# Some title")
 
 
+def test_metadata_field_metadata(tmp_path):
+    md = """---
+keywords: "k1"
+url: foo/bar
+metadata:
+ foo: bar
+ bar: 123
+---
+# Title
+
+Text.
+ """
+    f = tmp_path / "file.md"
+    f.write_text(md)
+    s = MarkdownDataContract.from_file(f)
+
+    assert "# Title" in s.md
+    assert s.metadata is not None
+    assert s.metadata["foo"] == "bar"
+    assert s.metadata["bar"] == 123
+    assert s.url == "foo/bar"
+
+    assert s.__hash__() == 21317556317919954558699657768736304700342060298586059611903002870732316103488, "Invalid hash"
+
+    # save and load again
+    f2 = tmp_path / "file2.json"
+
+    MarkdownDataContract.save_to_path(f2, s)
+
+    s2 = MarkdownDataContract.load_from_path(f2, MarkdownDataContract)
+
+    assert s.__hash__() == s2.__hash__(), "Invalid hash after write/load file"
+
+
 def test_utf8_encoding(tmp_path):
     """Test that UTF-8 encoded files are read correctly, especially on Windows."""
     f = tmp_path / "file.md"

tests/steps/embedding/e2e_test.py

@@ -3,25 +3,27 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # Standard library imports
+import logging
 import shutil
 from pathlib import Path
 
 import numpy as np
 import pytest
 
-from wurzel.utils import HAS_LANGCHAIN_CORE, HAS_REQUESTS
+from wurzel.utils import HAS_LANGCHAIN_CORE, HAS_REQUESTS, HAS_SPACY, HAS_TIKTOKEN
 
-if not HAS_LANGCHAIN_CORE or not HAS_REQUESTS:
-    pytest.skip("Embedding dependencies (langchain-core, requests) are not available", allow_module_level=True)
+if not HAS_LANGCHAIN_CORE or not HAS_REQUESTS or not HAS_SPACY or not HAS_TIKTOKEN:
+    pytest.skip("Embedding dependencies (langchain-core, requests, spacy, tiktoken) are not available", allow_module_level=True)
 
 from wurzel.exceptions import StepFailed
 from wurzel.step_executor import BaseStepExecutor
-
-# Local application/library specific imports
 from wurzel.steps import EmbeddingStep
 from wurzel.steps.embedding.huggingface import HuggingFaceInferenceAPIEmbeddings
 from wurzel.steps.embedding.step_multivector import EmbeddingMultiVectorStep
 
+SPLITTER_TOKENIZER_MODEL = "gpt-3.5-turbo"
+SENTENCE_SPLITTER_MODEL = "de_core_news_sm"
+
 
 @pytest.fixture(scope="module")
 def mock_embedding():
@@ -87,12 +89,23 @@ def test_embedding_step(mock_embedding, default_embedding_data, env):
 
     """
     env.set("EMBEDDINGSTEP__API", "https://example-embedding.com/embed")
+    env.set("EMBEDDINGSTEP__TOKEN_COUNT_MIN", "64")
+    env.set("EMBEDDINGSTEP__TOKEN_COUNT_MAX", "256")
+    env.set("EMBEDDINGSTEP__TOKEN_COUNT_BUFFER", "32")
+    env.set("EMBEDDINGSTEP__TOKENIZER_MODEL", SPLITTER_TOKENIZER_MODEL)
+    env.set("EMBEDDINGSTEP__SENTENCE_SPLITTER_MODEL", SENTENCE_SPLITTER_MODEL)
+
     EmbeddingStep._select_embedding = mock_embedding
     input_folder, output_folder = default_embedding_data
-    BaseStepExecutor(dont_encapsulate=False).execute_step(EmbeddingStep, [input_folder], output_folder)
+    step_res = BaseStepExecutor(dont_encapsulate=False).execute_step(EmbeddingStep, [input_folder], output_folder)
     assert output_folder.is_dir()
     assert len(list(output_folder.glob("*"))) > 0
 
+    step_output, step_report = step_res[0]
+
+    assert len(step_output) == 11, "Step outputs have wrong count."
+    assert step_report.results == 11, "Step report has wrong count of outputs."
+
 
 def test_mutlivector_embedding_step(mock_embedding, tmp_path, env):
     """Tests the execution of the `EmbeddingMultiVectorStep` with a mock input file.
@@ -137,3 +150,60 @@ def _select_embedding(*args, **kwargs) -> HuggingFaceInferenceAPIEmbeddings:
         with BaseStepExecutor() as ex:
             ex(InheritedStep, [inp], out)
     assert sf.value.message.endswith(EXPECTED_EXCEPTION)
+
+
+def test_embedding_step_log_statistics(mock_embedding, default_embedding_data, env, caplog):
+    """Tests the logging of descriptive statistics in the `EmbeddingStep` with a mock input file."""
+    env.set("EMBEDDINGSTEP__API", "https://example-embedding.com/embed")
+    env.set("EMBEDDINGSTEP__NUM_THREADS", "1")  # Ensure deterministic behavior with single thread
+    env.set("EMBEDDINGSTEP__TOKEN_COUNT_MIN", "64")
+    env.set("EMBEDDINGSTEP__TOKEN_COUNT_MAX", "256")
+    env.set("EMBEDDINGSTEP__TOKEN_COUNT_BUFFER", "32")
+    env.set("EMBEDDINGSTEP__TOKENIZER_MODEL", SPLITTER_TOKENIZER_MODEL)
+    env.set("EMBEDDINGSTEP__SENTENCE_SPLITTER_MODEL", SENTENCE_SPLITTER_MODEL)
+
+    EmbeddingStep._select_embedding = mock_embedding
+    input_folder, output_folder = default_embedding_data
+
+    with caplog.at_level(logging.INFO):
+        BaseStepExecutor(dont_encapsulate=False).execute_step(EmbeddingStep, [input_folder], output_folder)
+
+    # check if output log exists
+    assert "Distribution of char length" in caplog.text, "Missing log output for char length"
+    assert "Distribution of token length" in caplog.text, "Missing log output for token length"
+    assert "Distribution of chunks count" in caplog.text, "Missing log output for chunks count"
+
+    # check extras
+    char_length_record = None
+    token_length_record = None
+    chunks_count_record = None
+
+    for record in caplog.records:
+        if "Distribution of char length" in record.message:
+            char_length_record = record
+
+        if "Distribution of token length" in record.message:
+            token_length_record = record
+
+        if "Distribution of chunks count" in record.message:
+            chunks_count_record = record
+
+    expected_char_length_count = 11
+
+    # Check values if a small tolerance
+    expected_char_length_mean = pytest.approx(609.18, abs=0.1)
+    expected_token_length_mean = pytest.approx(257.18, abs=0.1)
+    expected_chunks_count_mean = pytest.approx(3.18, abs=0.2)
+
+    assert char_length_record.count == expected_char_length_count, (
+        f"Invalid char length count: expected {expected_char_length_count}, got {char_length_record.count}"
+    )
+    assert char_length_record.mean == expected_char_length_mean, (
+        f"Invalid char length mean: expected {expected_char_length_mean}, got {char_length_record.mean}"
+    )
+    assert token_length_record.mean == expected_token_length_mean, (
+        f"Invalid token length mean: expected {expected_token_length_mean}, got {token_length_record.mean}"
+    )
+    assert chunks_count_record.mean == expected_chunks_count_mean, (
+        f"Invalid chunks count mean: expected {expected_chunks_count_mean}, got {chunks_count_record.mean}"
+    )

tests/steps/simple_splitter/__init__.py

@@ -0,0 +1,42 @@
+# SPDX-FileCopyrightText: 2025 Deutsche Telekom AG (opensource@telekom.de)
+#
+# SPDX-License-Identifier: Apache-2.0
+import shutil
+from pathlib import Path
+
+import pytest
+
+from wurzel.utils import HAS_SPACY, HAS_TIKTOKEN
+
+if not HAS_SPACY or not HAS_TIKTOKEN:
+    pytest.skip("Simple splitter dependencies (spacy, tiktoken) are not available", allow_module_level=True)
+
+from wurzel.step_executor import BaseStepExecutor
+from wurzel.steps.splitter import SimpleSplitterStep
+
+
+@pytest.fixture
+def default_markdown_data(tmp_path):
+    mock_file = Path("tests/data/markdown.json")
+    input_folder = tmp_path / "input"
+    input_folder.mkdir()
+    shutil.copy(mock_file, input_folder)
+    output_folder = tmp_path / "out"
+    return (input_folder, output_folder)
+
+
+def test_simple_splitter_step(default_markdown_data, env):
+    """Tests the execution of the `SimpleSplitterStep` with a mock input file."""
+    env.set("SIMPLESPLITTERSTEP__TOKEN_COUNT_MIN", "64")
+    env.set("SIMPLESPLITTERSTEP__TOKEN_COUNT_MAX", "256")
+    env.set("SIMPLESPLITTERSTEP__TOKEN_COUNT_BUFFER", "32")
+
+    input_folder, output_folder = default_markdown_data
+    step_res = BaseStepExecutor(dont_encapsulate=False).execute_step(SimpleSplitterStep, [input_folder], output_folder)
+    assert output_folder.is_dir()
+    assert len(list(output_folder.glob("*"))) > 0
+
+    step_output, step_report = step_res[0]
+
+    assert len(step_output) == 11, "Step outputs have wrong count."
+    assert step_report.results == 11, "Step report has wrong count of outputs."

tests/steps/simple_splitter/e2e_simple_splitter_test.py

@@ -19,7 +19,7 @@
 
 
 class MarkdownDataContract(PydanticModel):
-    """A data contract of the input of the EmbeddingStep representing a document in Markdown format.
+    """A data contract of the input/output of the various pipeline steps representing a document in Markdown format.
 
     The document consists have the Markdown body (document content) and additional metadata (keywords, url).
     The metadata is optional.
@@ -47,11 +47,25 @@ class MarkdownDataContract(PydanticModel):
     Another text.
     ```
 
+    Example 3 (with extra metadata fields)
+    ```md
+    ---
+    keywords: "bread,butter"
+    url: "some/file/path.md"
+    metadata:
+        token_len: 123
+        char_len: 550
+    ---
+    # Some title
+
+    A short text.
+    ```
     """
 
     md: str
     keywords: str
     url: str  # Url of pydantic is buggy in serialization
+    metadata: dict[str, Any] | None = None
 
     @classmethod
     @pydantic.validate_call
@@ -61,6 +75,7 @@ def from_dict_w_function(cls, doc: dict[str, Any], func: Callable[[str], str]):
             md=func(doc["text"]),
             url=doc["metadata"]["url"],
             keywords=doc["metadata"]["keywords"],
+            metadata=doc["metadata"].get("metadata", None),
         )
 
     @classmethod
@@ -115,4 +130,5 @@ def from_file(cls, path: Path, url_prefix: str = "") -> Self:
             # Extract metadata fields or use default value
             url=metadata.get("url", url_prefix + str(path.absolute())),
             keywords=metadata.get("keywords", path.name.split(".")[0]),
+            metadata=metadata.get("metadata", None),
         )

wurzel/datacontract/common.py

@@ -116,11 +116,12 @@ def load_from_path(cls, path: Path, model_type: type[Union[Self, list[Self]]]) -
         raise NotImplementedError(f"Can not load {model_type}")
 
     def __hash__(self) -> int:
+        """Compute a hash based on all not-none field values."""
         # pylint: disable-next=not-an-iterable
         return int(
             hashlib.sha256(
                 bytes(
-                    "".join([getattr(self, name) for name in sorted(type(self).model_fields)]),
+                    "".join([str(getattr(self, name) or "") for name in sorted(type(self).model_fields)]),
                     encoding="utf-8",
                 ),
                 usedforsecurity=False,

wurzel/datacontract/datacontract.py

@@ -7,10 +7,12 @@
 # Standard library imports
 import os
 import re
+from collections import defaultdict
 from io import StringIO
 from logging import getLogger
 from typing import Optional, TypedDict
 
+import numpy as np
 from markdown import Markdown
 from pandera.typing import DataFrame
 from tqdm.auto import tqdm
@@ -87,9 +89,18 @@ def run(self, inpt: list[MarkdownDataContract]) -> DataFrame[EmbeddingResult]:
         splitted_md_rows = self._split_markdown(inpt)
         rows = []
         failed = 0
+        stats = defaultdict(list)
+
         for row in tqdm(splitted_md_rows, desc="Calculate Embeddings"):
             try:
                 rows.append(self._get_embedding(row))
+
+                # collect statistics
+                if row.metadata is not None:
+                    stats["char length"].append(row.metadata.get("char_len", 0))
+                    stats["token length"].append(row.metadata.get("token_len", 0))
+                    stats["chunks count"].append(row.metadata.get("chunks_count", 0))
+
             except EmbeddingAPIException as err:
                 log.warning(
                     f"Skipped because EmbeddingAPIException: {err.message}",
@@ -100,8 +111,47 @@ def run(self, inpt: list[MarkdownDataContract]) -> DataFrame[EmbeddingResult]:
             log.warning(f"{failed}/{len(splitted_md_rows)} got skipped")
         if failed == len(splitted_md_rows):
             raise StepFailed(f"all {len(splitted_md_rows)} embeddings got skipped")
+
+        # log statistics
+        for k, v in stats.items():
+            self.log_statistics(series=np.array(v), name=k)
+
         return DataFrame[EmbeddingResult](DataFrame[EmbeddingResult](rows))
 
+    def log_statistics(self, series: np.ndarray, name: str):
+        """Log descriptive statistics for all documents.
+
+        Parameters
+        ----------
+        series : np.ndarray
+            Numerical values representing the documents.
+        name : str
+            The name of the document metric.
+        """
+        stats = {
+            "count": len(series),
+            "mean": None,
+            "std": None,
+        }
+
+        if len(series) > 0:
+            stats.update(
+                {
+                    "mean": np.mean(series),
+                    "median": np.median(series),
+                    "std": np.std(series),
+                    "var": np.var(series),
+                    "min": np.min(series),
+                    "percentile_5": np.percentile(series, 5),
+                    "percentile_25": np.percentile(series, 25),
+                    "percentile_75": np.percentile(series, 75),
+                    "percentile_95": np.percentile(series, 95),
+                    "max": np.max(series),
+                }
+            )
+
+        log.info(f"Distribution of {name}: count={stats['count']}; mean={stats['mean']}; std={stats['std']}", extra=stats)
+
     def get_embedding_input_from_document(self, doc: MarkdownDataContract) -> str:
         """Clean the document such that it can be used as input to the embedding model.
 

wurzel/steps/embedding/step.py

@@ -24,7 +24,7 @@ class SplitterSettings(Settings):
     """Anything Embedding-related."""
 
     BATCH_SIZE: int = Field(100, gt=0)
-    NUM_THREADS: int = Field(4, gt=1)
+    NUM_THREADS: int = Field(4, ge=1)
     TOKEN_COUNT_MIN: int = Field(64, gt=0)
     TOKEN_COUNT_MAX: int = Field(1024, gt=1)
     TOKEN_COUNT_BUFFER: int = Field(32, gt=0)
@@ -94,13 +94,13 @@ def _split_markdown(self, markdowns: list[MarkdownDataContract]) -> list[Markdow
         """Creates data rows from a batch of markdown texts by splitting them and counting tokens."""
         rows = []
         skipped = 0
-        for s in markdowns:
+        for md_data_contract in markdowns:
             try:
-                rows.extend(self.splitter.split_markdown_document(s))
+                rows.extend(self.splitter.split_markdown_document(md_data_contract))
             except MarkdownException as err:
                 log.warning(
                     "skipped dokument ",
-                    extra={"reason": err.__class__.__name__, "doc": s},
+                    extra={"reason": err.__class__.__name__, "doc": md_data_contract},
                 )
                 skipped += 1
         if skipped == len(markdowns):