document_preprocessor.py

#!/usr/bin/env python3
'''Document preprocessing module'''

import argparse
import json
import logging
import operator
import os
import re
from functools import reduce
from typing import Callable, Sequence, TypeAlias, Generator

import pandoc  # type: ignore
import pandoc.types as T  # type: ignore
import tiktoken

from type_hints import DocumentJson, DocumentBlockInfo, DocumentChunkInfo, DocumentChunkBlockMapItem

PandocElement: TypeAlias = Sequence
PandocElementPath: TypeAlias = Sequence[int]


def get_element(li: PandocElement, indices: PandocElementPath):
    '''Access an element in a nested list by a list of indices'''
    return reduce(operator.getitem, indices, li)


def get_markdown(elt: PandocElement) -> str:
    '''Convert a Pandoc element to Markdown'''
    return pandoc.write(elt, options=[
        '--wrap=none',
        '--to=markdown+hard_line_breaks-smart-raw_attribute-header_attributes-link_attributes',
    ])


def list_startswith(li: Sequence, prefix: Sequence) -> bool:
    '''Check if a list starts with a prefix'''
    return os.path.commonprefix([li, prefix]) == prefix  # type: ignore


def process_document(
    elements: PandocElement,
    tokenizer_function: Callable,
    token_limit=500
) -> DocumentJson:
    '''Document preprocessing -- main code'''

    def dfs_to_block(elt: PandocElement, path: PandocElementPath) -> Generator[PandocElementPath, None, None]:
        '''
        Traverse the document tree and flatten it into a list of blocks.
        '''

        if isinstance(elt, (
            T.Para, T.Plain, T.Table, T.Header, T.HorizontalRule, T.CodeBlock, T.Figure, T.BlockQuote,
        )):
            # Leaf element, do not split further
            yield path
        elif isinstance(elt, T.BulletList):
            # Break the list into items
            for i, child in enumerate(elt[0]):
                yield from dfs_to_block(child, tuple(path) + (0, i))
        elif isinstance(elt, T.OrderedList):
            # Break the list into items
            for i, child in enumerate(elt[1]):
                yield from dfs_to_block(child, tuple(path) + (1, i))
        elif isinstance(elt, list):
            for i, child in enumerate(elt):
                child_path = tuple(path) + (i,)
                yield from dfs_to_block(child, child_path)
        else:
            raise NotImplementedError(f"Unsupported element: {type(elt)}")

    def init_blocks() -> tuple[list[PandocElementPath], list[str]]:
        blocks = list(dfs_to_block(elements, ()))
        block_texts = []

        for i, elt_idx in enumerate(blocks):
            # Trim blocks' element indices to shortest unique prefix
            prev_elt_idx = blocks[i - 1] if i > 0 else ()
            next_elt_idx = blocks[i + 1] if i < len(blocks) - 1 else ()

            prefix1 = os.path.commonprefix([prev_elt_idx, elt_idx])  # type: ignore
            prefix2 = os.path.commonprefix([next_elt_idx, elt_idx])  # type: ignore

            blocks[i] = elt_idx[:max(len(prefix1), len(prefix2)) + 1]

            # Precompute block strings and lengths
            elt = get_element(elements, blocks[i])
            block_texts.append(get_markdown(elt).strip())

        return blocks, block_texts

    def get_block_context_limit(blocks: list[PandocElementPath]) -> list[int]:
        # Determine to what subsequent block each block is part of a context
        context_limits = []

        for i, elt_idx in enumerate(blocks):
            elt = get_element(elements, elt_idx)
            limit = i

            if isinstance(elt, T.Header):
                heading_level = elt[0]

                for j in range(i + 1, len(blocks)):
                    elt_j = get_element(elements, blocks[j])

                    if (list_startswith(blocks[j], elt_idx[:-1])
                        and not(isinstance(elt_j, T.Header) and elt_j[0] <= heading_level)):
                        limit = j
                    else:
                        break
            else:
                text = get_markdown(elt)

                if (text.strip().endswith(':')
                    and i + 1 < len(blocks)
                    and list_startswith(blocks[i + 1], elt_idx[:-1])):

                    limit = i + 1

            context_limits.append(limit)

        return context_limits

    def init_chunks(blocks: list[PandocElementPath], block_texts: list[str]) -> list[frozenset[int]]:
        context_limit = get_block_context_limit(blocks)

        # Precompute block lengths
        block_lengths = [len(tokenizer_function(text)) for text in block_texts]

        # Create chunks
        chunks: list[frozenset[int]] = []

        for start_block_idx, elt_idx in enumerate(blocks):
            context = [j for j, limit in enumerate(context_limit[:start_block_idx]) if start_block_idx <= limit]
            n_tokens = block_lengths[start_block_idx] + sum(block_lengths[j] for j in context)

            # See if we can merge subsequent blocks to form longer chunks
            end_block_idx = start_block_idx + 1

            while True:
                elt_idx = tuple(elt_idx[:-1]) + (elt_idx[-1] + 1,)
                next_end = end_block_idx

                for j in range(end_block_idx, len(blocks)):
                    # Stay in the same level
                    if list_startswith(blocks[j], elt_idx):
                        next_end = max(next_end, context_limit[j] + 1)

                if end_block_idx == next_end:
                    break

                n_tokens += sum(block_lengths[j] for j in range(end_block_idx, next_end))

                if n_tokens < token_limit:
                    end_block_idx = next_end
                else:
                    break

            new_chunk = frozenset(range(start_block_idx, end_block_idx)) | frozenset(context)

            if all(new_chunk - ch for ch in chunks):
                chunks = [ch for ch in chunks if ch - new_chunk]
                chunks.append(new_chunk)

        return chunks

    def get_chunk_text(blocks: list[PandocElementPath], chunk: frozenset[int]):
        non_leaf_paths: set[Sequence[int]] = {()}
        leaf_paths = set()

        for blk_idx in chunk:
            elt_idx = blocks[blk_idx]

            for i in range(len(elt_idx)):
                non_leaf_paths.add(elt_idx[:i])

            leaf_paths.add(elt_idx)

        def dfs(elt, path):
            if path in leaf_paths:
                return elt
            elif path in non_leaf_paths or isinstance(elt, list):
                children = []

                for i, child in enumerate(elt):
                    children.append(dfs(child, path + (i,)))

                if isinstance(elt, list):
                    return children
                elif isinstance(elt, (T.Block, T.Inline)):
                    return elt.__class__(*children)
                else:
                    assert False
            else:
                placeholder_str = "%PLACEHOLDER%"

                if isinstance(elt, tuple):
                    return elt
                elif isinstance(elt, T.Block):
                    return T.Plain([T.Str(placeholder_str)])
                elif isinstance(elt, T.Inline):
                    return T.Str(placeholder_str)
                else:
                    assert False

        elements_with_context = dfs(elements, ())
        text = get_markdown(elements_with_context)

        previous_line = "%PLACEHOLDER%"
        processed_lines = []

        for line in text.split('\n'):
            if line != previous_line:
                processed_lines.append(line.replace("%PLACEHOLDER%", "..."))

            if line:
                previous_line = line

        # Strip extra newlines
        text = "\n".join(processed_lines).strip() + '\n'
        text = re.sub(r'[\r\n][\r\n]{2,}', '\n\n', text)

        return text

    def process() -> DocumentJson:
        # a block is defined by a unique element
        # a block never contains another block
        blocks, block_texts = init_blocks()

        # a chunk is defined by a list of block indices
        # chunks can overlap
        chunks = init_chunks(blocks, block_texts)

        block_info: list[DocumentBlockInfo] = []

        for element_indices, text in zip(blocks, block_texts):
            block_info.append({
                "element_indices": element_indices,
                "text": text,
            })

        chunk_info: list[DocumentChunkInfo] = []

        for c in chunks:
            text = get_chunk_text(blocks, c)

            # text range -> block index
            block_map: list[DocumentChunkBlockMapItem] = []

            for blk_idx in sorted(c):
                block_text = block_texts[blk_idx]
                re_pattern = r'\s+'.join(map(re.escape, block_text.split()))

                if m := re.search(re_pattern, text):
                    i_start, i_end = m.span()
                    block_map.append({
                        "index": blk_idx,
                        "text_range": (i_start, i_end),
                    })
                else:
                    raise ValueError(f"Block not found: {block_text}")

            chunk_info.append({
                "block_map": block_map,
                "text": text,
            })

        return {
            "blocks": block_info,
            "chunks": chunk_info,
        }

    return process()


def main():
    logging.basicConfig(format='%(asctime)s [%(levelname)s] %(message)s', level=logging.INFO)

    parser = argparse.ArgumentParser()
    parser.add_argument("workdirs", nargs="+", help="Input directories")
    parser.add_argument("--token-limit", type=int, default=500, help="Token limit")
    parser.add_argument("--model", default="gpt-4o", help="Model for which to encode the text")
    args = parser.parse_args()

    encoder = tiktoken.encoding_for_model(args.model)

    for d in args.workdirs:
        logging.info("Processing %s ...", d)

        _, elements = pandoc.read(file=os.path.join(d, 'content.md'))

        doc = process_document(elements, encoder.encode, args.token_limit)

        with open(os.path.join(d, 'document.json'), 'w', encoding='utf-8') as fout:
            json.dump(doc, fout, indent=2, sort_keys=True)



if __name__ == "__main__":
    main()