chunknorris 0.0.1__py3-none-any.whl

chunknorris/__init__.py

@@ -0,0 +1,4 @@
+from .chunkers import *
+from .custom_chunkers import *
+from .exceptions import *
+from .types import *

chunknorris/__init__.pyi

@@ -0,0 +1,4 @@
+from .chunkers import *
+from .custom_chunkers import *
+from .exceptions import *
+from .types import *

chunknorris/chunkers/__init__.py

@@ -0,0 +1,2 @@
+from .html_chunknorris import HTMLChunkNorris
+from .markdown_chunknorris import MarkdownChunkNorris

chunknorris/chunkers/__init__.pyi

@@ -0,0 +1,2 @@
+from .html_chunknorris import HTMLChunkNorris as HTMLChunkNorris
+from .markdown_chunknorris import MarkdownChunkNorris as MarkdownChunkNorris

chunknorris/chunkers/html_chunknorris.py

@@ -0,0 +1,25 @@
+from markdownify import markdownify
+
+from .markdown_chunknorris import MarkdownChunkNorris
+
+
+class HTMLChunkNorris(MarkdownChunkNorris):
+
+    def __call__(self, html_text: str, **kwargs) -> str:
+        text = HTMLChunkNorris.apply_markdownify(html_text)
+
+        return super().__call__(text, **kwargs)
+
+    @staticmethod
+    def apply_markdownify(html_text) -> str:
+        """Applies markdownify to the html text
+
+        Args:
+            html_text (str): the text of the html file
+
+        Returns:
+            str: the markdownified string
+        """
+        md_text = markdownify(html_text, strip=["figure", "img"], bullets="-*+")
+
+        return md_text

chunknorris/chunkers/html_chunknorris.pyi

@@ -0,0 +1,6 @@
+from .markdown_chunknorris import MarkdownChunkNorris as MarkdownChunkNorris
+
+class HTMLChunkNorris(MarkdownChunkNorris):
+    def __call__(self, html_text: str, **kwargs) -> str: ...
+    @staticmethod
+    def apply_markdownify(html_text) -> str: ...

chunknorris/chunkers/markdown_chunknorris.py

@@ -0,0 +1,543 @@
+import copy
+import json
+import re
+import tiktoken
+from unicodedata import normalize
+
+from ..exceptions.exceptions import ChunkSizeExceeded
+from ..types.types import Chunk, Chunks, TocTree
+
+
+class MarkdownChunkNorris:
+    def __init__(self):
+        self.tokenizer = tiktoken.get_encoding("cl100k_base")
+
+    def __call__(self, md_text: str, **kwargs) -> Chunks:
+        """Gets chunks from markdown string
+
+        Args:
+            md_text (str): the markdown string
+
+        Returns:
+            Chunks: the list of chunk's texts
+        """
+        toc_tree = self.get_toc_tree(md_text, **kwargs)
+        chunks = self.get_chunks(toc_tree, **kwargs)
+
+        return chunks
+
+    @staticmethod
+    def _check_string_argument_is_valid(argname: str, argvalue: str):
+        """Checks that an argument has a valid value
+
+        Args:
+            argname (str): the name of the argument
+            argvalue (str): the value of the argument
+        """
+        allowed_values_dict = {
+            "header_style": ["setext", "atx"],
+            "max_title_level_to_use": ["h1", "h2", "h3", "h4", "h5"],
+            "link_placement": [
+                "remove",
+                "end_of_chunk",
+                "in_sentence",
+                "leave_as_markdown",
+            ],
+            "chunk_tokens_exceeded_handling": ["raise_error", "split"],
+        }
+        allowed_values = allowed_values_dict[argname]
+        assert argvalue in allowed_values, ValueError(
+            f"Argument '{argname}' should be one of {allowed_values}. Got '{argvalue}'"
+        )
+
+    def _get_header_regex_patterns(self, header_style: str = "setext") -> dict:
+        """Get the header regex patterns depending on the header_style
+
+        Args:
+            header_style (str, optional): the markdown header style. Must be 'atx' or 'setext'.
+                Defaults to "setext".
+
+        Returns:
+            (dict) : a mapping between header name and regex patterns
+        """
+        self._check_string_argument_is_valid("header_style", header_style)
+        patterns = {
+            "h1": re.compile(r"(.+?)\n={3,}", re.MULTILINE),
+            "h2": re.compile(r"(.+?)\n-{3,}", re.MULTILINE),
+            "h3": re.compile(r"^(?:- )?(#{3} .+)", re.MULTILINE),
+            "h4": re.compile(r"^(?:- )?(#{4} .+)", re.MULTILINE),
+            "h5": re.compile(r"^(?:- )?(#{5} .+)", re.MULTILINE),
+        }
+        if header_style == "atx":
+            patterns["h1"] = re.compile(r"^(?:- )?(#{1} .+)", re.MULTILINE)
+            patterns["h2"] = re.compile(r"^(?:- )?(#{2} .+)", re.MULTILINE)
+
+        return patterns
+
+    def _convert_setext_to_atx(self, markdown_string: str) -> str:
+        """Converts headers from setext style to markdown style
+
+        Args:
+            markdown_string (str): the markdown string
+
+        Return:
+            str: the string with formatted headers
+        """
+        regex_patterns = self._get_header_regex_patterns("setext")
+        for match in re.finditer(regex_patterns["h1"], markdown_string):
+            markdown_string = markdown_string.replace(match[0], f"# {match[1]}")
+        for match in re.finditer(regex_patterns["h2"], markdown_string):
+            markdown_string = markdown_string.replace(match[0], f"## {match[1]}")
+
+        return markdown_string
+
+    def get_toc_tree(
+        self,
+        markdown_string: str,
+        max_title_level_to_use: str = "h4",
+        header_style="setext",
+        **kwargs,
+    ) -> TocTree:
+        """Builds the table of content tree based on header
+
+        Args:
+            markdown_string (str): the string, as markdown
+            max_title_level_to_use (str, optional): the maximum title level to use. Headers with lower
+                level than this won't be considered as headers. Defaults to "h4".
+            header_style (str, optional): the type of header format. Defaults to "setext".
+
+        Returns:
+            TocTree: the table of content
+        """
+        MarkdownChunkNorris._check_string_argument_is_valid(
+            "header_style", header_style
+        )
+        MarkdownChunkNorris._check_string_argument_is_valid(
+            "max_title_level_to_use", max_title_level_to_use
+        )
+        if header_style == "setext":
+            markdown_string = self._convert_setext_to_atx(markdown_string)
+        title_types_to_consider = [
+            f"h{i}" for i in range(1, int(max_title_level_to_use[1]) + 1)
+        ]
+        regex_patterns = self._get_header_regex_patterns("atx")
+        lines = markdown_string.splitlines()
+        tree = {
+            "title": "",
+            "children": [],
+            "content": "",
+            "id": -1,
+            "line_index": -1,
+            "level": -1,
+            "parent": {},
+        }
+        current_level = -1
+        current_node = tree
+        id_cntr = 0
+        for line_idx, line in enumerate(lines):
+            for level, title_type in enumerate(title_types_to_consider):
+                match = re.match(regex_patterns[title_type], line)
+                if match:
+                    title = match.group(1)
+                    while level <= current_level:
+                        current_node = current_node["parent"]
+                        current_level = current_node["level"]
+                    new_node = {
+                        "title": title,
+                        "children": [],
+                        "content": "",
+                        "id": id_cntr,
+                        "line_index": line_idx,
+                        "level": level,
+                        "parent": current_node,
+                    }
+                    current_node["children"].append(new_node)
+                    current_node = new_node
+                    current_level = level
+                    id_cntr += 1
+                    break
+            if not match:
+                current_node["content"] += line + "\n"
+
+        MarkdownChunkNorris._cleanup_tree_texts(tree)
+
+        return tree
+
+    @staticmethod
+    def save_toc_tree(toc_tree: TocTree, output_path: str = "toc_tree.json"):
+        """Save the toc tree as json
+
+        Args:
+            toc_tree (TocTree): the toc tree to save
+            output_path (str, optional): the output path. Defaults to "toc_tree.json".
+        """
+        dumpable_dict = copy.deepcopy(toc_tree)
+        MarkdownChunkNorris._remove_circular_refs(dumpable_dict)
+        with open(output_path, "w", encoding="UTF-8") as f:
+            json.dump(dumpable_dict, f)
+
+    @staticmethod
+    def _remove_circular_refs(toc_tree_element: TocTree):
+        """Recursively removes the circular ref in dict
+        (used to save as json).
+
+        Args:
+            toc_tree_element (TocTree): the element of toc tree
+        """
+        if "parent" in toc_tree_element:
+            del toc_tree_element["parent"]
+        if "children" in toc_tree_element:
+            for child in toc_tree_element["children"]:
+                MarkdownChunkNorris._remove_circular_refs(child)
+
+    @staticmethod
+    def _cleanup_tree_texts(tree: TocTree):
+        """Cleans up the texts at each 'content' key of the toc tree
+        Uses recursion to go through each child
+        Args:
+            tree (TocTree): the toc tree
+        """
+        text = tree["content"]
+        text = normalize("NFKD", text)
+        # remove special characters
+        special_chars = ["**", r"\*"]
+        for char in special_chars:
+            text = text.replace(char, "")
+        # remove white spaces and newlines
+        text = re.sub(r"\n\s*\n", "\n", text)
+        text = text.rstrip().lstrip()
+        tree["content"] = text
+        if tree["children"]:
+            for child in tree["children"]:
+                tree = MarkdownChunkNorris._cleanup_tree_texts(child)
+
+    @staticmethod
+    def get_title_by_id(toc_tree: TocTree, id: int) -> TocTree:
+        """Gets a toc tree using its id
+
+        Args:
+            toc_tree (TocTree): the toc tree
+            id (int): the id of the title we are looking for. Defaults to int.
+
+        Returns:
+            TocTree: the element of title we were looking for
+        """
+        if toc_tree.get("id") == id:
+            return toc_tree
+        for child in toc_tree.get("children", []):
+            target = MarkdownChunkNorris.get_title_by_id(child, id)
+            if target:
+                return target
+
+    def get_chunks(self, toc_tree: TocTree, **kwargs) -> Chunks:
+        """Wrapper that build the chunk's texts, check
+        that they fit in size, replace links formatting.
+
+        Args:
+            toc_tree (TocTree): the toc tree of the document
+
+        Returns:
+            Chunks: the chunks text, formatted
+        """
+        chunks = MarkdownChunkNorris.get_chunks_texts(toc_tree, **kwargs)
+        chunks = [MarkdownChunkNorris.change_links_format(c, **kwargs) for c in chunks]
+        chunks = MarkdownChunkNorris.remove_small_chunks(chunks, **kwargs)
+        chunks = self.remove_small_chunks(chunks, **kwargs)
+
+        return chunks
+
+    @staticmethod
+    def get_chunks_texts(
+        toc_tree_element: TocTree, already_ok_chunks: Chunks | None = None, **kwargs
+    ) -> Chunks:
+        """Uses the toc tree to build the chunks. Uses recursion.
+        Method :
+        - build the chunk (= titles from sections above + section content + content of subsections)
+        - if the chunk is too big:
+            - save the section as title + content (if section has content)
+            - subdivide section recursively using subsections
+        - else save it as is
+
+        Args:
+            toc_tree_element (TocTree): _description_
+            already_ok_chunks (Chunks, optional): the chunks already built.
+                Used for recursion. Defaults to None.
+
+        Returns:
+            Chunks: list of chunk's texts
+        """
+        if not already_ok_chunks:
+            already_ok_chunks = []
+
+        current_chunk = MarkdownChunkNorris.build_chunk_text(toc_tree_element)
+
+        if MarkdownChunkNorris._chunk_is_too_big(current_chunk, **kwargs):
+            if toc_tree_element["content"]:
+                parents = MarkdownChunkNorris.get_parents_headers(toc_tree_element)
+                current_chunk = "\n".join(
+                    parents + [toc_tree_element["title"], toc_tree_element["content"]]
+                )
+                already_ok_chunks.append(current_chunk)
+            for child in toc_tree_element["children"]:
+                already_ok_chunks = MarkdownChunkNorris.get_chunks_texts(
+                    child, already_ok_chunks, **kwargs
+                )
+        else:
+            already_ok_chunks.append(current_chunk)
+
+        return already_ok_chunks
+
+    @staticmethod
+    def build_chunk_text(toc_tree_element: TocTree) -> Chunk:
+        """Builds a chunk by apposing the text of headers
+        and recursively getting the content of children
+
+        Args:
+            toc_tree_element (TocTree): the toc tree element
+
+        Returns:
+            str: the chunk content. parent's headers + content
+        """
+        parents = MarkdownChunkNorris.get_parents_headers(toc_tree_element)
+        content = MarkdownChunkNorris._build_chunk_content(toc_tree_element)
+
+        return "\n".join(parents) + "\n" + content
+
+    @staticmethod
+    def _build_chunk_content(toc_tree_element: TocTree) -> str:
+        """Builds a chunk content (i.e without headers above)
+        from a toc tree. It uses the toc tree's content, and recursively
+        adds the header and content of its children
+
+        Args:
+            toc_tree_element (TocTree): the toc tree (or element of toc tree)
+
+        Returns:
+            str: the text of the chunk (without the headers of parents)
+        """
+        text = toc_tree_element.get("title") + "\n" + toc_tree_element.get("content")
+        if toc_tree_element.get("children"):
+            for child in toc_tree_element.get("children"):
+                text += "\n" + MarkdownChunkNorris._build_chunk_content(child)
+
+        return text
+
+    @staticmethod
+    def get_parents_headers(toc_tree_element: TocTree) -> list[str]:
+        """Gets a list of the titles that are parent
+        of the provided toc tree element. The list
+        is ordered in descending order in terms of header level.
+
+        Args:
+            toc_tree_element (TocTree): the toc tree element
+
+        Returns:
+            list[str]: the list of header's text
+        """
+        headers = []
+        while toc_tree_element.get("parent"):
+            toc_tree_element = toc_tree_element.get("parent")
+            headers.append(toc_tree_element.get("title"))
+        # remove empty string header, such as root header
+        headers = [h for h in headers if h]
+
+        return list(reversed(headers))
+
+    @staticmethod
+    def _chunk_is_too_big(
+        chunk: str, max_chunk_word_length: int = 200, **kwargs
+    ) -> bool:
+        """Returns True if the chunk is bigger than the value
+        specified as max_chunk_length in terms of word count
+
+        Args:
+            chunk (str): a chunk of text
+            max_chunk_length (int): the max size of the chunk in words. Default to 200.
+
+        Returns:
+            bool: True if the chunk is too big, else false
+        """
+        return len(chunk.split()) > max_chunk_word_length
+
+    @staticmethod
+    def _header_has_children(toc_tree_element: TocTree) -> bool:
+        """Mostly used for code comprehension.
+        Returns True if the title has children.
+
+        Args:
+            toc_tree_element (TocTree): the header to check for children
+
+        Returns:
+            bool: True if the header has children
+        """
+        return bool(toc_tree_element.get("children"))
+
+    @staticmethod
+    def change_links_format(text, link_placement: str = "in_sentence", **kwargs) -> str:
+        """Removes the markdown format of the links in the text.
+        The links are treated as specified by 'link_position':
+        - remove : links are removed
+        - in_sentence : the link is placed in the sentence, between parenthesis
+        - end_of_chunk : all links are added at the end of the text
+        - leave_as_markdown: leave links as markdown format
+
+        Args:
+            text (str): the text to find the links in
+            link_placement (str, optional): How the links should be handled. Defaults to end_of_chunk.
+
+        Returns:
+            str: the formated text
+        """
+        MarkdownChunkNorris._check_string_argument_is_valid(
+            "link_placement", link_placement
+        )
+        if link_placement == "leave_as_markdown":
+            return text
+
+        patterns = {
+            "image_as_link": r"\[!\[(.*?)\]\(.*?\)\]\((.*?)\)",
+            "image": r"!\[(.*?)\]\((https?:[^\s'()]+).*?\)",
+            "link": r"\[(.+?)\]\((https?:.+?)\)",
+        }
+        for pattern in patterns.values():
+            matches = re.finditer(pattern, text)
+            replacements = [(m[0], m[1], m[2]) for m in matches]
+            if replacements is not None:
+                text = MarkdownChunkNorris._handle_link_replacements(
+                    text, replacements, link_placement=link_placement
+                )
+
+        return text
+
+    @staticmethod
+    def _handle_link_replacements(
+        text: str,
+        replacements: list[tuple[str, str, str]],
+        link_placement: str = "in_sentence",
+    ) -> str:
+        """Handles the replacement of links in the text
+        according to specified format
+
+        Args:
+            text (str): the text to replace the links ind
+            replacements (list[tuple[str, str, str]]): the replacements. Is is a list of
+                [(entire regex match, link's text, link's url)]
+            link_placement (str, optional): _description_. Defaults to "in_sentence".
+
+        Returns:
+            str: the text with links modified
+        """
+        for i, m in enumerate(replacements):
+            match link_placement:
+                case "remove":
+                    text = text.replace(m[0], m[1])
+                case "end_of_chunk":
+                    if i == 0:
+                        text += "\nPour plus d'informations:\n"
+                    text = text.replace(m[0], m[1])
+                    text += f"- {m[1]}: {m[2]}"
+                case "in_sentence":
+                    text = text.replace(m[0], f"{m[1]} (lien : {m[2]})")
+
+        return text
+
+    @staticmethod
+    def remove_small_chunks(
+        chunks: Chunks, min_chunk_word_count: int = 15, **kwargs
+    ) -> Chunks:
+        """Removes chunks that have less words than the specified limit
+
+        Args:
+            chunks (Chunks): _description_
+            min_chunk_tokens (int, optional): the minimum size a chunk is allowed to be,
+                in words. Chunks with less words than this are dicarded. Defaults to 15.
+        Returns:
+            Chunks: the chunks with more words than the specified threshold
+        """
+        return [c for c in chunks if len(c.split()) >= min_chunk_word_count]
+
+    def split_big_chunks(
+        self,
+        chunks: Chunks,
+        max_chunk_tokens: int = 8191,
+        chunk_tokens_exceeded_handling: str = "raise_error",
+        **kwargs,
+    ) -> Chunks:
+        """Checks that the chunks do not exceed the token limit, considered as a hard limit
+        If chunk_tokens_exceeded_handling is:
+        - "raise_error" -> it will raise an error in case a chunk to big is found
+        for it to be investigated.
+        - "split" -> Chunks exceeding the max size will be split to fit max_chunk_tokens
+
+        Args:
+            chunks (Chunks): The chunks obtained from the get_chunks() method
+            max_chunk_tokens (int, optional): the maximum size a chunk is allowed to be,
+                in tokens. Defaults to 8191.
+            chunk_tokens_exceeded_handling (str, optional): whether or not error should be raised if a big
+                chunk is encountered, or split. Defaults to raise_error.
+        """
+        MarkdownChunkNorris._check_string_argument_is_valid(
+            "chunk_tokens_exceeded_handling", chunk_tokens_exceeded_handling
+        )
+        splitted_chunks = []
+        for chunk in chunks:
+            if self.get_token_count(chunk) < max_chunk_tokens:
+                splitted_chunks.append(chunk)
+            else:
+                match chunk_tokens_exceeded_handling:
+                    case "raise_error":
+                        raise ChunkSizeExceeded(
+                            (
+                                f"Found chunk bigger than the specified token limit {max_chunk_tokens}:",
+                                "You can disable this error and allow dummy splitting of this chunk by passing 'raise_error=False'",
+                                f"The chunk : {chunk}",
+                            )
+                        )
+                    case "split":
+                        splitted_chunk = self._dummy_split_big_chunk(
+                            chunk, max_chunk_tokens
+                        )
+                        splitted_chunks.extend(splitted_chunk)
+
+        return splitted_chunks
+
+    def get_token_count(self, text: Chunk) -> int:
+        """Get the token count of a chunk
+
+        Args:
+            text (Chunk): the text to get the token count from
+
+        Returns:
+            int: the token count
+        """
+        return len(self.tokenizer.encode(text))
+
+    def _dummy_split_big_chunk(
+        self,
+        chunk: Chunk,
+        max_chunk_tokens: int = 8191,
+    ) -> Chunks:
+        """Splits the chunk so that the subchunks fit un max_chunk_size
+
+        Args:
+            chunk (Chunk): _description_
+            max_chunk_tokens (int, optional): maximum chunk size. Defaults to 8191.
+
+        Returns:
+            Chunks: the subdivided chunks
+        """
+        token_count = self.get_token_count(chunk)
+        if token_count < max_chunk_tokens:
+            return [chunk]
+
+        split_count = (token_count // max_chunk_tokens) + 1
+        split_token_size = token_count // split_count
+        tokenized_text = self.tokenizer.encode(chunk)
+        splitted_text = [
+            self.tokenizer.decode(
+                tokenized_text[i * split_token_size : (i + 1) * split_token_size]
+            )
+            for i in range(split_count)
+        ]
+
+        return splitted_text

chunknorris/chunkers/markdown_chunknorris.pyi

@@ -0,0 +1,26 @@
+from ..exceptions.exceptions import ChunkSizeExceeded as ChunkSizeExceeded
+from ..types.types import Chunk as Chunk, Chunks as Chunks, TocTree as TocTree
+from _typeshed import Incomplete
+
+class MarkdownChunkNorris:
+    tokenizer: Incomplete
+    def __init__(self) -> None: ...
+    def __call__(self, md_text: str, **kwargs) -> Chunks: ...
+    def get_toc_tree(self, markdown_string: str, max_title_level_to_use: str = 'h4', header_style: str = 'setext', **kwargs) -> TocTree: ...
+    @staticmethod
+    def save_toc_tree(toc_tree: TocTree, output_path: str = 'toc_tree.json'): ...
+    @staticmethod
+    def get_title_by_id(toc_tree: TocTree, id: int) -> TocTree: ...
+    def get_chunks(self, toc_tree: TocTree, **kwargs) -> Chunks: ...
+    @staticmethod
+    def get_chunks_texts(toc_tree_element: TocTree, already_ok_chunks: Chunks | None = None, **kwargs) -> Chunks: ...
+    @staticmethod
+    def build_chunk_text(toc_tree_element: TocTree) -> Chunk: ...
+    @staticmethod
+    def get_parents_headers(toc_tree_element: TocTree) -> list[str]: ...
+    @staticmethod
+    def change_links_format(text, link_placement: str = 'in_sentence', **kwargs) -> str: ...
+    @staticmethod
+    def remove_small_chunks(chunks: Chunks, min_chunk_word_count: int = 15, **kwargs) -> Chunks: ...
+    def split_big_chunks(self, chunks: Chunks, max_chunk_tokens: int = 8191, chunk_tokens_exceeded_handling: str = 'raise_error', **kwargs) -> Chunks: ...
+    def get_token_count(self, text: Chunk) -> int: ...

chunknorris/custom_chunkers/__init__.py

@@ -0,0 +1 @@
+from .wikit_chunknorris import WikitChunkNorris

chunknorris/custom_chunkers/__init__.pyi

@@ -0,0 +1 @@
+from .wikit_chunknorris import WikitChunkNorris as WikitChunkNorris

chunknorris/custom_chunkers/wikit_chunknorris.py

@@ -0,0 +1,178 @@
+import glob
+import os
+import json
+from argparse import ArgumentParser
+from ..chunkers.html_chunknorris import HTMLChunkNorris
+from ..exceptions.exceptions import ChunkNorrisException
+from ..types.types import Chunks
+
+
+class WikitChunkNorris(HTMLChunkNorris):
+
+    @staticmethod
+    def read_file(filepath: str, return_full_content: bool = False) -> str:
+        """Reads a html or json file
+
+        Args:
+            filepath (str): path to a file. must end with .json or .html
+            return_full_content (bool): Only applies to JSON files. Indicates whether or not
+                the entire content of the file is returned or just the text content.
+        Returns:
+            str: the text, mardkownified
+        """
+        try:
+            with open(filepath, "r", encoding="UTF-8") as f:
+                if filepath.endswith(".json"):
+                    file = json.load(f)
+                    if not return_full_content:
+                        file = file["hasPart"][0]["text"]
+                elif filepath.endswith(".html"):
+                    file = f.read()
+                else:
+                    raise Exception("Can only open JSON or HTML files")
+        except Exception as e:
+            raise ChunkNorrisException(f"Can't open file '{filepath}': {e}") from e
+
+        return file
+
+    def chunk_file(self, filepath: str, output_filepath: str = None, **kwargs):
+        """Chunks a json file and save the chunked file
+
+        Args:
+            filepath (str): path to the file
+            output_filepath (str): path of the output file
+            make_dirs (bool): whether of not to make the path if folders don't exist
+        """
+        filepath = os.path.normpath(filepath)
+        html_text = WikitChunkNorris.read_file(filepath)
+        chunks = self(html_text, **kwargs)
+        file_content = WikitChunkNorris.format_output(filepath, chunks)
+
+        if not output_filepath:
+            output_dir = os.path.dirname(filepath) + "-chunked"
+            output_filepath = os.path.join(output_dir, os.path.basename(filepath))
+        else:
+            output_dir = os.path.dirname(output_filepath)
+        if not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+
+        with open(output_filepath, "w", encoding="UTF-8") as f:
+            json.dump(file_content, f, ensure_ascii=False)
+
+    def chunk_directory(self, input_dir: str, output_dir: str = None, **kwargs) -> None:
+        """Chunks the json files of entire directory, recursively
+
+        Args:
+            input_dir (str): the path to directory
+            output_dir (str): the directory where chunks will be saved
+        """
+        input_dir = os.path.normpath(input_dir)
+        if not output_dir:
+            output_dir = f"{input_dir}-chunked"
+
+        filepaths = [
+            subdir
+            for dir in os.walk(input_dir)
+            for subdir in glob.glob(os.path.join(dir[0], "*.json"))
+        ]
+        for fp in filepaths:
+            self.chunk_file(fp, fp.replace(input_dir, output_dir), **kwargs)
+
+    @staticmethod
+    def format_output(filepath: str, chunks: Chunks) -> dict:
+        """Formats the chunks according to the input json file
+        i.e places the chunks inside the key [hasPart]
+
+        Args:
+            filepath (str): path toward the json file
+            chunks (Chunks): the chunks
+
+        Returns:
+            dict: the formatted file content
+        """
+        file_content = WikitChunkNorris.read_file(filepath, return_full_content=True)
+        chunks = [
+            {
+                "@type": "DocumentChunk",
+                "text": c,
+                "position": i,
+            }
+            for i, c in enumerate(chunks)
+        ]
+
+        file_content["hasPart"] = chunks
+
+        return file_content
+
+
+def parse_arguments():
+    """Parse the given command-line arguments."""
+
+    parser = ArgumentParser(
+        description="Chunking a folder of json files containing a HTML test under the key ['hasPart'][0]['text']"
+    )
+
+    parser.add_argument(
+        "--input_dir",
+        type=str,
+        required=True,
+        help="Path to folder containing the files to chunk",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="Path to output folder where the chunked files will be stored",
+    )
+    parser.add_argument(
+        "--max_title_level_to_use",
+        type=str,
+        choices=["h1", "h2", "h3", "h4", "h5"],
+        default="h4",
+        help="The maximum level of titles to use for chunking",
+    )
+    parser.add_argument(
+        "--max_chunk_tokens",
+        type=int,
+        default=8191,
+        help="Hard limit of the token size of a chunk. Chunks bigger than that will be handled according to chunk_tokens_exceeded_handling",
+    )
+    parser.add_argument(
+        "--chunk_tokens_exceeded_handling",
+        type=str,
+        choices=["split", "raise_error"],
+        default="raise_error",
+        help="Whether a big chunk with not headers should be split or error should be raised",
+    )
+    parser.add_argument(
+        "--link_placement",
+        type=str,
+        choices=["remove", "end_of_chunk", "in_sentence", "end_of_sentence"],
+        default="end_of_chunk",
+        help="Where the links are placed in the chunks",
+    )
+    parser.add_argument(
+        "--max_chunk_word_count",
+        type=int,
+        default=250,
+        help="Soft limit of chunk size. Chunks bigger than this limit will be subdivided with lower level headers.",
+    )
+    parser.add_argument(
+        "--min_chunk_word_count",
+        type=int,
+        default=15,
+        help="Minium amount a word a chunk must have. If lower, the chunk is discarded.",
+    )
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_arguments()
+
+    wcn = WikitChunkNorris()
+    wcn.chunk_directory(**vars(args))
+
+
+if __name__ == "__main__":
+    main()

chunknorris/custom_chunkers/wikit_chunknorris.pyi

@@ -0,0 +1,14 @@
+from ..chunkers.html_chunknorris import HTMLChunkNorris as HTMLChunkNorris
+from ..exceptions.exceptions import ChunkNorrisException as ChunkNorrisException
+from ..types.types import Chunks as Chunks
+
+class WikitChunkNorris(HTMLChunkNorris):
+    @staticmethod
+    def read_file(filepath: str, return_full_content: bool = False) -> str: ...
+    def chunk_file(self, filepath: str, output_filepath: str = None, **kwargs): ...
+    def chunk_directory(self, input_dir: str, output_dir: str = None, **kwargs) -> None: ...
+    @staticmethod
+    def format_output(filepath: str, chunks: Chunks) -> dict: ...
+
+def parse_arguments(): ...
+def main() -> None: ...

chunknorris/exceptions/__init__.py

@@ -0,0 +1 @@
+from .exceptions import *

chunknorris/exceptions/__init__.pyi

@@ -0,0 +1 @@
+from .exceptions import *

chunknorris/exceptions/exceptions.py

@@ -0,0 +1,8 @@
+class ChunkNorrisException(Exception):
+    def __init__(self, message):
+        pass
+
+
+class ChunkSizeExceeded(Exception):
+    def __init__(self, message):
+        pass

chunknorris/exceptions/exceptions.pyi

@@ -0,0 +1,5 @@
+class ChunkNorrisException(Exception):
+    def __init__(self, message) -> None: ...
+
+class ChunkSizeExceeded(Exception):
+    def __init__(self, message) -> None: ...

chunknorris/types/types.py

@@ -0,0 +1,12 @@
+Chunk = str
+
+Chunks = list[Chunk]
+
+class TocTree():
+    id: int
+    text: str
+    level: int
+    line_index: int
+    content: str
+    parents: dict
+    children: list[dict]

chunknorris/types/types.pyi

@@ -0,0 +1,11 @@
+Chunk = str
+Chunks = list[Chunk]
+
+class TocTree:
+    id: int
+    text: str
+    level: int
+    line_index: int
+    content: str
+    parents: dict
+    children: list[dict]

chunknorris/utils/utils.py

@@ -0,0 +1,64 @@
+import re
+
+alphabets = "([A-Za-z])"
+prefixes = "(Mr|St|Mrs|Ms|Dr)[.]"
+suffixes = "(Inc|Ltd|Jr|Sr|Co)"
+starters = "(Mr|Mrs|Ms|Dr|Prof|Capt|Cpt|Lt|He\s|She\s|It\s|They\s|Their\s|Our\s|We\s|But\s|However\s|That\s|This\s|Wherever)"
+acronyms = "([A-Z][.][A-Z][.](?:[A-Z][.])?)"
+websites = "[.](com|net|org|io|gov|edu|me)"
+digits = "([0-9])"
+multiple_dots = r"\.{2,}"
+
+
+def split_into_sentences(text: str) -> list[str]:
+    """
+    Split the text into sentences.
+
+    If the text contains substrings "<prd>" or "<stop>", they would lead
+    to incorrect splitting because they are used as markers for splitting.
+
+    :param text: text to be split into sentences
+    :type text: str
+
+    :return: list of sentences
+    :rtype: list[str]
+    """
+    text = " " + text + "  "
+    text = text.replace("\n", " ")
+    text = re.sub(prefixes, "\\1<prd>", text)
+    text = re.sub(websites, "<prd>\\1", text)
+    text = re.sub(digits + "[.]" + digits, "\\1<prd>\\2", text)
+    text = re.sub(
+        multiple_dots, lambda match: "<prd>" * len(match.group(0)) + "<stop>", text
+    )
+    if "Ph.D" in text:
+        text = text.replace("Ph.D.", "Ph<prd>D<prd>")
+    text = re.sub("\s" + alphabets + "[.] ", " \\1<prd> ", text)
+    text = re.sub(acronyms + " " + starters, "\\1<stop> \\2", text)
+    text = re.sub(
+        alphabets + "[.]" + alphabets + "[.]" + alphabets + "[.]",
+        "\\1<prd>\\2<prd>\\3<prd>",
+        text,
+    )
+    text = re.sub(alphabets + "[.]" + alphabets + "[.]", "\\1<prd>\\2<prd>", text)
+    text = re.sub(" " + suffixes + "[.] " + starters, " \\1<stop> \\2", text)
+    text = re.sub(" " + suffixes + "[.]", " \\1<prd>", text)
+    text = re.sub(" " + alphabets + "[.]", " \\1<prd>", text)
+    if "”" in text:
+        text = text.replace(".”", "”.")
+    if '"' in text:
+        text = text.replace('."', '".')
+    if "!" in text:
+        text = text.replace('!"', '"!')
+    if "?" in text:
+        text = text.replace('?"', '"?')
+    text = text.replace(".", ".<stop>")
+    text = text.replace("?", "?<stop>")
+    text = text.replace("!", "!<stop>")
+    text = text.replace("<prd>", ".")
+    sentences = text.split("<stop>")
+    sentences = [s.strip() for s in sentences]
+    if sentences and not sentences[-1]:
+        sentences = sentences[:-1]
+
+    return sentences

chunknorris/utils/utils.pyi

@@ -0,0 +1,10 @@
+alphabets: str
+prefixes: str
+suffixes: str
+starters: str
+acronyms: str
+websites: str
+digits: str
+multiple_dots: str
+
+def split_into_sentences(text: str) -> list[str]: ...

chunknorris-0.0.1.dist-info/LICENCE

@@ -0,0 +1,15 @@
+ChunkNorris - A package for reliable chunking of documents
+Copyright (C) 2024 Wikit.ai
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as
+published by the Free Software Foundation, either version 3 of the
+License, or (at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program.  If not, see <https://www.gnu.org/licenses/>

chunknorris-0.0.1.dist-info/METADATA

@@ -0,0 +1,142 @@
+Metadata-Version: 2.1
+Name: chunknorris
+Version: 0.0.1
+Summary: A package for chunking documents from various formats
+Author-email: Wikit <dev@wikit.ai>
+Project-URL: Homepage, https://gitlab.com/wikit/research-and-development/chunk-norris
+Project-URL: Issues, https://gitlab.com/wikit/research-and-development/chunk-norris/-/issues
+Keywords: chunk,document,split,html,markdown,pdf,header
+Classifier: Natural Language :: English
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Framework :: Pytest
+Classifier: License :: OSI Approved :: GNU Affero General Public License v3
+Classifier: Topic :: Text Processing :: Markup :: Markdown
+Classifier: Topic :: Text Processing :: Markup :: HTML
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENCE
+Requires-Dist: markdownify >=0.11.6
+Requires-Dist: tiktoken >=0.5.2
+Requires-Dist: PyMuPDF >=1.23.16
+
+# Chunk Norris
+
+## Goal
+
+This project aims at improving the method of chunking documents from various sources (HTML, PDFs, ...)
+An optimized chunking method might lead to smaller chunks, meaning :
+- **Better relevancy of chunks** (and thus easier identification of useful chunks through embedding cosine similarity)
+- **Less errors** because of chunks exceeding the API limit in terms of number of tokens
+- **Less hallucinations** of generation models because of superfluous information in the prompt
+- **Reduced cost** as the prompt would have reduced size
+
+## Installation
+
+Using Pypi, just run the following command :
+```pip install chunknorris```
+
+## Chunkers
+
+The package features multiple ***chunkers*** that can be used independently depending on the type of document needed.
+
+All chunkers follow a similar logic :
+- Extract table of contents (= headers)
+- Build chunks using the text content of a part, and put the titles of the parts it belongs to on top
+
+![](images/chunk_method.png)
+
+### MarkdownChunkNorris
+
+This chunker is meant to be used **on markdown-formatted text**. 
+
+Note: When calling the chunker, **you need to specify the header style** of your markdown text ([ATX or Setext](https://golem.ph.utexas.edu/~distler/maruku/markdown_syntax.html#header)). By default it will consider "Setext" heading style.
+
+#### Usage
+
+```py
+from chunkers import MarkdownChunkNorris
+
+text = """
+# This is a header
+This is a text
+## This is another header
+And another text
+## With this final header
+And this last text
+"""
+chunker = MarkdownChunkNorris()
+header_style = "atx" # or "setext" depending on headers in your text
+chunks = chunker(text, header_style=header_style)
+```
+
+### HTMLChunkNorris
+
+This chunker is meant to be used **on html-formatted text**. Behind the scene, it uses markdownify to transform the text to markdown with "setex"-style headers and uses MarkdownChunkNorris to process it.
+
+#### Usage
+
+```py
+from chunkers import HTMLChunkNorris
+
+text = """
+<h1>This is 1st level heading</h1>
+<p>This is a test paragraph.</p>
+<h2>This is 2nd level heading</h2>
+<p>This is a test paragraph.</p>
+<h2>This is another level heading</h2>
+<p>This is another test paragraph.</p>
+"""
+hcn = HTMLChunkNorris()
+chunks = hcn(text)
+```
+
+### Advanced usage of chunkers
+
+Additionally, the chunkers can take a number of argument allowing to modifiy its behavior:
+
+```py
+from chunkers import MarkdownChunkNorris
+
+mystring = "# header\nThis is a markdown string"
+
+chunker = MarkdownChunkNorris() # or any other chunker
+chunks = chunker(
+    mystring,
+    max_title_level_to_use="h3",
+    max_chunk_word_length=200,
+    link_placement="in_sentence",
+    max_chunk_tokens=8191,
+    chunk_tokens_exceeded_handling="split",
+    min_chunk_wordcount=15,
+    )
+```
+
+***max_title_level_to_use*** 
+(str): The maximum (included) level of headers take into account for chunking. For example, if "h3" is set, then "h4" and "h5" titles won't be used. Must be a string of type "hx" with x being the title level. Defaults to "h4".
+
+***max_chunk_word_length***
+(int): The maximum size (soft limit, in words) a chunk can be. Chunk bigger that this size will be chunked using lower level headers, until no lower level headers are available. Defaults to 200.
+
+***link_placement***
+(str): How the links should be handled. Defaults to in_sentence.
+Options :
+- "remove" : text is kept but links are removed
+- "end_of_chunk" : adds a paragraph at the end of the chunk containing all the links
+- "in_sentence" : the links is added between parenthesis inside the sentence
+
+***max_chunk_tokens***
+(int): The hard maximum of number of token a chunk can be. Chunks bigger by this limit will be handler according to chunk_tokens_exceeded_handling. Defaults to 8191. 
+
+***chunk_tokens_exceeded_handling***
+(str): how the chunks bigger that than specified by max_chunk_tokens should be handled. Default to "raise_error".
+Options: 
+- "raise_error": raises an error, indicated the chunk could not be split according to headers
+- "split": split the chunks arbitrarily sothat each chunk has a size lower than max_chunk_tokens
+
+***min_chunk_wordcount***
+(int): Minimum number of words to consider keeping the chunks. Chunks with less words will be discarded. Defaults to 15.
+
+### PDFChunkNorris
+
+#TODO:

chunknorris-0.0.1.dist-info/RECORD

@@ -0,0 +1,27 @@
+chunknorris/__init__.py,sha256=GMddZWF-AhcCso9CLeuZzDYLH1wF_mWB_Z1UALPS6i4,102
+chunknorris/__init__.pyi,sha256=GMddZWF-AhcCso9CLeuZzDYLH1wF_mWB_Z1UALPS6i4,102
+chunknorris/chunkers/__init__.py,sha256=pm6GiSXWpn09xJP1xlKbnU7bJglY4l_3g-CBfr0adVE,100
+chunknorris/chunkers/__init__.pyi,sha256=Pn5uEk6TwEE_dMb1yzwPZeIcbwMB8zRPVdhKO8ntKg0,142
+chunknorris/chunkers/html_chunknorris.py,sha256=X9LYqd2maej_7bFUa0pY_CAGWdanM_7_tUY6yoyRUR0,663
+chunknorris/chunkers/html_chunknorris.pyi,sha256=M2xQe5xvoysJCHB1HtGpQ7jv4sj-JFB7vpOJegw0LI4,250
+chunknorris/chunkers/markdown_chunknorris.py,sha256=ZdB1NpZASTDiEMfv3CBR8FRvrGl8Rd6r-GZcmBqCRBk,20268
+chunknorris/chunkers/markdown_chunknorris.pyi,sha256=kk8I3La__qcvCyRVtwxD-5NuxFn721LnZRZhUVDx2WM,1486
+chunknorris/custom_chunkers/__init__.py,sha256=qMlMdmT4KHJnyRr42ZWT0g-RnhDdwDCENdJ9BztLF00,48
+chunknorris/custom_chunkers/__init__.pyi,sha256=sley_LOReM3qVfDtSqKa_13soCe32eDBHGZUgAoiASI,68
+chunknorris/custom_chunkers/wikit_chunknorris.py,sha256=IA-7Xb0ogXmxgqp8ITJM-P7n4clvVaSRAV_DRtT92nE,5887
+chunknorris/custom_chunkers/wikit_chunknorris.pyi,sha256=YlZTYx8iAVqNHbjGw8kbSQZ_kw550qPzYjUNNTuPdok,652
+chunknorris/exceptions/__init__.py,sha256=nWnJwEphtAjFsYg_dIYrrbWdkIB0G2Nu89d7d6DJGKc,25
+chunknorris/exceptions/__init__.pyi,sha256=y6AJu3xWHud92ZK_pfU3WzDj8gLIYvXfFNJ-phZmjJo,26
+chunknorris/exceptions/exceptions.py,sha256=OUj1qxRcQU0RdYTJMs3L2abOarD16nhjkAPTMRhr9V0,169
+chunknorris/exceptions/exceptions.pyi,sha256=HZeW31MY2y1BijHfSlHElqpTkn1jxawO6LvdMa1Ft-s,166
+chunknorris/types/types.py,sha256=ii8PICCROmuRSZSH2lM_iAr4PYBMBEI5ePm1S9eG8Cw,173
+chunknorris/types/types.pyi,sha256=iW3elE5TLioEZLL6HUSPP9LH8izJCWE0JY-73eqAe1I,170
+chunknorris/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+chunknorris/utils/__init__.pyi,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+chunknorris/utils/utils.py,sha256=nA-qCm22dAYRtynCJwWa7S9AK8SCYoDzG6Xh296uX1g,2288
+chunknorris/utils/utils.pyi,sha256=ApwdRg48ZViJhdr-g91E6F76Qdj5EBfqVVKHQ6qtSJI,171
+chunknorris-0.0.1.dist-info/LICENCE,sha256=Z4Dj4xyHOkzx-Ggiig2yGWeOy-gnR7OF_hkJ6kCcCNw,720
+chunknorris-0.0.1.dist-info/METADATA,sha256=hGylMSJp_O-uaIRAWXGybjRILsPxdeGF14a59eMpEZk,5200
+chunknorris-0.0.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+chunknorris-0.0.1.dist-info/top_level.txt,sha256=zYBuOKW7poeXPcU-OLQF5PBbdbrw2aZbWgtPuge0x7Y,12
+chunknorris-0.0.1.dist-info/RECORD,,

chunknorris-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: bdist_wheel (0.43.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

chunknorris-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+chunknorris