utils/commons/indexed_datasets.py

import pickle
from bisect import bisect
from copy import deepcopy
import numpy as np
import gzip


def int2bytes(i: int, *, signed: bool = False) -> bytes:
    length = ((i + ((i * signed) < 0)).bit_length() + 7 + signed) // 8
    return i.to_bytes(length, byteorder='little', signed=signed)


def bytes2int(b: bytes, *, signed: bool = False) -> int:
    return int.from_bytes(b, byteorder='little', signed=signed)


def load_index_data(data_file):
    index_data_size = bytes2int(data_file.read(32))
    index_data = data_file.read(index_data_size)
    index_data = pickle.loads(index_data)
    data_offsets = deepcopy(index_data['offsets'])
    id2pos = deepcopy(index_data.get('id2pos', {}))
    meta = deepcopy(index_data.get('meta', {}))
    return data_offsets, id2pos, meta


class IndexedDataset:
    def __init__(self, path, unpickle=True):
        self.path = path
        self.root_data_file = open(f"{path}.data", 'rb', buffering=-1)
        try:
            self.byte_offsets, self.id2pos, self.meta = load_index_data(self.root_data_file)
            self.data_files = [self.root_data_file]
        except:
            self.__init__old(path)
            self.meta = {}
        self.gzip = self.meta.get('gzip', False)
        if 'chunk_begin' not in self.meta:
            self.meta['chunk_begin'] = [0]
        for i in range(len(self.meta['chunk_begin'][1:])):
            self.data_files.append(open(f"{self.path}.{i + 1}.data", 'rb'))
        self.unpickle = unpickle

    def __init__old(self, path):
        self.path = path
        index_data = np.load(f"{path}.idx", allow_pickle=True).item()
        self.byte_offsets = index_data['offsets']
        self.id2pos = index_data.get('id2pos', {})
        self.data_files = [open(f"{path}.data", 'rb', buffering=-1)]

    def __getitem__(self, i):
        if self.id2pos is not None and len(self.id2pos) > 0:
            i = self.id2pos[i]
        self.check_index(i)
        
        # chunk_id = bisect(self.meta['chunk_begin'][1:], self.byte_offsets[i])
        # if chunk_id == 0:
        #     data_file = open(f"{self.path}.data", 'rb', buffering=-1)
        # else:
        #     data_file = open(f"{self.path}.{chunk_id}.data", 'rb', buffering=-1)
        # data_file.seek(self.byte_offsets[i] - self.meta['chunk_begin'][chunk_id])
        # b = data_file.read(self.byte_offsets[i + 1] - self.byte_offsets[i])
        # data_file.close()
        
        chunk_id = bisect(self.meta['chunk_begin'][1:], self.byte_offsets[i])
        data_file = self.data_files[chunk_id]
        data_file.seek(self.byte_offsets[i] - self.meta['chunk_begin'][chunk_id])
        b = data_file.read(self.byte_offsets[i + 1] - self.byte_offsets[i])

        unpickle = self.unpickle
        if unpickle:
            if self.gzip:
                b = gzip.decompress(b)
            item = pickle.loads(b)
        else:
            item = b
        return item

    def __del__(self):
        for data_file in self.data_files:
            data_file.close()

    def check_index(self, i):
        if i < 0 or i >= len(self.byte_offsets) - 1:
            raise IndexError('index out of range')

    def __len__(self):
        return len(self.byte_offsets) - 1

    def __iter__(self):
        self.iter_i = 0
        return self

    def __next__(self):
        if self.iter_i == len(self):
            raise StopIteration
        else:
            item = self[self.iter_i]
            self.iter_i += 1
            return item


class IndexedDatasetBuilder:
    def __init__(self, path, append=False, max_size=1024 * 1024 * 1024 * 64,
                 default_idx_size=1024 * 1024 * 16, gzip=False):
        self.path = self.root_path = path
        self.default_idx_size = default_idx_size
        if append:
            self.data_file = open(f"{path}.data", 'r+b')
            self.data_file.seek(0)
            self.byte_offsets, self.id2pos, self.meta = load_index_data(self.data_file)
            self.data_file.seek(0)
            self.data_file.write(bytes(default_idx_size))
            self.data_file.seek(self.byte_offsets[-1])
            self.gzip = self.meta['gzip']
        else:
            self.data_file = open(f"{path}.data", 'wb')
            self.data_file.seek(default_idx_size)
            self.byte_offsets = [default_idx_size]
            self.id2pos = {}
            self.meta = {}
            self.meta['chunk_begin'] = [0]
            self.gzip = self.meta['gzip'] = gzip
        self.root_data_file = self.data_file
        self.max_size = max_size
        self.data_chunk_id = 0

    def add_item(self, item, id=None, use_pickle=True):
        if self.byte_offsets[-1] > self.meta['chunk_begin'][-1] + self.max_size:
            if self.data_file != self.root_data_file:
                self.data_file.close()
            self.data_chunk_id += 1
            self.data_file = open(f"{self.path}.{self.data_chunk_id}.data", 'wb')
            self.data_file.seek(0)
            self.meta['chunk_begin'].append(self.byte_offsets[-1])
        if not use_pickle:
            s = item
        else:
            s = pickle.dumps(item)
            if self.gzip:
                s = gzip.compress(s, 1)
        bytes = self.data_file.write(s)
        if id is not None:
            self.id2pos[id] = len(self.byte_offsets) - 1
        self.byte_offsets.append(self.byte_offsets[-1] + bytes)

    def finalize(self):
        self.root_data_file.seek(0)
        s = pickle.dumps({'offsets': self.byte_offsets, 'id2pos': self.id2pos, 'meta': self.meta})
        assert len(s) < self.default_idx_size, (len(s), self.default_idx_size)
        len_bytes = int2bytes(len(s))
        self.root_data_file.write(len_bytes)
        self.root_data_file.seek(32)
        self.root_data_file.write(s)
        self.root_data_file.close()
        try:
            self.data_file.close()
        except:
            pass


if __name__ == "__main__":
    import random
    from tqdm import tqdm

    # builder = IndexedDatasetBuilder(ds_path, append=True)
    # for i in tqdm(range(size)):
    #     builder.add_item(items[i], i + size)
    # builder.finalize()
    # ds = IndexedDataset(ds_path)
    # for i in tqdm(range(1000)):
    #     idx = random.randint(size, 2 * size - 1)
    #     assert (ds[idx]['a'] == items[idx - size]['a']).all()
    #     idx = random.randint(0, size - 1)
    #     assert (ds[idx]['a'] == items[idx]['a']).all()

    ds_path = '/tmp/indexed_ds_example'
    size = 100
    items = [{"a": np.random.normal(size=[10000, 10]),
              "b": np.random.normal(size=[10000, 10])} for i in range(size)]
    builder = IndexedDatasetBuilder(ds_path, max_size=1024 * 1024 * 40)
    builder.meta['lengths'] = [1, 2, 3]
    for i in tqdm(range(size)):
        builder.add_item(pickle.dumps(items[i]), i, use_pickle=False)
    builder.finalize()
    ds = IndexedDataset(ds_path)
    assert ds.meta['lengths'] == [1, 2, 3]
    for i in tqdm(range(1000)):
        idx = random.randint(0, size - 1)
        assert (ds[idx]['a'] == items[idx]['a']).all()

    # builder = IndexedDataset2Builder(ds_path, append=True)
    # builder.meta['lengths'] = [1, 2, 3, 5, 6, 7]
    # for i in tqdm(range(size)):
    #     builder.add_item(items[i], i + size)
    # builder.finalize()
    # ds = IndexedDataset2(ds_path)
    # assert ds.meta['lengths'] == [1, 2, 3, 5, 6, 7]
    # for i in tqdm(range(1000)):
    #     idx = random.randint(size, 2 * size - 1)
    #     assert (ds[idx]['a'] == items[idx - size]['a']).all()
    #     idx = random.randint(0, size - 1)
    #     assert (ds[idx]['a'] == items[idx]['a']).all()