def _partition_indexer()

in aidial_adapter_bedrock/llm/truncate_prompt.py [0:0]


def _partition_indexer(chunks: List[int]) -> Callable[[int], List[int]]:
    """
    Returns a function that maps an index to indices of its partition.
    """
    mapping: dict[int, List[int]] = {}
    offset = 0
    for size in chunks:
        chunk = list(range(offset, offset + size))
        for idx in range(size):
            mapping[offset + idx] = chunk
        offset += size

    return mapping.__getitem__