data_packager

Author: Heli Qi Affiliation: NAIST Date: 2022.11

main(src_path, feat_type, comp_chunk_ext, chunk_size=None, ncpu=8, remove_ori_data=False)

Parameters:

Name	Type	Description	Default
src_path	str		required
feat_type	str		required
comp_chunk_ext	str		required
chunk_size	int or str		None
ncpu	int		8
remove_ori_data	bool		False

Returns:

Source code in speechain/datasets/pyscripts/data_packager.py

def main(
    src_path: str,
    feat_type: str,
    comp_chunk_ext: str,
    chunk_size: int or str = None,
    ncpu: int = 8,
    remove_ori_data: bool = False,
):
    """

    Args:
        src_path:
        feat_type:
        comp_chunk_ext:
        chunk_size:
        ncpu:
        remove_ori_data:

    Returns:

    """
    # --- 0. Information Initialization --- #
    # compressed data folder creating
    assert comp_chunk_ext in [
        "npz",
        "hdf5",
    ], f"Data compression file extension should be one of {['npz', 'hdf5']}, but got {comp_chunk_ext}."
    comp_save_path = os.path.join(src_path, f"{feat_type}_{comp_chunk_ext}")
    os.makedirs(comp_save_path, exist_ok=True)

    # data path reading
    idx2data = os.path.join(src_path, f"idx2{feat_type}")
    if os.path.exists(idx2data):
        # ndmin=2 keeps the loaded array 2-dimensional even when the file has only one line
        idx2data = np.loadtxt(idx2data, dtype=str, delimiter=" ", ndmin=2)
        idx2data = dict(zip(idx2data[:, 0], idx2data[:, 1]))
    else:
        raise RuntimeError(
            "Please create 'idx2wav' or 'idx2feat' before data compression!"
        )

    # data length reading
    idx2data_len = os.path.join(src_path, f"idx2{feat_type}_len")
    # sort idx2data in descending order by their length if data_len is given
    if os.path.exists(idx2data_len):
        # ndmin=2 keeps the loaded array 2-dimensional even when the file has only one line
        idx2data_len = np.loadtxt(idx2data_len, dtype=str, delimiter=" ", ndmin=2)
        idx2data_len = dict(sorted(idx2data_len, key=lambda x: int(x[1]), reverse=True))
        idx2data = {idx: idx2data[idx] for idx in idx2data_len.keys()}
    else:
        idx2data_len = None
        warnings.warn(
            f"Data length file 'idx2{feat_type}_len' is not found in {src_path}, "
            f"so the data compressed into the same chunk may not have the similar lengths."
        )

    # different default values for fixed-length chunk and fixed-number chunk
    if chunk_size is None:
        chunk_size = "40m" if idx2data_len is not None else "2k"
    # turn the readable value to the raw integer value
    if isinstance(chunk_size, str):
        try:
            chunk_size = parse_readable_number(chunk_size)
        except (ValueError, AssertionError) as e:
            raise ValueError(
                f"Invalid --chunk_size {chunk_size!r}: {e}. Expected a plain integer "
                "or a readable number such as '2k', '40m', '1b500m' "
                "(units: b=1e9, m=1e6, k=1e3, h=1e2)."
            ) from e

    # --- 1. Collect Chunk Information --- #
    # loop each wav file
    chunk_dict, curr_chunk_num, curr_chunk_size = dict(), 0, 0
    for idx, data_path in idx2data.items():
        # create the current chunk sub-dict
        if curr_chunk_num not in chunk_dict.keys():
            chunk_dict[curr_chunk_num] = dict()

        # update the current data instance to the current chunk
        chunk_dict[curr_chunk_num][idx] = idx2data[idx]
        curr_chunk_size += int(idx2data_len[idx]) if idx2data_len is not None else 1

        # move to the next chunk if the current chunk is full
        if curr_chunk_size >= chunk_size:
            curr_chunk_num += 1
            curr_chunk_size = 0

    # --- 2. Write the collected data instances into compressed chunk files --- #
    chunk_list = list(chunk_dict.items())
    with Pool(ncpu) as executor:
        save_chunk_func = partial(
            save_chunk,
            feat_type=feat_type,
            comp_chunk_ext=comp_chunk_ext,
            save_path=comp_save_path,
            remove_ori_data=remove_ori_data,
        )
        idx2data_comp_chunk = executor.starmap(save_chunk_func, chunk_list)

    total_chunk_num = len(chunk_list)
    chunk_memories = [
        os.path.getsize(
            os.path.join(comp_save_path, f"chunk_{chunk_num}.{comp_chunk_ext}")
        )
        for chunk_num in range(total_chunk_num)
    ]
    print(
        f"Successfully save {len(chunk_list)} chunk files in {comp_save_path}!\n"
        f"Average size: {get_readable_memory(sum(chunk_memories) / total_chunk_num)}; "
        f"Total size: {get_readable_memory(sum(chunk_memories))}."
    )

    # --- 3. Save the new data addressed after compression --- #
    idx2data_comp = dict()
    for chunk in idx2data_comp_chunk:
        idx2data_comp.update(chunk.items())
    np.savetxt(
        os.path.join(src_path, f"idx2{feat_type}_{comp_chunk_ext}"),
        sorted(idx2data_comp.items(), key=lambda x: x[0]),
        fmt="%s",
    )