grn/utils/compress_tokens.py

import torch
import numpy as np

def save_packed_tensor(filename, tensor):
    """use np.savez_compressed to save compressed tensor and its shape"""
    if tensor.dtype != torch.bool:
        raise TypeError("Input tensor must be of dtype torch.bool")
    shape_array = np.array(tensor.shape)
    packed_data = np.packbits(tensor.numpy())
    np.savez_compressed(filename, shape=shape_array, data=packed_data)

def load_packed_tensor(filename):
    """read .npz file and decompress tensor"""
    with np.load(filename) as loader:
        shape = loader['shape']
        packed_data = loader['data']
    numel = np.prod(shape)
    unpacked_data = np.unpackbits(packed_data, count=numel)
    restored_tensor = torch.from_numpy(unpacked_data.reshape(shape))
    return restored_tensor