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)