count_token.py

from multiprocessing import Pool

import fire
import jsonlines as jl
from loguru import logger
from tqdm import tqdm
from transformers import AutoTokenizer

TOKENIZER = None


def count_token_single(sample, text_keys):
    global TOKENIZER
    num = 0
    for key in text_keys:
        num += len(TOKENIZER.tokenize(sample[key]))
    return num


def prepare_tokenizer(tokenizer_method):
    global TOKENIZER
    logger.info('Loading tokenizer from HuggingFace...')
    TOKENIZER = AutoTokenizer.from_pretrained(tokenizer_method,
                                              trust_remote_code=True)


def main(data_path,
         text_keys='text',
         tokenizer_method='EleutherAI/pythia-6.9b-deduped',
         num_proc=1):
    """
    Count the number of tokens for given dataset and tokenizer.

    :param data_path: path to the input dataset. Only support 'jsonl' now.
    :param text_keys: field keys that will be considered into token counts.
    :param tokenizer_method: name of the Hugging Face tokenizer.
    :param num_proc: number of processes to count tokens.
    """
    prepare_tokenizer(tokenizer_method)

    if isinstance(text_keys, str):
        text_keys = [text_keys]

    with jl.open(data_path) as reader:
        token_count = 0
        result_list = []
        with Pool(num_proc) as p:
            for sample in tqdm(reader):
                result_list.append(
                    p.apply_async(count_token_single,
                                  args=(
                                      sample,
                                      text_keys,
                                  )))
            for res in tqdm(result_list):
                token_count += res.get()

        logger.info(f'Total num of tokens: {token_count}')


if __name__ == '__main__':
    fire.Fire(main)