Method remove_overlaps

langroid/vector_store/base.py:347–404

Given a collection of windows, where each window is a sequence of ids, identify groups of overlapping windows, and for each overlapping group, order the chunk-ids using topological sort so they appear in the original order in the text. Args: wind

(windows: List[List[str]])

    @staticmethod
    def remove_overlaps(windows: List[List[str]]) -> List[List[str]]:
        """
        Given a collection of windows, where each window is a sequence of ids,
        identify groups of overlapping windows, and for each overlapping group,
        order the chunk-ids using topological sort so they appear in the original
        order in the text.

        Args:
            windows (List[int|str]): List of windows, where each window is a
                sequence of ids.

        Returns:
            List[int|str]: List of windows, where each window is a sequence of ids,
                and no two windows overlap.
        """
        ids = set(id for w in windows for id in w)
        # id -> {win -> # pos}
        id2win2pos: Dict[str, Dict[int, int]] = {id: {} for id in ids}

        for i, w in enumerate(windows):
            for j, id in enumerate(w):
                id2win2pos[id][i] = j

        n = len(windows)
        # relation between windows:
        order = np.zeros((n, n), dtype=np.int8)
        for i, w in enumerate(windows):
            for j, x in enumerate(windows):
                if i == j:
                    continue
                if len(set(w).intersection(x)) == 0:
                    continue
                id = list(set(w).intersection(x))[0]  # any common id
                if id2win2pos[id][i] > id2win2pos[id][j]:
                    order[i, j] = -1  # win i is before win j
                else:
                    order[i, j] = 1  # win i is after win j

        # find groups of windows that overlap, like connected components in a graph
        groups = components(np.abs(order))

        # order the chunk-ids in each group using topological sort
        new_windows = []
        for g in groups:
            # find total ordering among windows in group based on order matrix
            # (this is a topological sort)
            _g = np.array(g)
            order_matrix = order[_g][:, _g]
            ordered_window_indices = topological_sort(order_matrix)
            ordered_window_ids = [windows[i] for i in _g[ordered_window_indices]]
            flattened = [id for w in ordered_window_ids for id in w]
            flattened_deduped = list(dict.fromkeys(flattened))
            # Note we are not going to split these, and instead we'll return
            # larger windows from concatenating the connected groups.
            # This ensures context is retained for LLM q/a
            new_windows += [flattened_deduped]

        return new_windows