wiki百科之将词转换为索引表示

import os
import json

# 现在已经有准备好的BERT维基百科训练语料, 已经分割为train_wiki.txt和test_wiki.txt
# 语料来源: https://github.com/brightmart/nlp_chinese_corpus
# 在准备好的文件里, 写成了下面的格式, 每一行是一条string, 可以eval为python dict
# 分别对应着两句有着上下文关系的句子,
# 示例:
# "{'text1': '眼蛱蝶族（学名：Junoniini）是蛱蝶科蛱蝶亚科中的一个族。', 
#   'text2': '此分类的物种在始新世末至渐新世初开始形成。'}"
# 在这个项目里, BERT的训练中, 由./BERT/dataset/wiki_dataset.py文件中的脚本读取txt文件,
# 并动态随机做Masked LM和next sentence的mini batch

# 在这里主要是演示怎样制作用来训练的字典, 用来做tokenize, 也就是把汉字转换为token

# 加载所有的语料
# 注意这里可能会很慢, 可能需要等到5分钟
with open("train_wiki.txt", "r", encoding="utf-8") as f:
    all_wiki_corpus = [i for i in f.readlines()]
with open("test_wiki.txt", "r", encoding="utf-8") as f:
    all_wiki_corpus += [i for i in f.readlines()]
print(len(all_wiki_corpus))

# 因为这里上下句有重复的, 所以需要去重, 之后制作字典
# 注意这里可能会很慢, 可能需要等到5分钟
all_text = []
for dic in all_wiki_corpus:
    dic = eval(dic)
    all_text += [v for _, v in dic.items()]
all_text = list(set(all_text))
print(len(all_text))

all_text[333]

'在音乐方面，它更常指作品的类型和风格的更替。'

# 我们要制作字典, 首先要制作一个记录所有字出现频率的dict, 然后可以舍去出现频率非常低的字, 也可以不舍去
def get_word2tf(corpus_list):
    # word2tf是记录字频的dict
    word2tf = {}
    for text in corpus_list:
        for char in list(text):
            char = char.lower()
            word2tf = update_dic(char, word2tf)
    return word2tf

def update_dic(char, word2tf):
    if word2tf.get(char) is None:
        word2tf[char] = 1
    else:
        word2tf[char] += 1
    return word2tf

# 注意这里可能会很慢, 可能需要等到5-10分钟
word2tf = get_word2tf(all_text)

print(len(word2tf))
# 这里可以根据需要舍去字频较低的字, 我们这里不舍去任何东西, 因为只有19211个字...

19211

# 我们要训练BERT, 所以我们会有一些特殊的token, 例如#CLS#, #PAD#(用来补足长度)等等,
# 所以我们留出前20个token做备用, 实际字的token从序号20开始

# word2idx是我们将要制作的字典
word2idx = {}
# 定义一些特殊token
pad_index = 0 # 用来补长度和空白
unk_index = 1 # 用来表达未知的字, 如果字典里查不到
cls_index = 2 #CLS#
sep_index = 3 #SEP#
mask_index = 4 # 用来做Masked LM所做的遮罩
num_index = 5 # (可选) 用来替换语句里的所有数字, 例如把 "23.9" 直接替换成 #num#
word2idx["#PAD#"] = pad_index
word2idx["#UNK#"] = unk_index
word2idx["#SEP#"] = sep_index
word2idx["#CLS#"] = cls_index
word2idx["#MASK#"] = mask_index
word2idx["#NUM#"] = num_index

idx = 20
for char, v in word2tf.items():
    word2idx[char] = idx
    idx += 1
print(len(word2idx))

19217

# 注意!! 我们在训练BERT的时候, 实际需要初始化的字向量矩阵的维度是 [19211+20, embedding_dim]
# 不要忘记我们预留的20个特殊token的空间

# 写入json
with open('bert_word2idx.json', 'w+', encoding='utf-8') as f:
    f.write(json.dumps(word2idx, ensure_ascii=False))

# 至此字典制作完毕

{"#PAD#": 0, "#UNK#": 1, "#SEP#": 3, "#CLS#": 2, "#MASK#": 4, "#NUM#": 5, "n": 20, "中": 21, "山": 22, "路": 23, "位": 24, "于": 25, "松": 26, "江": 27, "老": 28, "城": 29, "区": 30, "心": 31, "地": 32, "带": 33, "，": 34, "主": 35, "要": 36, "有": 37,