def test(opt):
# 数据
dataloader = get_dataloader(opt)
_data = dataloader.dataset._data
word2ix, ix2word = _data['word2ix'], _data['ix2word']
sos = word2ix.get(_data.get('sos'))
eos = word2ix.get(_data.get('eos'))
unknown = word2ix.get(_data.get('unknown'))
voc_length = len(word2ix)
#定义模型
encoder = EncoderRNN(opt, voc_length)
decoder = LuongAttnDecoderRNN(opt, voc_length)
#加载模型
if opt.model_ckpt == None:
raise ValueError('model_ckpt is None.')
return False
checkpoint = torch.load(opt.model_ckpt, map_location=lambda s, l: s)
encoder.load_state_dict(checkpoint['en'])
decoder.load_state_dict(checkpoint['de'])
with torch.no_grad():
#切换模式
encoder = encoder.to(opt.device)
decoder = decoder.to(opt.device)
encoder.eval()
decoder.eval()
#定义seracher
searcher = GreedySearchDecoder(encoder, decoder)
return searcher, sos, eos, unknown, word2ix, ix2word
def eval(**kwargs):
opt = Config()
for k, v in kwargs.items(): #设置参数
setattr(opt, k, v)
# 数据
dataloader = get_dataloader(opt)
_data = dataloader.dataset._data
word2ix, ix2word = _data['word2ix'], _data['ix2word']
sos = word2ix.get(_data.get('sos'))
eos = word2ix.get(_data.get('eos'))
unknown = word2ix.get(_data.get('unknown'))
voc_length = len(word2ix)
#定义模型
encoder = EncoderRNN(opt, voc_length)
decoder = LuongAttnDecoderRNN(opt, voc_length)
#加载模型
if opt.model_ckpt == None:
raise ValueError('model_ckpt is None.')
return False
checkpoint = torch.load(opt.model_ckpt, map_location=lambda s, l: s)
encoder.load_state_dict(checkpoint['en'])
decoder.load_state_dict(checkpoint['de'])
with torch.no_grad():
#切换模式
encoder = encoder.to(opt.device)
decoder = decoder.to(opt.device)
encoder.eval()
decoder.eval()
#定义seracher
searcher = GreedySearchDecoder(encoder, decoder)
while (1):
input_sentence = input('> ')
if input_sentence == 'q' or input_sentence == 'quit': break
cop = re.compile("[^\u4e00-\u9fa5^a-z^A-Z^0-9]") #分词处理正则
input_seq = jieba.lcut(cop.sub("", input_sentence)) #分词序列
input_seq = input_seq[:opt.max_input_length] + ['</EOS>']
input_seq = [word2ix.get(word, unknown) for word in input_seq]
tokens = generate(input_seq, searcher, sos, eos, opt)
output_words = ''.join([ix2word[token.item()] for token in tokens])
print('BOT: ', output_words)
def eval():
parameter = Config()
# 加载参数
save_dir = parameter.save_dir
loadFilename = parameter.model_ckpt
pretrained_embedding_path = parameter.pretrained_embedding_path
dropout = parameter.dropout
hidden_size = parameter.hidden_size
num_layers = parameter.num_layers
attn_model = parameter.method
max_input_length = parameter.max_input_length
max_generate_length = parameter.max_generate_length
embedding_dim = parameter.embedding_dim
#加载embedding
voc = read_voc_file('./data/voc.pkl')
embedding = get_weight(voc,pretrained_embedding_path)
#输入
inputs = get_input_line('./test/test.txt')
input_batches, lengths = get_batch_id(inputs)
#
encoder = EncoderRNN(hidden_size, embedding, num_layers, dropout)
decoder = LuongAttnDecoderRNN(attn_model,embedding,hidden_size,len(voc),num_layers,dropout)
if loadFilename == None:
raise ValueError('model_ckpt is None.')
return False
checkpoint = torch.load(loadFilename, map_location=lambda s, l: s)
print(checkpoint['plt'])
encoder.load_state_dict(checkpoint['en'])
decoder.load_state_dict(checkpoint['de'])
answer =[]
with torch.no_grad():
encoder.to(device)
decoder.to(device)
#切换到测试模式
encoder.eval()
decoder.eval()
search = GreedySearchDecoder(encoder, decoder)
for input_batch in input_batches:
#print(input_batch)
token,score = generate(input_batch, search, GO_ID, EOS_ID, device)
print(token)
answer.append(token)
print(answer)
return answer
def model_fn(model_dir):
logger.info('Loading the model.')
model_info = {}
with open(os.path.join(model_dir, 'model_info.pth'), 'rb') as f:
# If loading a model trained on GPU to CPU
if torch.cuda.device_count() < 1:
checkpoint = torch.load(f, map_location=torch.device('cpu'))
else:
checkpoint = torch.load(f)
#have to save these hyper parameters
hidden_size = model_info['hidden_size']
encoder_n_layers = model_info['encoder_n_layers']
decoder_n_layers = model_info['decoder_n_layers']
dropout = model_info['dropout']
attn_model = model_info['attn_model']
voc = model_info['voc']
# Initialize word embeddings
embedding = nn.Embedding(voc.num_words, hidden_size)
embedding.load_state_dict(checkpoint['embedding'])
# Initialize encoder & decoder models
encoder = EncoderRNN(hidden_size, embedding, encoder_n_layers, dropout)
decoder = LuongAttnDecoderRNN(attn_model, embedding, hidden_size,
voc.num_words, decoder_n_layers, dropout)
encoder.load_state_dict(checkpoint['en'])
decoder.load_state_dict(checkpoint['de'])
# Set dropout layers to eval mode
encoder.eval()
decoder.eval()
searcher = GreedySearchDecoder(encoder, decoder, device)
return {'searcher': searcher, 'voc': voc}
encoder.train()
decoder.train()
# 初始化优化器
print('Building optimizers ...')
encoder_optimizer = optim.Adam(encoder.parameters(), lr=learning_rate)
decoder_optimizer = optim.Adam(decoder.parameters(),
lr=learning_rate * decoder_learning_ratio)
if loadFilename:
encoder_optimizer.load_state_dict(encoder_optimizer_sd)
decoder_optimizer.load_state_dict(decoder_optimizer_sd)
# 运行训练迭代
print("Starting Training!")
trainIters(model_name, voc, pairs, encoder, decoder, encoder_optimizer,
decoder_optimizer, embedding, encoder_n_layers, decoder_n_layers,
save_dir, n_iteration, batch_size, print_every, save_every, clip,
corpus_name, loadFilename)
#eval model
encoder.eval()
decoder.eval()
# 初始化探索模块
searcher = GreedySearchDecoder(encoder, decoder)
# 开始聊天（取消注释并运行以下行开始）
evaluateInput(encoder, decoder, searcher, voc)
print("*************")
def main():
USE_CUDA = torch.cuda.is_available()
device = torch.device("cuda" if USE_CUDA else "cpu")
# load dict
corpus_name = "cornell movie-dialogs corpus"
corpus = os.path.join("data", corpus_name)
datafile = os.path.join(corpus, "formatted_movie_lines.txt")
voc, pairs = loadPrepareData(corpus_name, datafile)
# model parameters
save_dir = os.path.join("data", "save")
model_name = 'cb_model'
attn_model = 'dot'
encoder_n_layers = 2
decoder_n_layers = 2
hidden_size = 500
checkpoint_iter = 4000
loadFilename = os.path.join(
save_dir, model_name, corpus_name,
'{}-{}_{}'.format(encoder_n_layers, decoder_n_layers, hidden_size),
'{}_checkpoint.tar'.format(checkpoint_iter))
# Load model if a loadFilename is provided
if loadFilename:
# If loading on same machine the model was trained on
checkpoint = torch.load(loadFilename)
# If loading a model trained on GPU to CPU
# checkpoint = torch.load(loadFilename, map_location=torch.device('cpu'))
encoder_sd = checkpoint['en']
decoder_sd = checkpoint['de']
encoder_optimizer_sd = checkpoint['en_opt']
decoder_optimizer_sd = checkpoint['de_opt']
embedding_sd = checkpoint['embedding']
voc.__dict__ = checkpoint['voc_dict']
print('Building encoder and decoder ...')
# Initialize word embeddings
embedding = nn.Embedding(voc.num_words, hidden_size)
if loadFilename:
embedding.load_state_dict(embedding_sd)
# Initialize encoder & decoder models
encoder = EncoderRNN(hidden_size, embedding, encoder_n_layers, dropout=0)
decoder = LuongAttnDecoderRNN(attn_model,
embedding,
hidden_size,
voc.num_words,
decoder_n_layers,
dropout=0)
if loadFilename:
encoder.load_state_dict(encoder_sd)
decoder.load_state_dict(decoder_sd)
# Use appropriate device
encoder = encoder.to(device)
decoder = decoder.to(device)
print('Models built and ready to go!')
# Set dropout layers to eval mode
encoder.eval()
decoder.eval()
# Initialize search module
searcher = GreedySearchDecoder(encoder, decoder, device)
# Begin chatting (uncomment and run the following line to begin)
evaluateInput(device, encoder, decoder, searcher, voc)