utils/zh_util.py

"""
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))