语言模型建模是针对句子建模,主要分为掩码语言模型和自回归语言模型。
我们从corpus中截取一句话作为例子,这句话是bpe分词后的句子:
1994 年 5 月 17 日 安全 理事会 第 33@@ 77 次 会议 通过
掩码语言模型的主要机制是:例如将33@@ 77这两个分词做掩码,这样句子变成了
1994 年 5 月 17 日 安全 理事会 第 [MASK] [MASK] 会议 通过
将带有掩码的句子输入神经网络,让模型预测掩码位置的正确分词,例如上文正确的分词为 33@@ 、77.在分词的词表中,训练的目的是为了使得预测这句话的mask时,33@@、77分词的概率最高,从而选取这两个分词作为正确的预测结果。
自回归语言模型的机制是:给出这句话的前i-1个词,来预测第i个词;随后加入第i个词,预测这句话的第i+1个词。例如:
1994 -> 年
1994 年 -> 5
1994 年 5 -> 月
…
1994 年 5 月 17 日 安全 理事会 第 33@@ 77 次 会议 -> 通过
在自回归语言模型主要有两个问题:
1.句首的词怎么来,比如本句话的1994.
2.句末的分词“通过”生成之后,是继续生成还是结束?
因此,对句子的处理我们要加两个特殊标记:句首加[SOS](start of sentence),句末加[EOS](end of sentence).例如:
[SOS] 1994 年 5 月 17 日 安全 理事会 第 33@@ 77 次 会议 -> 通过 [EOS]
[SOS]作触发器,当遇到时开始了第一个词的生成以及整个句子的生成;[EOS]是结束的作用,当分类产生的概率[EOS]最大的时候,就认为不需要继续生成了。
模型的搭建
首先我们对分词后的语料按句长升序排序,设置最长句长为254,这样加上 eos 和 sos 后最大句长为256。
#sorti.py
#encoding: utf-8
import sys
def handle(srcf, srts, max_len=256): #设置句长最大256,大于则丢弃
# {length: {src}} 外层dict,内层set
_maxlen = max_len - 2 #减去[SOS]和[EOS]
data = {}
with open(srcf, "rb") as fsrc:
for ls in fsrc:
ls = ls.strip()
if ls:
ls = ls.decode("utf-8")
_ = len(ls.split()) #获取句子的分词个数
if _ <= max_len:
if _ in data: #若已有这个长度在data中
if ls not in data[_]: #去重,重复的跳过
data[_].add(ls) #不重复的添加
else:
data[_] = set([ls]) #转化成set去重
ens = "\n".encode("utf-8")
with open(srts, "wb") as fsrc: #写入
for _l in sorted(data.keys()): #按照句子长度从小到大排
lset = data[_l] #取出句长对应的set
fsrc.write("\n".join(lset).encode("utf-8")) #在每个句子间插入换行符
fsrc.write(ens) #每个句子后插入换行
if __name__ == "__main__":
handle(*sys.argv[1:3])
在命令行执行并查看结果文件:
:~/nlp/lm$ python sorti.py ../token/zh.bpe zh.srt
:~/nlp/lm$ less zh.srt
由于自回归语言模型的机制,我们通过已有的分词预测句子中下一个分词,则不断将预测的分词加入句子来预测下一个分词,那么输入就是变长的。为了解决这个问题,我们使用卷积操作,设置每次固定的上下文长度ctx_len
,来预测下一个分词。
假设ctx_len=4
,如果我们有序列0 1 2 3 4 5 6 7 8,则有:
0 1 2 3 -> 4
1 2 3 4 -> 5
2 3 4 5 -> 6
3 4 5 6 -> 7
4 5 6 7 -> 8
因此,若每个batch中句子的长度为seql,则我们每次需要取seql - ctx_len
个长度的分词,共取ctx_len次。
>>> import torch
>>> a = torch.arange(9,dtype=torch.long)
>>> ctx_len = 4
>>> nlen = 9 - ctx_len
>>> _ = [a.narrow(0,i,nlen) for i in range(ctx_len)]
>>> _
[tensor([0, 1, 2, 3, 4]), tensor([1, 2, 3, 4, 5]), tensor([2, 3, 4, 5, 6]), tensor([3, 4, 5, 6, 7])]
>>> b = torch.stack(_,dim=-1)
>>> b
tensor([[0, 1, 2, 3],
[1, 2, 3, 4],
[2, 3, 4, 5],
[3, 4, 5, 6],
[4, 5, 6, 7]])
#CNNLM.py
#encoding: utf-8
import torch
from torch import nn
class NNLayer(nn.Module):
def __init__(self, isize, hsize, dropout,norm_residual=True,
**kwargs):
super(NNLayer, self,).__init__() ##调用父类的初始化函数
self.net = nn.Sequential(nn.Linear(isize, hsize),
nn.ReLU(inplace=True), #设置relu激活函数,inplace=True在原始张量上进行
nn.Dropout(p=dropout, inplace=False),#设置丢弃率防止过拟合,同时创建一个新的张量
nn.Linear(hsize, isize, bias=False), nn.Dropout(p=dropout, inplace=True))
self.normer = nn.LayerNorm(isize) #做归一化
self.norm_residual = norm_residual #设置变量存储做判断
def forward(self, input):
_ = self.normer(input) #稳定之后的结果
return (_ if self.norm_residual else input) + self.net(_)
#如果参数初始化做的好,就用LayerNorm后的值,否则用原始值
class NNLM(nn.Module):
def __init__(self, vcb_size, isize, hsize, dropout,
nlayer, ctx_len=4, bindemb=True, **kwargs): #有多少个词,就分多少类,类别数为vcb_size
super(NNLM, self).__init__()
self.emb = nn.Embedding(vcb_size, isize,
padding_idx=0) #<pad>的索引为0
self.comp = nn.Linear(ctx_len * isize, isize, bias=False) #将4个词的向量降维为isize
self.drop = nn.Dropout(p=dropout, inplace=True) #embedding后dropout
self.nets = nn.Sequential(*[NNLayer(isize, hsize, dropout)
for _ in range(nlayer)])
self.classifier = nn.Linear(isize, vcb_size)
if bindemb:
self.classifier.weight = self.emb.weight#将emb的权重赋给分类器
self.normer = nn.LayerNorm(isize)
self.out_normer = nn.LayerNorm(isize)
self.ctx_len = ctx_len #保存上下文长度
# input: (bsize, seql-1) 句数,句长-1 由于最后一个词是预测不作为输入
def forward(self, input):
out = self.emb(input)
# out: (bsize, seql-1, isize)
_nlen = input.size(1) - self.ctx_len + 1# 每次取_nlen长度的分词数
_ = []
for i in range(self.ctx_len):
_.append(out.narrow(1, i, _nlen)) #从seql这一维取,每次从i开始取,取_nlen长度
out = torch.cat(_,dim=-1)
#[(bsize, _nlen, isize)] -> (bsize, _nlen, ctx_len * isize)
out = self.normer(self.drop(out))
out = self.comp(out) #降维
#(bsize, _nlen, ctx_len * isize) -> (bsize, _nlen, isize)
out = self.normer(out) #使用归一化,使模长均匀
out = self.out_normer(self.nets(out))
return self.classifier(out) #分类产生参数
分词词典映射
#vcb.py
#encoding: utf-8
import sys
#预定义规则
init_vcb ={"<pad>": 0, "<sos>": 1, "<eos>": 2} #将占位符、开始标志和结束标志的ID映射为0、1、2
init_token_id = len(init_vcb) #设置其他词的ID从3开始编号
def get_vcb(*files): #词频统计
#{word: freq}
vcb = {}
for file in files:
with open(file, "rb") as f:
for line in f:
line = line.strip()
if line:
for word in line.decode("utf-8").split():
vcb[word] = vcb.get(word,0) + 1
return vcb
# freq word0 word1 word2 将同一频率的词放在一起
def save_vcb(vcb, file):
# {freq: [word]}
rs = {}
for word, freq in vcb.items():
if freq in rs: #若此词频项已存在,则添加该词到list
rs[freq].append(word)
else: #若此项不存在,则将该词作为list第一个词
rs[freq] = [word]
with open(file, "wb") as f:
for freq in sorted(rs.keys(), reverse=True):#从大到小读rs的频率,reverse=True说明按照从大到小排序
f.write(("%d\t%s\n"%(freq,"\t".join(rs[freq]))).#将同一词频的词合成一个字符串,词间用制表符分开
encode("utf-8"))
def load_vcb(file, vanilla = False, init_vcb = init_vcb, init_token_id = init_token_id):#为每个分词建立索引
#vanilla变量是为了区分是否使用了<pad>,若使用则编号从1开始否则从0开始
vcb, cid = ({}, 0) if vanilla else (init_vcb.copy(),
init_token_id)
with open(file, "rb") as f:
for line in f:
line = line.strip()
if line:
for word in line.decode("utf-8").split()[1:]: #舍弃词频,只需要保留分词
vcb[word] = cid
cid += 1
#为每个分词建立索引
return vcb
reverse_vcb = lambda vcb: {v: k for k, v in vcb.items()}
#使用匿名函数将 词->映射 转换为 映射->词
if __name__ == "__main__":
save_vcb(get_vcb(*sys.argv[1:-1]), sys.argv[-1])
#最后一个参数为存词频————分词的文件
我们对已经排好序的句子进行词频的收集,在命令行执行:
:~/nlp/lm$ python vcb.py zh.srt zh.vcb
:~/nlp/lm$ less zh.vcb
我们可以看到收集好的词典按照词频从高到低保存了所有分词。
文本转张量
#mkh5.py
#encoding: utf-8
import sys
from h5py import File as h5File
import numpy
from vcb import load_vcb
#导入加载词典的函数
def batch_loader(fsrc, max_tokens = 4096, **kwargs):#返回一批一batch的数据,设置每个batch最多存放2048个子词
ri = []
mlen = n = 0 #n记录当前收集了多少条句子,mlen记录当前收集的句子长度
with open(fsrc, "rb") as fs:
for ls in fs:
ls = ls.strip()
if ls:
ls =ls.decode("utf-8").split()
_l = len(ls) #当前行中的分词个数
_mlen = max(_l, mlen) #当前行或当前batch中句子的长度
_n = n + 1
if (_n * _mlen) > max_tokens: #如果把添加了这句话的 句数*分词数量 大于最大值则不能放
if ri: #如果ri不为空
yield ri, mlen #返回ri和原来的句子长度
ri, mlen, n = [ls], _l, 1 #返回后重新初始化,将本句加入新的batch
else: #如果不超过当前长度,则将此句添加到batch中
ri.append(ls)
mlen, n = _mlen, _n #更新句子长度与句子数量
if ri: #最后若仍然有数据,则返回为一个新的batch
yield ri, mlen
def map_instance(lin, vcb, sos_id=1, eos_id=2):
rs = [sos_id] #添加开始标志
rs.extend([vcb[_word] for _word in lin if _word in vcb])
rs.append(eos_id) #添加结束标志
return rs
def batch_mapper(fsrc, vcbs, **kwargs): #将分词变索引
for ri, mlen in batch_loader(fsrc, **kwargs):
yield [map_instance(_s, vcbs) for _s in ri], mlen + 2
#遍历每个batch中的句子,返回每个batch中每个分词的索引,以及batch长度
def pad_batch(lin, mlen, pad_id = 0):#补<pad>的函数
rs = []
for lu in lin: #每个batch中的每句
_d = mlen - len(lu) #当前此句需要补<pad>的个数
if _d > 0:
lu.extend([pad_id for _ in range(_d)])#extend函数用来拼接两个列表。补_d个<pad>的索引0
rs.append(lu)
return rs #返回的是均已对齐的每个batch
def batch_padder(fsrc, vcbs, **kwargs):
for ri, mlen in batch_mapper(fsrc, vcbs, **kwargs):
yield pad_batch(ri, mlen) #返回的是每个已补齐的batch,以及batch中句子的长度
def handle(fsrc, fvcbs, frs, **kwargs):
vcbs = load_vcb(fvcbs, vanilla = False)
with h5File(frs, "w", libver = 'latest', track_order = False) as h5f:#libver使用最新的,track_order表示无需记录顺序
src_grp = h5f.create_group("src", track_order=False) #创建两个组,分别放句子和标签
for i, ri in enumerate(batch_padder(fsrc, vcbs, **kwargs)):
ri = numpy.array(ri, dtype = numpy.int32) #转化成numpy数组并设置数据类型,target的数据很小,所以我们只需要int16存储
src_grp.create_dataset(str(i), data=ri, compression="gzip",
compression_opts=9, shuffle=True ) #设置压缩存储节省空间,压缩等级设置为最大压缩代价9
h5f["nword"] = numpy.array([len(vcbs)], dtype=numpy.int32) #存储总词数
h5f["ndata"] = numpy.array([i + 1], dtype=numpy.int32) #存储总batch数
if __name__ == "__main__":
handle(*sys.argv[1:4])
在命令行执行:
:~/nlp/lm$ python mkh5.py zh.srt zh.vcb train.h5
:~/nlp/lm$ h5ls -d train.h5
ndata Dataset {1}
Data:
83685
nword Dataset {1}
Data:
42585
src Group
可以看到我们总的训练集张量共有42585个分词以及83685个batch。
模型的训练
#train.py
#encoding: utf-8
import torch
from torch import nn
from CNNLM import NNLM #导入模型
from h5py import File as h5File #读训练数据
from math import sqrt
from random import shuffle #使输入数据乱序,使模型更均衡
from lrsch import SqrtDecayLR
from tqdm import tqdm
train_data = "train.h5"#之前已经张量转文本的h5文件
isize = 64
hsize = isize * 2 #设置初始参数
dropout = 0.3 #设置丢弃率
nlayer = 4 #设置层数
gpu_id = -1 #设置是否使用gpu
lr = 1e-3 #设置初始学习率
max_run = 8 #设置训练轮数
#early_stop = 16 #设置早停轮数
ctx_len = 4 #设置上下文长度
nreport = 5000 #每训练5000个batch打印一次
tokens_optm = 25000 #设置更新参数的词数阈值
def init_model_parameters(modin): #初始化模型参数
with torch.no_grad(): #不是训练不用求导
for para in modin.parameters():
if para.dim() > 1: #若维度大于1,说明是权重参数
_ = 1.0 / sqrt(para.size(-1))
para.uniform_(-_,_) #均匀分布初始化
for _m in modin.modules(): #遍历所有小模型
if isinstance(_m, nn.Linear):#如果小模型是linear类型
if _m.bias is not None: #初始化bias
_m.bias.zero_()
elif isinstance(_m, nn.LayerNorm):#初始化LayerNorm参数
_m.weight.fill_(1.0)
_m.bias.zero_()
return modin
def train(train_data, tl, model, lossf, optm, cuda_device,
nreport=nreport, tokens_optm=tokens_optm, ctx_len=ctx_len):#nreport为每训练一部分词打一次epoch
model.train() #设置模型在训练的模式
src_grp = train_data["src"] #从输入数据中取出句子
_l = 0.0 #_l用来存当前loss
_t = 0 #_t用来存句子数
_lb = 0.0
_tb = 0
_tom = 0
_min_len_thres = ctx_len + 1
for _i, _id in tqdm(enumerate(tl, 1)):
seq_batch = torch.from_numpy(src_grp[_id][()])
#seq_batch:[bsize, seql]
_seqlen = seq_batch.size(-1) #取出每个batch的句长
if _seqlen < _min_len_thres: #至少有ctx_len个词才能运行,否则跳过
continue
if cuda_device is not None:
seq_batch = seq_batch.to(cuda_device, non_blocking=True)
#将数据放在同一gpu上
seq_batch = seq_batch.long() #数据转换为long类型
seq_i = seq_batch.narrow(1, 0, _seqlen - 1) #训练数据读取bsize个seql-1的数据
#seq_i:[bsize, seql-1]
seq_o = seq_batch.narrow(1, ctx_len, _seqlen - ctx_len) #预测数据读取ctx_len~seql-1的预测数据
#seq_o:[bsize, seql-ctx_len]
#out: {bsize, seql-ctx_len, vcb_size} vcb_size即预测类别数
out = model(seq_i) #获得模型结果
loss = lossf(out.view(-1, out.size(-1)), seq_o.contiguous().view(-1))
#转换out维度为[bsize*(seql-ctx_len),vcb_size],seq_o:[bsize*(seql-ctx_len)]
_lossv = loss.item()
_l += _lossv #整个训练集的loss
_lb += _lossv #每个batch的loss
_n = seq_o.ne(0).int().sum().item() #seq_o中不是<pad>的位置的数量
_t += _n #整个训练集的分词数
_tb += _n #每个batch的分词数
_tom += _n
loss.backward() #反向传播求导
if _tom > tokens_optm: #当词数大于时更新参数
optm.step() #参数的更新
optm.zero_grad(set_to_none=True)#参数更新后清空梯度
_tom = 0
if _i % nreport == 0: #每训练5000个batch打印一次
print("Average loss over %d tokens: %.2f"%(_tb, _lb/_tb))
_lb = 0.0
_tb = 0
save_model(model, "checkpoint.pt") #暂存检查点模型
return _l / _t #返回总的loss
def save_model(modin, fname): #保存模型所有内容 权重、偏移、优化
torch.save({name: para.cpu() for name, para in
model.named_parameters()}, fname)
t_data = h5File(train_data, "r")#以读的方式打开训练数据
vcb_size = t_data["nword"][()].tolist()[0] #将返回的numpy的ndarray转为list
#在我们的h5文件中存储了nword: 总词数
model = NNLM(vcb_size, isize, hsize, dropout, nlayer, ctx_len=ctx_len)
model = init_model_parameters(model) #在cpu上初始化模型
lossf = nn.CrossEntropyLoss(reduction='sum', ignore_index=0,
label_smoothing=0.1)
#设置ignore_index=0,即忽略<pad>的影响
if (gpu_id >= 0) and torch.cuda.is_available(): #如果使用gpu且设备支持cuda
cuda_device = torch.device("cuda", gpu_id) #配置gpu
torch.set_default_device(cuda_device)
else:
cuda_device = None
if cuda_device is not None: #如果要用gpu
model.to(cuda_device) #将模型和损失函数放在gpu上
lossf.to(cuda_device)
optm = torch.optim.Adam(model.parameters(), lr=lr,
betas=(0.9, 0.98), eps=1e-08)
#使用model.parameters()返回模型所有参数,
lrm = SqrtDecayLR(optm, lr) #将优化器和初始学习率传入
tl = [str(_) for _ in range(t_data["ndata"][()].item())] #获得字符串构成的训练数据的list
for i in range(1, max_run + 1):
shuffle(tl) #使数据乱序
_tloss = train(t_data, tl, model, lossf, optm,
cuda_device) #获取每轮训练的损失
print("Epoch %d: train loss %.2f"%(i, _tloss)) #打印日志
save_model(model, "eva.pt")
lrm.step() #每轮训练后更新学习率
t_data.close()
在训练脚本中,我们要注意训练集很大,所以设置每nreport 个batch打印一次结果。更新参数时也设置一个阈值来多轮更新一次而不是每个batch更新一次。
由于整体训练模型开销大,因此我们每nreport个batch暂存一次检查点模型方便解码和后续工作的推进。
我们训练的过程就是将句子的前seql-1个词输入模型(seql为句子长度,即句子分词个数),将训练集句子第ctx_len以后的词作为标签(模型从前ctx_len个词预测第ctx_len+1的词,再添加再预测以此类推),直到下标为 seql-ctx_len-2
到 seql-2 预测下标为 seql-1 的词。至此句长为seql的0~seql-1的词都过了训练。计算预测分词和标签的交叉熵损失。
模型的预测
预测数据准备
我们选取两句话让模型预测之后的句子:
审议 了 1994 年 5 月 13 日 秘书长 的 报告 ( S / 1994 / 565 ) ,
重申 其 1993 年 9 月 29 日 关于 联合国 行动 安全 的 第 868 ( 1993 ) 号 决议 ,
我们做测试用的句子需满足句长(即分词个数)大于5,即大于ctx_len+1.
我们将其存入文件中,然后保存为h5文件:
:~/nlp/lm$ less test.bpe
审议 了 1994 年 5 月 13 日 秘书长 的 报告 ( S / 1994 / 565 ) ,
重申 其 1993 年 9 月 29 日 关于 联合国 行动 安全 的 第 868 ( 1993 ) 号 决议 ,
由于这两句是长度升序排列,所以省去排序的步骤。
:~/nlp/lm$ python mkh5.py test.bpe zh.vcb test.h5
:~/nlp/lm$ h5ls -d test.h5/src
0 Dataset {2, 23}
Data:
1, 159, 14, 1246, 12, 112, 20, 352, 35, 125, 3, 34, 11, 532, 13, 1246, 13, 12210, 10, 4, 2, 0,
0, 1, 489, 39, 1459, 12, 165, 20, 600, 35, 41, 33, 78, 79, 3, 16, 20380, 11, 1459, 10, 63, 67,
4, 2
我们可以看到第一句话的张量是:[ 1, 159, 14, 1246, 12, 112, 20, 352, 35, 125, 3, 34, 11, 532, 13, 1246, 13, 12210, 10, 4, 2, 0, 0,]
第二句话的张量是:[1, 489, 39, 1459, 12, 165, 20, 600, 35, 41, 33, 78, 79, 3, 16, 20380, 11, 1459, 10, 63, 67, 4, 2]
他们都满足:开头是<sos>
,索引为1;结尾是<eos>
,索引为2。第一句话由于要和第二句长度对齐构成一个batch,在末尾补了两个<pad>
,索引为0.
而我们预测续写句子时,在测试集我们首先不需要结尾标记2以及为补齐句子长度而在末尾添加的0.于是我们将pad占位符放在句子的开头,对mkh5.py脚本做以下修改:
#mkh5test.py
def pad_batch(lin, mlen, pad_id = 0):#补<pad>的函数
rs = []
for lu in lin: #每个batch中的每句
_d = mlen - len(lu) #当前此句需要补<pad>的个数
if _d > 0:
_ = [pad_id for _ in range(_d)]
#lu.extend([pad_id for _ in range(_d)])#extend函数用来拼接两个列表。补_d个<pad>
_.extend(lu) #将原句放在<pad>后
else:
_ = lu
rs.append(_)
return rs #返回的是均已对齐的每个batch
def map_instance(lin, vcb, sos_id=1, eos_id=2):
rs = [sos_id]
rs.extend([vcb[_word] for _word in lin if _word in vcb])
# rs.append(eos_id) 测试时句末无需加eos
return rs
再次保存h5文件后:
:~/nlp/lm$ python mktesth5.py test.bpe zh.vcb test.h5
:~/nlp/lm$ h5ls -d test.h5/src
0 Dataset {2, 22}
Data:
0, 0, 1, 159, 14, 1246, 12, 112, 20, 352, 35, 125, 3, 34, 11, 532, 13, 1246, 13, 12210, 10, 4,
1, 489, 39, 1459, 12, 165, 20, 600, 35, 41, 33, 78, 79, 3, 16, 20380, 11, 1459, 10, 63, 67, 4
第一句话张量:[ 0, 0, 1, 159, 14, 1246, 12, 112, 20, 352, 35, 125, 3, 34, 11, 532, 13, 1246, 13, 12210, 10, 4,]
第二句话张量:[ 1, 489, 39, 1459, 12, 165, 20, 600, 35, 41, 33, 78, 79, 3, 16, 20380, 11, 1459, 10, 63, 67, 4]
解码与预测
模型的解码
#CNNLM.py
# input: (bsize, seql)
def decode(self, input, maxlen=50): #设置最多生成50个分词
rs = input
bsize =input.size(0)
done_trans = None #记录是否完成生成
for i in range(maxlen):
out = rs.narrow(1, rs.size(1) - self.ctx_len, self.ctx_len) #利用后四个词预测下一个词
# (bsize, ctx_len, isize)
out = self.emb(out).view(bsize, -1)
# (bsize, ctx_len*isize)
out = self.normer(self.comp(self.drop(out)))
out = self.out_normer(self.nets(out))
out = self.classifier(out).argmax(-1,keepdim=True) #取最后一维分数最高的索引
#这一步对应分类,keepdim=True保持维度便于拼接
# out:(bsize, vcb_size) -> (bsize, 1)
rs = torch.cat([rs, out], dim=1) #将预测的词拼接到原句后,在第一维度即seql后
_eos = out.eq(2) #当遇到<eos>解码停止
# _eos:(bsize, 1)
if done_trans is None:
done_trans = _eos
else:
done_trans |= _eos #将_eos中的True赋给done_trans
#_eos中的元素如果为True,则说明在该索引位置上out值为2即结束标志
if done_trans.all().item(): #当全都为True,说明此batch中所有句子预测都为<eos>,即解码完成
break
return rs
首先,我们需要在模型文件的NNLM类中添加decode方法用于解码,decode方法将输入句子的最后ctx_len个分词作为模型的输入,预测下一个分词后将其拼接到句子中,下次再取后三个以及新生成的分词作为输入预测;
重复此步骤直至生成超过50个分词或者这一个batch中生成全部为<eos>
结束标志,则停止。
self.classifier(out).argmax(-1,keepdim=True)
这句话在执行argmax函数前,classifier分类器分类后的size是(bsize, vcb_size),所以执行argmax函数在最后一维取的下标索引即为分词在词表中的下标。
模型的预测
#predict.py
#encoding: utf-8
import sys
import torch
from CNNLM import NNLM #读模型
from h5py import File as h5File #读文件
from vcb import load_vcb, reverse_vcb #获取词表
isize = 64
hsize = isize * 2 #设置初始参数
dropout = 0.3 #设置丢弃率
nlayer = 4 #设置层数
gpu_id = -1 #设置是否使用gpu
ctx_len = 4 #设置上下文长度
def extract(lin, vcb): #提取结果的函数
rs = []
for lu in lin:
if lu > 1:
if lu == 2:
break
else:
rs.append(vcb[lu]) #返回索引对应词典中的分词
return " ".join(rs) #返回空格分隔的解码后的字符串
test_data = sys.argv[1]
test_file = h5File(test_data, "r") #读验证集
vcb_size = test_file["nword"][()].tolist()[0] #获取总词数
tgt_vcb = reverse_vcb(load_vcb(sys.argv[2], vanilla=False))
#vanilla设置为false,读取词表时需考虑到特殊标记:0,1,2
model = NNLM(vcb_size, isize, hsize, dropout, nlayer, ctx_len=ctx_len)
model_file = sys.argv[-1] #获取模型
with torch.no_grad(): #避免求导
_ = torch.load(model_file) #加载词典
for name, para in model.named_parameters():
if name in _:
para.copy_(_[name]) #从词典里取出name的参数
if (gpu_id >= 0) and torch.cuda.is_available():
cuda_device = torch.device("cuda", gpu_id)
torch.set_default_device(cuda_device)
else:
cuda_device = None
if cuda_device is not None:
model.to(cuda_device) #判断是否使用cuda
model.eval()
src_grp = test_file["src"]
ens = "\n".encode("utf-8")
with torch.no_grad(), open(sys.argv[3],"wb") as f: #解码避免求导,将预测标签按行写入文件
for _ in range(test_file["ndata"][()].item()):#每个batch上遍历
seq_batch = torch.from_numpy(src_grp[str(_)][()])
if cuda_device is not None:
seq_batch = seq_batch.to(cuda_device, non_blocking=True)
seq_batch = seq_batch.long() #s数据类型转换
output = model.decode(seq_batch).tolist() #将解码后的numpy转为list
output = [extract(_, tgt_vcb) for _ in output] #将张量转为文本
f.write("\n".join(output).encode("utf-8"))
f.write(ens) #每个batch间还应有换行
test_file.close()
在命令行执行:
:~/nlp/lm$ python predict.py test.h5 zh.vcb pred.txt checkpoint.pt
:~/nlp/lm$ less pred.txt
可以看到模型对测试句子做了续写。