""" TODO: 繁体、简体、语种、 """ import os import json from collections import Counter from utils.log_util import logger from utils.text_util import is_zh_char, is_all_zh, has_zh, is_all_digit, has_digit, get_zh_count, get_digit_count, get_space_count CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) zh_tokens = [line.strip() for line in open(os.path.join(CURRENT_DIR, "vocab.jd.txt.v2"), "r", encoding="utf-8") if is_zh_char(line.strip())] def to_unicode(text): return ''.join(r'\u{:04X}'.format(ord(chr)) for chr in text) def zh_iterator(): for idx in range(ord(u'\u4e00'), ord(u'\u9fa5')): yield (chr(idx)) def get_coding_length(tokenizer, vocab, filter=None): """ 计算编码长度。（有些中文汉字被解码成多个token） """ all_length = [] for word in vocab: if len(word) > 1: continue if filter is not None and filter(word): continue try: tokens = tokenizer.encode(word) except Exception as e: print(e) all_length.append(len(tokens)) # if len(tokens.ids) > 1: # if len(tokens) > 3: # print(word, tokens) dist_length = Counter(all_length) mean_length = round(sum(all_length) / len(all_length), 2) return dist_length, mean_length def remove_special_char(): """ :return: """ # bert词典有 ##开头的 # byteBPE词典有带空格的 # decode_str = decode_str.strip().replace("#", "") # TODO, 按类型 pass cache = {} def iter_vocab(tokenizer, from_cache=True, cache_dir="stats/iter_vocab"): """ 由于速度较快，建议不采用文件缓存。 :param tokenizer: :param from_cache: :return: """ cache_dir = os.path.join(CURRENT_DIR, f"../{cache_dir}") os.makedirs(cache_dir, exist_ok=True) name = tokenizer.alias # L1 cache if from_cache and name in cache: logger.info(f"load {name} from cache") return cache[name] # L2 cache: not recommended # has_zh_token_stats = {"total_tokens": 0, "mean_token_length": 0} # all_zh_token_stats = {"total_tokens": 0, "mean_token_length": 0} # has_number_token_stats = {"total_tokens": 0, "mean_token_length": 0} # all_number_token_stats = {"total_tokens": 0, "mean_token_length": 0} has_zh_tokens = [] all_zh_tokens = [] has_digit_tokens = [] all_digit_tokens = [] has_space_tokens = [] all_space_tokens = [] # zh_tags = ["all_zh", "has_zh"] # digit_tags = ["all_digit", "has_digit"] # zh_token_count = {"total": 0, "包含1个中文单字": 0, "中文多字": 0} # symbol_count = 0 all_single_zh_tokens = set() zh_symbol_count = 0 buffer = [] for token_id in range(tokenizer.vocab_size): decode_str = tokenizer.decode([token_id], skip_special_tokens=False) token = tokenizer.convert_ids_to_tokens([token_id], skip_special_tokens=False)[0] # tokenizer.convert_tokens_to_string(tokens) tags = [] if token is None: # 有些词典有空的id（不连续） continue if isinstance(token, bytes): token = token.decode("utf-8", errors="ignore") digit_count = get_digit_count(decode_str) if is_all_zh(decode_str): tags.append("all_zh") all_zh_tokens.append(decode_str) elif has_zh(decode_str): tags.append("has_zh") has_zh_tokens.append(decode_str) if is_all_digit(decode_str): tags.append("all_digit") all_digit_tokens.append(decode_str) elif has_digit(decode_str): tags.append("has_digit") has_digit_tokens.append(decode_str) space_count = get_space_count(decode_str) zh_count = get_zh_count(decode_str) buffer.append(json.dumps( {"id": token_id, "token": token, "token_decode": decode_str, "token_dumps": json.dumps(token), "token_unicode": to_unicode(token), "token_len": len(decode_str), "zh_count": zh_count, # 包含汉字的数目 "tags": tags, "zh_symbol_count": zh_symbol_count, "": "", }, ensure_ascii=False) + "\n") # if zh_count >= 1: # zh_token_count["total"] += 1 # if zh_count > 1: # zh_token_count["中文多字"] += 1 # else: # zh_token_count["中文单字"] += 1 # all_single_zh_tokens.add(decode_str.strip().replace("#", "")) # # zh_token_count["中文单字-去重后"] = len(all_single_zh_tokens) dist_length, mean_length = get_coding_length(tokenizer, zh_tokens, filter=lambda k: not is_zh_char(k)) # TODO: 繁体字，简体字 result = { "name": name, "impl": str(tokenizer.__class__), "vocab_size": tokenizer.vocab_size, "中文token数": len(has_zh_tokens), "中文token的平均长度": None, "纯中文token的平均长度": None, "中文标点数": zh_symbol_count, "中文汉字编码长度均值": mean_length, "中文汉字编码长度分布": json.dumps(dist_length), "纯数字token数": digit_count, "纯数字token的平均长度": None, "纯中文token数": None, "纯space的token数": space_count, "纯space的token的平均长度": None, } out_path = os.path.join(cache_dir, f"{name}.vocab.jsonl") logger.info(f"saving vocab to {out_path}") with open(out_path, "w", encoding="utf-8") as f_out: f_out.write(json.dumps(result, ensure_ascii=False) + "\n") for line in buffer: f_out.write(line) cache[name] = result return result if __name__ == "__main__": # test_coding_length(jd_vocab_tokens, filter=lambda k: not is_chinese(k)) # test_coding_length(zh_punc) # test_coding_length(zh_iterator()) # from vocab.chatglm2_6b import tokenizer; name = "chatglm2_6b" # from vocab.chatglm_6b import tokenizer; name="chatglm_6b" # from vocab.baichuan2 import tokenizer; name="baichuan2" from vocab.gpt_4 import tokenizer; name="gpt4" # from vocab.gpt2 import tokenizer; name="gpt2" # from vocab.qwen1_5_14b_chat import tokenizer; name="qwen1_5_14b_chat" # from vocab.gpt_nexo_20b import tokenizer; name="gpt_nexo_20b" # from vocab.fastchat_t5_3b import tokenizer; name="fastchat_t5_3b" print(iter_vocab(tokenizer, name=name))