Method add_window_ids

langroid/parsing/parser.py:168–202

Chunks may belong to multiple docs, but for each doc, they appear consecutively. Add window_ids in metadata

(self, chunks: List[Document])

        return self.tokenizer.decode(tokens[:max_tokens])

    def add_window_ids(self, chunks: List[Document]) -> None:
        """Chunks may belong to multiple docs, but for each doc,
        they appear consecutively. Add window_ids in metadata"""

        # discard empty chunks
        chunks = [c for c in chunks if c.content.strip() != ""]
        if len(chunks) == 0:
            return
        # The original metadata.id (if any) is ignored since it will be same for all
        # chunks and is useless. We want a distinct id for each chunk.
        # ASSUMPTION: all chunks c of a doc have same c.metadata.id !
        orig_ids = [c.metadata.id for c in chunks]
        ids = [ObjectRegistry.new_id() for c in chunks]
        id2chunk = {id: c for id, c in zip(ids, chunks)}

        # group the ids by orig_id
        # (each distinct orig_id refers to a different document)
        orig_id_to_ids: Dict[str, List[str]] = {}
        for orig_id, id in zip(orig_ids, ids):
            if orig_id not in orig_id_to_ids:
                orig_id_to_ids[orig_id] = []
            orig_id_to_ids[orig_id].append(id)

        # now each orig_id maps to a sequence of ids within a single doc

        k = self.config.n_neighbor_ids
        for orig, ids in orig_id_to_ids.items():
            # ids are consecutive chunks in a single doc
            n = len(ids)
            window_ids = [ids[max(0, i - k) : min(n, i + k + 1)] for i in range(n)]
            for i, _ in enumerate(ids):
                c = id2chunk[ids[i]]
                c.metadata.window_ids = window_ids[i]
                c.metadata.id = ids[i]
                c.metadata.is_chunk = True

    def split_simple(self, docs: List[Document]) -> List[Document]:
        if len(self.config.separators) == 0: