PyTorch serialized list

pytorch_serialized_list.py

import torch


# Modified from https://github.com/ppwwyyxx/RAM-multiprocess-dataloader
class PyTorchStrList:
    def __init__(self, items: list[str]):
        data = [torch.frombuffer(x.encode(), dtype=torch.uint8) for x in items]
        lengths = [0] + [x.shape[0] for x in data]
        self.data = torch.cat(data, 0)
        self.index = torch.tensor(lengths).cumsum_(0)

    def __getitem__(self, i: int) -> str:
        assert i >= 0
        start, end = self.index[i : i + 2].tolist()
        return memoryview(self.data[start:end].numpy()).tobytes().decode()

    def __len__(self) -> int:
        return self.index.shape[0] - 1