语言模型建模是针对句子建模，主要分为掩码语言模型和自回归语言模型。
我们从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 sysdef 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] #取出句长对应的setfsrc.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 nnclass 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_sizesuper(NNLM, self).__init__()self.emb = nn.Embedding(vcb_size, isize,padding_idx=0)                #<pad>的索引为0self.comp = nn.Linear(ctx_len * isize, isize, bias=False) #将4个词的向量降维为isizeself.drop = nn.Dropout(p=dropout, inplace=True) #embedding后dropoutself.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) + 1return vcb# freq word0 word1 word2 将同一频率的词放在一起
def save_vcb(vcb, file):# {freq: [word]}rs = {}for word, freq in vcb.items():if freq in rs: #若此词频项已存在，则添加该词到listrs[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] = cidcid += 1#为每个分词建立索引return vcbreverse_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 + 1if (_n * _mlen) > max_tokens: #如果把添加了这句话的 句数*分词数量 大于最大值则不能放if ri:       #如果ri不为空yield ri, mlen  #返回ri和原来的句子长度ri, mlen, n = [ls], _l, 1 #返回后重新初始化，将本句加入新的batchelse: #如果不超过当前长度，则将此句添加到batch中ri.append(ls)mlen, n = _mlen, _n #更新句子长度与句子数量if ri: #最后若仍然有数据，则返回为一个新的batchyield ri, mlendef 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 rsdef 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>的索引0rs.append(lu)return rs #返回的是均已对齐的每个batchdef 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 ) #设置压缩存储节省空间，压缩等级设置为最大压缩代价9h5f["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-8import 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 tqdmtrain_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 modindef train(train_data, tl, model, lossf, optm, cuda_device,nreport=nreport, tokens_optm=tokens_optm, ctx_len=ctx_len):#nreport为每训练一部分词打一次epochmodel.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 + 1for _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个词才能运行，否则跳过continueif 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 += _nloss.backward()                 #反向传播求导if _tom > tokens_optm:          #当词数大于时更新参数optm.step()                     #参数的更新optm.zero_grad(set_to_none=True)#参数更新后清空梯度_tom = 0if _i % nreport == 0:   #每训练5000个batch打印一次print("Average loss over %d tokens: %.2f"%(_tb, _lb/_tb))_lb = 0.0_tb = 0save_model(model, "checkpoint.pt") #暂存检查点模型return _l / _t #返回总的lossdef save_model(modin, fname): #保存模型所有内容 权重、偏移、优化torch.save({name: para.cpu() for name, para inmodel.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且设备支持cudacuda_device = torch.device("cuda", gpu_id)  #配置gputorch.set_default_device(cuda_device)
else:cuda_device = Noneif cuda_device is not None:                     #如果要用gpumodel.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())] #获得字符串构成的训练数据的listfor 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:_ = lurs.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) 测试时句末无需加eosreturn 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 = inputbsize =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 = _eoselse:done_trans |= _eos #将_eos中的True赋给done_trans#_eos中的元素如果为True，则说明在该索引位置上out值为2即结束标志if done_trans.all().item(): #当全都为True，说明此batch中所有句子预测都为<eos>，即解码完成breakreturn 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-8import 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:breakelse: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,2model = 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)       #判断是否使用cudamodel.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转为listoutput = [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

可以看到模型对测试句子做了续写。