Class LMDBDataset

codegeex/mindspore/src/sat_dataset.py:18–53

class LMDBDataset(Dataset):
    def __init__(self, path, process_fn=None):
        import lmdb

        self.path = path
        self.env = lmdb.open(
            path,
            max_readers=32,
            readonly=True,
            lock=False,
            readahead=False,
            meminit=False,
        )
        self.process_fn = process_fn
        if not self.env:
            raise IOError("Cannot open lmdb dataset", path)

        with self.env.begin(write=False) as txn:
            self.length = int(txn.get("length".encode("utf-8")).decode("utf-8"))

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        # print(f"Get {self.path}: {idx}")
        with self.env.begin(write=False) as txn:
            key = str(idx).encode("utf-8")
            # row = pickle.loads(txn.get(key))
            try:
                row = pickle.loads(txn.get(key))
            except TypeError:
                raise IndexError("Index out of range")
            if self.process_fn:
                return self.process_fn(row)
            else:
                return row


class PadDataset(Dataset):