Spaces:
Running
Running
| """ | |
| TODO: | |
| 1. 繁体、简体、语种、 | |
| 2. 确认 bert的space token数目 | |
| 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, language_ranges | |
| CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) | |
| 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(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[str] = None, | |
| return_type: Optional[Literal["dict", "dataframe"]] = "dataframe" | |
| ) -> Union[pd.DataFrame, dict]: | |
| """ | |
| """ | |
| logger.info(f"columns: {columns}, tokenizer_filter: {tokenizer_filter}") | |
| stats = {} | |
| 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)}") | |