File size: 1,383 Bytes
31f25cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 | import json
from transformers import AutoTokenizer
from tqdm import tqdm
meta_json = './shell/playground/data/meta/sft_opc436k_scale_math_1m_smoltalk_1m_tulu_1m.json'
with open(meta_json, "r", encoding="utf-8") as file:
meta = json.load(file)
tokenizer = AutoTokenizer.from_pretrained(
'./shell/playground/ckpt/Qwen2.5-3B',
add_eos_token=True, # for |im_end|
trust_remote_code=True,
use_fast=False
)
token_dict = {}
for data_name, data_item in meta.items():
annotation = data_item['annotation']
print(f'processing... {annotation}')
anno_token_cnt = 0
with open(annotation, 'r') as f:
raw_data = f.readlines()
for data_item in tqdm(raw_data):
data_item = json.loads(data_item)
sample = ''
for conversation in data_item['conversations']:
sample += conversation['value']
input_ids = tokenizer(
[sample],
return_tensors='np',
padding=False,
truncation=False,
).input_ids
anno_token_cnt += input_ids.shape[1]
token_dict[annotation] = anno_token_cnt
print(f'{annotation} total {anno_token_cnt}')
print(token_dict)
with open("utils/data_process/data_token_stastic.json", "w", encoding="utf-8") as f:
json.dump(token_dict, f, ensure_ascii=False, indent=4)
|