File size: 7,384 Bytes
7c73423
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
"""

TODO:

1. add more language

2. check space count of bert

3. add token_impl

4.

"""
import os
import json
import numpy as np
import pandas as pd
from collections import Counter, defaultdict
from vocab import tokenizer_factory
from typing import Optional, Union, Literal
from utils.log_util import logger
from utils.text_util import contains_digit, get_space_count
from utils.lang_util import detect_language_by_unicode, language_ranges

CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))

default_columns = ["digit", "zh"]

def _to_unicode(text):
    return ''.join(r'\u{:04X}'.format(ord(chr)) for chr in text)


def _get_coding_length(tokenizer, vocab, filter=None):
    """

    oov character may be tokenized into more than one 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


cache = {}


def _dist(token_lens):
    """

    :param token_lens:

    :return: min,median,max of token_lens

    """
    if not token_lens:
        return "-"
    return f"{min(token_lens)},{round(np.median(token_lens))},{max(token_lens)}"


def iter_vocab(

        tokenizer_name: str,

        from_cache: bool = True,

        cache_dir: str = "stats",

) -> Union[pd.DataFrame, dict]:
    """

    :param tokenizer_name:

    :param from_cache:

    :param cache_dir:

    :return:

    """
    tokenizer_config = tokenizer_factory.get_tokenizer_config(tokenizer_name)

    cache_dir = os.path.join(CURRENT_DIR, cache_dir)
    os.makedirs(cache_dir, exist_ok=True)

    # load from cache
    cache_path = os.path.join(cache_dir, "character_stats.json")
    if not cache and os.path.exists(cache_path):
        with open(cache_path, "r", encoding="utf-8") as f_tmp:
            cache.update(json.load(f_tmp))
    if from_cache and tokenizer_name in cache:
        # logger.info(f"load {tokenizer_config.name_or_path} from cache")
        return cache[tokenizer_name]

    tokenizer = tokenizer_factory.get_tokenizer(tokenizer_name)

    tokens_by_lang = {lang[1]: [] for lang in language_ranges.keys()}
    digit_tokens = []
    space_tokens = []
    byte_tokens = []

    buffer = []
    for token_id in range(tokenizer.vocab_size):
        # for token_id in tokenizer.get_vocab():
        # for token_id in range(len(tokenizer)):
        decode_str = tokenizer.decode([token_id], skip_special_tokens=False)
        token = tokenizer.convert_ids_to_tokens([token_id], skip_special_tokens=False)[0]
        tags = []
        if token is None:  # 有些词典有空的id(不连续)
            continue
        if isinstance(token, bytes):
            token = token.decode("utf-8", errors="ignore")

        if hasattr(tokenizer, "sp_model"):  # 基于 sentencepiece 包
            if tokenizer.sp_model.is_byte(token_id):
                tags.append("is_byte")
                byte_tokens.append(token)

        language_tags = detect_language_by_unicode(decode_str)
        for language in language_tags:
            tokens_by_lang[language[1]].append(decode_str)

        if contains_digit(decode_str):
            tags.append("digit")
            digit_tokens.append(decode_str)

        space_count = get_space_count(decode_str)
        if space_count > 0:
            space_tokens.append(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),
            },
            ensure_ascii=False) + "\n")

    result = {
        "tokenizer": tokenizer_factory.get_name_with_hyperlink(tokenizer_name),
        "organization": tokenizer_config.org,
        # "impl": str(tokenizer.__class__),
        # "vocab_size-": tokenizer.vocab_size,  # vocab_size_without_added_token
        "vocab_size": len(tokenizer),

        # "中文汉字编码长度均值": mean_length,   # 不用统计,因为字典包含中文字符多,一般就意味着 中文汉字编码长度短。
        # "中文汉字编码长度分布": json.dumps(dist_length),

        "num(digit)": len(digit_tokens),
        "len(digit)": _dist([len(token) for token in digit_tokens]),
        "num(space)": len(space_tokens),
        "len(space)": _dist([len(token) for token in space_tokens]),

        # "num(byte)": len(byte_tokens)
    }

    for lang, tokens in tokens_by_lang.items():
        result[f"num({lang})"] = len(tokens)
        result["len(" + lang + ")"] = _dist([len(token) for token in tokens])

    out_path = os.path.join(cache_dir, f"iter_vocab/{tokenizer_name.replace('/', '_')}.vocab.jsonl")
    with open(out_path, "w", encoding="utf-8") as f_out:
        for line in buffer:
            f_out.write(line)
    len_before = len(cache)
    cache[tokenizer_name] = result
    len_after = len(cache)
    logger.info(f"saving {tokenizer_name} to memory and file cache: {len_before}->{len_after}")
    with open(cache_path, "w", encoding="utf-8") as f_out:
        f_out.write(json.dumps(cache, ensure_ascii=False, indent=2))
    return result


def to_dataframe(stats, columns):
    table = []
    for stat in stats.values():
        filtered_stat = {}
        for k, v in stat.items():
            if not k.startswith("num") and not k.startswith("len"):
                filtered_stat[k] = v
            if any(column in k for column in columns):
                k = k.replace("ja-kana", "kana")
                filtered_stat[k] = v
        table.append(filtered_stat)
    df = pd.DataFrame(table)
    return df


def get_character_table(

        tokenizer_filter: Optional[str] = None,

        columns: Optional[list] = None,

        return_type: Optional[Literal["dict", "dataframe"]] = "dataframe"

) -> Union[pd.DataFrame, dict]:
    """

    """
    logger.info(f"columns: {columns}, tokenizer_filter: {tokenizer_filter}")
    stats = {}
    if columns is None:
        columns = default_columns
    if tokenizer_filter is not None:
        tokenizer_names = [tokenizer_config.name_or_path for tokenizer_config in tokenizer_factory.all_tokenizer_configs
                           if tokenizer_filter.lower() in tokenizer_config.name_or_path.lower()]
    else:
        tokenizer_names = tokenizer_factory.all_tokenizer_names

    for tokenizer_name in tokenizer_names:
        stat = iter_vocab(tokenizer_name)
        stats[tokenizer_name] = stat

    if return_type == "dataframe":
        stats = to_dataframe(stats, columns)
    return stats


if __name__ == "__main__":
    # aa = get_character_table(tokenizer_filter="baichuan")
    df = get_character_table()
    logger.info(f"\n{df.to_markdown(index=False)}")