"""

vocab size: 106029

中文汉字数：54230, 中文标点数: 549
"""

import json
from transformers import AutoTokenizer, BloomTokenizerFast


# tokenizer = AutoTokenizer.from_pretrained("tokenizer", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("moss-moon-003-sft", trust_remote_code=True)

print("vocab size:", tokenizer.vocab_size)

# tokens = tokenizer.encode("中国\nabcde<eoc>")
tokens = tokenizer.encode("中<eoc>")
decode_line = tokenizer.decode(tokens)
for token in tokens:
    print(token, tokenizer.decode([token]))


def id2token(ids):
    return tokenizer.convert_ids_to_tokens(ids)

def test_token():
    for word in "中国解决方法黑白侗，。！？；":
        encoding = tokenizer.encode(word)
        for token_id in encoding:
            decode_str = tokenizer.decode([token_id])  # 特殊字符解码后会统一变成 �，对应 "\ufffd"
            token = id2token([token_id])
            print(word, token_id, decode_str, json.dumps(decode_str), token, json.dumps(token))


test_token()