out = MDSWriter(
        columns={"input_ids": f"ndarray:int32:{args.seq_len}",
                 "token_type_ids": f"ndarray:int8:{args.seq_len}",
                 "attention_mask": f"ndarray:int8:{args.seq_len}",
                 "special_tokens_mask": f"ndarray:int8:{args.seq_len}"},
        out=out_path,
        compression=None
    )

def parse_data_2_mds_format(tokenized_dataset):
    input_ids = np.array(tokenized_dataset['input_ids']).astype(np.int32)
    token_type_ids = np.array(tokenized_dataset['token_type_ids']).astype(np.int8)
    attention_mask = np.array(tokenized_dataset['attention_mask']).astype(np.int8)
    special_tokens_mask = np.array(tokenized_dataset['special_tokens_mask']).astype(np.int8)
    return {'input_ids': input_ids, 'token_type_ids': token_type_ids,
            'attention_mask': attention_mask, 'special_tokens_mask': special_tokens_mask}

 with Pool(processes=args.mds_num_workers) as inner_pool:
                with tqdm(total=len(tokenized_datasets), desc="Writing Out MDS File") as pbar:
                    for result in inner_pool.imap(parse_data_2_mds_format, tokenized_datasets):
                        out.write(result)
                        pbar.update()