graphrag-chunking 3.0.0__tar.gz

graphrag_chunking-3.0.0/.gitignore

@@ -0,0 +1,65 @@
+# Python Artifacts
+python/*/lib/
+dist/
+build/
+*.egg-info/
+
+# Test Output
+.coverage
+coverage/
+licenses.txt
+examples_notebooks/*/data
+tests/fixtures/cache
+tests/fixtures/*/cache
+tests/fixtures/*/output
+output/lancedb
+
+
+# Random
+.DS_Store
+*.log*
+.venv
+venv/
+.conda
+.tmp
+packages/graphrag-llm/notebooks/metrics
+packages/graphrag-llm/notebooks/cache
+
+.env
+build.zip
+
+.turbo
+
+__pycache__
+
+.pipeline
+
+# Azurite
+temp_azurite/
+__azurite*.json
+__blobstorage*.json
+__blobstorage__/
+
+# Getting started example
+ragtest/
+.ragtest/
+.pipelines
+.pipeline
+
+
+# mkdocs
+site/
+
+# Docs migration
+docsite/
+.yarn/
+.pnp*
+
+# PyCharm
+.idea/
+
+# Jupyter notebook
+.ipynb_checkpoints/
+
+# Root build assets
+packages/*/LICENSE

graphrag_chunking-3.0.0/LICENSE

@@ -0,0 +1,21 @@
+    MIT License
+
+    Copyright (c) Microsoft Corporation.
+
+    Permission is hereby granted, free of charge, to any person obtaining a copy
+    of this software and associated documentation files (the "Software"), to deal
+    in the Software without restriction, including without limitation the rights
+    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+    copies of the Software, and to permit persons to whom the Software is
+    furnished to do so, subject to the following conditions:
+
+    The above copyright notice and this permission notice shall be included in all
+    copies or substantial portions of the Software.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+    SOFTWARE

graphrag_chunking-3.0.0/PKG-INFO

@@ -0,0 +1,59 @@
+Metadata-Version: 2.4
+Name: graphrag-chunking
+Version: 3.0.0
+Summary: Chunking utilities for GraphRAG
+Project-URL: Source, https://github.com/microsoft/graphrag
+Author: Mónica Carvajal
+Author-email: Alonso Guevara Fernández <alonsog@microsoft.com>, Andrés Morales Esquivel <andresmor@microsoft.com>, Chris Trevino <chtrevin@microsoft.com>, David Tittsworth <datittsw@microsoft.com>, Dayenne de Souza <ddesouza@microsoft.com>, Derek Worthen <deworthe@microsoft.com>, Gaudy Blanco Meneses <gaudyb@microsoft.com>, Ha Trinh <trinhha@microsoft.com>, Jonathan Larson <jolarso@microsoft.com>, Josh Bradley <joshbradley@microsoft.com>, Kate Lytvynets <kalytv@microsoft.com>, Kenny Zhang <zhangken@microsoft.com>, Nathan Evans <naevans@microsoft.com>, Rodrigo Racanicci <rracanicci@microsoft.com>, Sarah Smith <smithsarah@microsoft.com>
+License: MIT
+License-File: LICENSE
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: <3.14,>=3.11
+Requires-Dist: graphrag-common==3.0.0
+Requires-Dist: pydantic~=2.10
+Description-Content-Type: text/markdown
+
+# GraphRAG Chunking
+
+This package contains a collection of text chunkers, a core config model, and a factory for acquiring instances.
+
+## Examples
+
+### Basic sentence chunking with nltk
+
+The SentenceChunker class splits text into individual sentences by identifying sentence boundaries. It takes input text and returns a list where each element is a separate sentence, making it easy to process text at the sentence level.
+
+```python
+chunker = SentenceChunker()
+chunks = chunker.chunk("This is a test. Another sentence.")
+print(chunks) # ["This is a test.", "Another sentence."]
+```
+
+### Token chunking
+
+The TokenChunker splits text into fixed-size chunks based on token count rather than sentence boundaries. It uses a tokenizer to encode text into tokens, then creates chunks of a specified size with configurable overlap between chunks.
+
+```python
+tokenizer = tiktoken.get_encoding("o200k_base")
+chunker = TokenChunker(size=3, overlap=0, encode=tokenizer.encode, decode=tokenizer.decode)
+chunks = chunker.chunk("This is a random test fragment of some text")
+print(chunks) # ["This is a", " random test fragment", " of some text"]
+```
+
+### Using the factory via helper util
+
+The create_chunker factory function provides a configuration-driven approach to instantiate chunkers by accepting a ChunkingConfig object that specifies the chunking strategy and parameters. This allows for more flexible and maintainable code by separating chunker configuration from direct instantiation.
+
+```python
+tokenizer = tiktoken.get_encoding("o200k_base")
+config = ChunkingConfig(
+    strategy="tokens",
+    size=3,
+    overlap=0
+)
+chunker = create_chunker(config, tokenizer.encode, tokenizer.decode)
+...
+```

graphrag_chunking-3.0.0/README.md

@@ -0,0 +1,41 @@
+# GraphRAG Chunking
+
+This package contains a collection of text chunkers, a core config model, and a factory for acquiring instances.
+
+## Examples
+
+### Basic sentence chunking with nltk
+
+The SentenceChunker class splits text into individual sentences by identifying sentence boundaries. It takes input text and returns a list where each element is a separate sentence, making it easy to process text at the sentence level.
+
+```python
+chunker = SentenceChunker()
+chunks = chunker.chunk("This is a test. Another sentence.")
+print(chunks) # ["This is a test.", "Another sentence."]
+```
+
+### Token chunking
+
+The TokenChunker splits text into fixed-size chunks based on token count rather than sentence boundaries. It uses a tokenizer to encode text into tokens, then creates chunks of a specified size with configurable overlap between chunks.
+
+```python
+tokenizer = tiktoken.get_encoding("o200k_base")
+chunker = TokenChunker(size=3, overlap=0, encode=tokenizer.encode, decode=tokenizer.decode)
+chunks = chunker.chunk("This is a random test fragment of some text")
+print(chunks) # ["This is a", " random test fragment", " of some text"]
+```
+
+### Using the factory via helper util
+
+The create_chunker factory function provides a configuration-driven approach to instantiate chunkers by accepting a ChunkingConfig object that specifies the chunking strategy and parameters. This allows for more flexible and maintainable code by separating chunker configuration from direct instantiation.
+
+```python
+tokenizer = tiktoken.get_encoding("o200k_base")
+config = ChunkingConfig(
+    strategy="tokens",
+    size=3,
+    overlap=0
+)
+chunker = create_chunker(config, tokenizer.encode, tokenizer.decode)
+...
+```

graphrag_chunking-3.0.0/graphrag_chunking/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""System-level chunking package."""

graphrag_chunking-3.0.0/graphrag_chunking/bootstrap_nltk.py

@@ -0,0 +1,31 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""Bootstrap definition."""
+
+import warnings
+
+# Ignore warnings from numba
+warnings.filterwarnings("ignore", message=".*The 'nopython' keyword.*")
+warnings.filterwarnings("ignore", message=".*Use no seed for parallelism.*")
+
+initialized_nltk = False
+
+
+def bootstrap():
+    """Bootstrap definition."""
+    global initialized_nltk
+    if not initialized_nltk:
+        import nltk
+        from nltk.corpus import wordnet as wn
+
+        nltk.download("punkt")
+        nltk.download("punkt_tab")
+        nltk.download("averaged_perceptron_tagger")
+        nltk.download("averaged_perceptron_tagger_eng")
+        nltk.download("maxent_ne_chunker")
+        nltk.download("maxent_ne_chunker_tab")
+        nltk.download("words")
+        nltk.download("wordnet")
+        wn.ensure_loaded()
+        initialized_nltk = True

graphrag_chunking-3.0.0/graphrag_chunking/chunk_strategy_type.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""Chunk strategy type enumeration."""
+
+from enum import StrEnum
+
+
+class ChunkerType(StrEnum):
+    """ChunkerType class definition."""
+
+    Tokens = "tokens"
+    Sentence = "sentence"

graphrag_chunking-3.0.0/graphrag_chunking/chunker.py

@@ -0,0 +1,24 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing the 'Chunker' class."""
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import Any
+
+from graphrag_chunking.text_chunk import TextChunk
+
+
+class Chunker(ABC):
+    """Abstract base class for document chunkers."""
+
+    @abstractmethod
+    def __init__(self, **kwargs: Any) -> None:
+        """Create a chunker instance."""
+
+    @abstractmethod
+    def chunk(
+        self, text: str, transform: Callable[[str], str] | None = None
+    ) -> list[TextChunk]:
+        """Chunk method definition."""

graphrag_chunking-3.0.0/graphrag_chunking/chunker_factory.py

@@ -0,0 +1,77 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing 'ChunkerFactory', 'register_chunker', and 'create_chunker'."""
+
+from collections.abc import Callable
+
+from graphrag_common.factory.factory import Factory, ServiceScope
+
+from graphrag_chunking.chunk_strategy_type import ChunkerType
+from graphrag_chunking.chunker import Chunker
+from graphrag_chunking.chunking_config import ChunkingConfig
+
+
+class ChunkerFactory(Factory[Chunker]):
+    """Factory for creating Chunker instances."""
+
+
+chunker_factory = ChunkerFactory()
+
+
+def register_chunker(
+    chunker_type: str,
+    chunker_initializer: Callable[..., Chunker],
+    scope: ServiceScope = "transient",
+) -> None:
+    """Register a custom chunker implementation.
+
+    Args
+    ----
+        - chunker_type: str
+            The chunker id to register.
+        - chunker_initializer: Callable[..., Chunker]
+            The chunker initializer to register.
+    """
+    chunker_factory.register(chunker_type, chunker_initializer, scope)
+
+
+def create_chunker(
+    config: ChunkingConfig,
+    encode: Callable[[str], list[int]] | None = None,
+    decode: Callable[[list[int]], str] | None = None,
+) -> Chunker:
+    """Create a chunker implementation based on the given configuration.
+
+    Args
+    ----
+        - config: ChunkingConfig
+            The chunker configuration to use.
+
+    Returns
+    -------
+        Chunker
+            The created chunker implementation.
+    """
+    config_model = config.model_dump()
+    if encode is not None:
+        config_model["encode"] = encode
+    if decode is not None:
+        config_model["decode"] = decode
+    chunker_strategy = config.type
+
+    if chunker_strategy not in chunker_factory:
+        match chunker_strategy:
+            case ChunkerType.Tokens:
+                from graphrag_chunking.token_chunker import TokenChunker
+
+                register_chunker(ChunkerType.Tokens, TokenChunker)
+            case ChunkerType.Sentence:
+                from graphrag_chunking.sentence_chunker import SentenceChunker
+
+                register_chunker(ChunkerType.Sentence, SentenceChunker)
+            case _:
+                msg = f"ChunkingConfig.strategy '{chunker_strategy}' is not registered in the ChunkerFactory. Registered types: {', '.join(chunker_factory.keys())}."
+                raise ValueError(msg)
+
+    return chunker_factory.create(chunker_strategy, init_args=config_model)

graphrag_chunking-3.0.0/graphrag_chunking/chunking_config.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""Parameterization settings for the default configuration."""
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from graphrag_chunking.chunk_strategy_type import ChunkerType
+
+
+class ChunkingConfig(BaseModel):
+    """Configuration section for chunking."""
+
+    model_config = ConfigDict(extra="allow")
+    """Allow extra fields to support custom cache implementations."""
+
+    type: str = Field(
+        description="The chunking type to use.",
+        default=ChunkerType.Tokens,
+    )
+    encoding_model: str | None = Field(
+        description="The encoding model to use.",
+        default=None,
+    )
+    size: int = Field(
+        description="The chunk size to use.",
+        default=1200,
+    )
+    overlap: int = Field(
+        description="The chunk overlap to use.",
+        default=100,
+    )
+    prepend_metadata: list[str] | None = Field(
+        description="Metadata fields from the source document to prepend on each chunk.",
+        default=None,
+    )

graphrag_chunking-3.0.0/graphrag_chunking/create_chunk_results.py

@@ -0,0 +1,32 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing 'create_chunk_results' function."""
+
+from collections.abc import Callable
+
+from graphrag_chunking.text_chunk import TextChunk
+
+
+def create_chunk_results(
+    chunks: list[str],
+    transform: Callable[[str], str] | None = None,
+    encode: Callable[[str], list[int]] | None = None,
+) -> list[TextChunk]:
+    """Create chunk results from a list of text chunks. The index assignments are 0-based and assume chunks were not stripped relative to the source text."""
+    results = []
+    start_char = 0
+    for index, chunk in enumerate(chunks):
+        end_char = start_char + len(chunk) - 1  # 0-based indices
+        result = TextChunk(
+            original=chunk,
+            text=transform(chunk) if transform else chunk,
+            index=index,
+            start_char=start_char,
+            end_char=end_char,
+        )
+        if encode:
+            result.token_count = len(encode(result.text))
+        results.append(result)
+        start_char = end_char + 1
+    return results

graphrag_chunking-3.0.0/graphrag_chunking/sentence_chunker.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing 'SentenceChunker' class."""
+
+from collections.abc import Callable
+from typing import Any
+
+import nltk
+
+from graphrag_chunking.bootstrap_nltk import bootstrap
+from graphrag_chunking.chunker import Chunker
+from graphrag_chunking.create_chunk_results import create_chunk_results
+from graphrag_chunking.text_chunk import TextChunk
+
+
+class SentenceChunker(Chunker):
+    """A chunker that splits text into sentence-based chunks."""
+
+    def __init__(
+        self, encode: Callable[[str], list[int]] | None = None, **kwargs: Any
+    ) -> None:
+        """Create a sentence chunker instance."""
+        self._encode = encode
+        bootstrap()
+
+    def chunk(
+        self, text: str, transform: Callable[[str], str] | None = None
+    ) -> list[TextChunk]:
+        """Chunk the text into sentence-based chunks."""
+        sentences = nltk.sent_tokenize(text.strip())
+        results = create_chunk_results(
+            sentences, transform=transform, encode=self._encode
+        )
+        # nltk sentence tokenizer may trim whitespace, so we need to adjust start/end chars
+        for index, result in enumerate(results):
+            txt = result.text
+            start = result.start_char
+            actual_start = text.find(txt, start)
+            delta = actual_start - start
+            if delta > 0:
+                result.start_char += delta
+                result.end_char += delta
+                # bump the next to keep the start check from falling too far behind
+                if index < len(results) - 1:
+                    results[index + 1].start_char += delta
+                    results[index + 1].end_char += delta
+        return results

graphrag_chunking-3.0.0/graphrag_chunking/text_chunk.py

@@ -0,0 +1,29 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""The TextChunk dataclass."""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class TextChunk:
+    """Result of chunking a document."""
+
+    original: str
+    """Raw original text chunk before any transformation."""
+
+    text: str
+    """The final text content of this chunk."""
+
+    index: int
+    """Zero-based index of this chunk within the source document."""
+
+    start_char: int
+    """Character index where the raw chunk text begins in the source document."""
+
+    end_char: int
+    """Character index where the raw chunk text ends in the source document."""
+
+    token_count: int | None = None
+    """Number of tokens in the final chunk text, if computed."""

graphrag_chunking-3.0.0/graphrag_chunking/token_chunker.py

@@ -0,0 +1,69 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing 'TokenChunker' class."""
+
+from collections.abc import Callable
+from typing import Any
+
+from graphrag_chunking.chunker import Chunker
+from graphrag_chunking.create_chunk_results import create_chunk_results
+from graphrag_chunking.text_chunk import TextChunk
+
+
+class TokenChunker(Chunker):
+    """A chunker that splits text into token-based chunks."""
+
+    def __init__(
+        self,
+        size: int,
+        overlap: int,
+        encode: Callable[[str], list[int]],
+        decode: Callable[[list[int]], str],
+        **kwargs: Any,
+    ) -> None:
+        """Create a token chunker instance."""
+        self._size = size
+        self._overlap = overlap
+        self._encode = encode
+        self._decode = decode
+
+    def chunk(
+        self, text: str, transform: Callable[[str], str] | None = None
+    ) -> list[TextChunk]:
+        """Chunk the text into token-based chunks."""
+        chunks = split_text_on_tokens(
+            text,
+            chunk_size=self._size,
+            chunk_overlap=self._overlap,
+            encode=self._encode,
+            decode=self._decode,
+        )
+        return create_chunk_results(chunks, transform=transform, encode=self._encode)
+
+
+def split_text_on_tokens(
+    text: str,
+    chunk_size: int,
+    chunk_overlap: int,
+    encode: Callable[[str], list[int]],
+    decode: Callable[[list[int]], str],
+) -> list[str]:
+    """Split a single text and return chunks using the tokenizer."""
+    result = []
+    input_tokens = encode(text)
+
+    start_idx = 0
+    cur_idx = min(start_idx + chunk_size, len(input_tokens))
+    chunk_tokens = input_tokens[start_idx:cur_idx]
+
+    while start_idx < len(input_tokens):
+        chunk_text = decode(list(chunk_tokens))
+        result.append(chunk_text)  # Append chunked text as string
+        if cur_idx == len(input_tokens):
+            break
+        start_idx += chunk_size - chunk_overlap
+        cur_idx = min(start_idx + chunk_size, len(input_tokens))
+        chunk_tokens = input_tokens[start_idx:cur_idx]
+
+    return result

graphrag_chunking-3.0.0/graphrag_chunking/transformers.py

@@ -0,0 +1,25 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A collection of useful built-in transformers you can use for chunking."""
+
+from collections.abc import Callable
+from typing import Any
+
+
+def add_metadata(
+    metadata: dict[str, Any],
+    delimiter: str = ": ",
+    line_delimiter: str = "\n",
+    append: bool = False,
+) -> Callable[[str], str]:
+    """Add metadata to the given text, prepending by default. This utility writes the dict as rows of key/value pairs."""
+
+    def transformer(text: str) -> str:
+        metadata_str = (
+            line_delimiter.join(f"{k}{delimiter}{v}" for k, v in metadata.items())
+            + line_delimiter
+        )
+        return text + metadata_str if append else metadata_str + text
+
+    return transformer

graphrag_chunking-3.0.0/pyproject.toml

@@ -0,0 +1,43 @@
+[project]
+name = "graphrag-chunking"
+version = "3.0.0"
+description = "Chunking utilities for GraphRAG"
+authors = [
+    {name = "Alonso Guevara Fernández", email = "alonsog@microsoft.com"},
+    {name = "Andrés Morales Esquivel", email = "andresmor@microsoft.com"},
+    {name = "Chris Trevino", email = "chtrevin@microsoft.com"},
+    {name = "David Tittsworth", email = "datittsw@microsoft.com"},
+    {name = "Dayenne de Souza", email = "ddesouza@microsoft.com"},
+    {name = "Derek Worthen", email = "deworthe@microsoft.com"},
+    {name = "Gaudy Blanco Meneses", email = "gaudyb@microsoft.com"},
+    {name = "Ha Trinh", email = "trinhha@microsoft.com"},
+    {name = "Jonathan Larson", email = "jolarso@microsoft.com"},
+    {name = "Josh Bradley", email = "joshbradley@microsoft.com"},
+    {name = "Kate Lytvynets", email = "kalytv@microsoft.com"},
+    {name = "Kenny Zhang", email = "zhangken@microsoft.com"},
+    {name = "Mónica Carvajal"},
+    {name = "Nathan Evans", email = "naevans@microsoft.com"},
+    {name = "Rodrigo Racanicci", email = "rracanicci@microsoft.com"},
+    {name = "Sarah Smith", email = "smithsarah@microsoft.com"},
+]
+license = {text = "MIT"}
+readme = "README.md"
+requires-python = ">=3.11,<3.14"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+]
+dependencies = [
+    "graphrag-common==3.0.0",
+    "pydantic~=2.10",
+]
+
+[project.urls]
+Source = "https://github.com/microsoft/graphrag"
+
+[build-system]
+requires = ["hatchling>=1.27.0,<2.0.0"]
+build-backend = "hatchling.build"
+