INNER CODE UNIT · Python

count_all_token

HIT-SCIR/Chinese-Mixtral-8x7B · data/utils.py:30

def count_all_token(
        sequence_length: int = 2048,
):
    # TODO refactor with `count_token`
    statistics = {}

    with open("./data/datasets.toml", "rb") as f:
        ds_info = tomli.load(f)

    for ds_name, ds in ds_info.items():
        for split in ds["splits"]:
            try:
                ds_path = (ds["root"].format(DATA_DIR="./data", name=ds_name) + "/" +
                           ds["encoded"].format(name=ds_name, split=split))
                print(f"Loading {ds_name} ...")
                dataset = load_from_disk(ds_path)

                statistics[ds_name + "-" + split] = len(dataset) * sequence_length / 10 ** 9

View source record →