class Generator(object): ''' generator for testing ''' def __init__(self): self.tool = Tool(hparams.sens_num, hparams.key_len, hparams.sen_len, hparams.poem_len, 0.0) self.tool.load_dic(hparams.vocab_path, hparams.ivocab_path) vocab_size = self.tool.get_vocab_size() print("vocabulary size: %d" % (vocab_size)) PAD_ID = self.tool.get_PAD_ID() B_ID = self.tool.get_B_ID() assert vocab_size > 0 and PAD_ID >= 0 and B_ID >= 0 self.hps = hparams._replace(vocab_size=vocab_size, pad_idx=PAD_ID, bos_idx=B_ID) # load model model = MixPoetAUS(self.hps) # load trained model utils.restore_checkpoint(self.hps.model_dir, device, model) self.model = model.to(device) self.model.eval() #utils.print_parameter_list(self.model) # load poetry filter print("loading poetry filter...") self.filter = PoetryFilter(self.tool.get_vocab(), self.tool.get_ivocab(), self.hps.data_dir) print("--------------------------") def generate_one(self, keyword, length, factor_label1, factor_label2, beam_size=20, verbose=1, manu=False): ''' generate one poem according to the inputs: keyword: a topic word length: the length of each line, 5 or 7 factor_label: label for the two factors, when factor_label = -1, the model infers an appropriate class in terms of the keyword verbose: 0, 1, 2, 3 ''' assert length == 5 or length == 7 key_state = self.get_key_state(keyword) # infer labels when factor_label = -1 if factor_label1 == -1: factor_label1 = self.get_inferred_factor(key_state, 1) if verbose >= 1: print("inferred label1: %d" % (factor_label1[0].item())) else: factor_label1 = torch.tensor([factor_label1], dtype=torch.long, device=device) if factor_label2 == -1: factor_label2 = self.get_inferred_factor(key_state, 2) if verbose >= 1: print("inferred label2: %d" % (factor_label2[0].item())) else: factor_label2 = torch.tensor([factor_label2], dtype=torch.long, device=device) # get decoder initial state dec_init_state = self.get_dec_init_state(key_state, factor_label1, factor_label2, length) context = torch.zeros((1, self.hps.context_size), dtype=torch.float, device=device) # (B, context_size) # initialize beam pool beam_pool = PoetryBeam(beam_size, length, self.tool.get_B_ID(), self.tool.get_E_ID(), self.tool.get_UNK_ID(), self.filter.get_level_cids(), self.filter.get_oblique_cids()) # beam search poem = [] self.filter.reset(length, verbose) for step in range(0, self.hps.sens_num): # generate each line if verbose >= 1: print("\ngenerating step: %d" % (step)) # get the rhythm pattern and rhyme id of te current line _, rhythms, rhyme = self.filter.get_pattern(step) # pos_tensor = self.tool.pos2tensor(step) pos_tensor = pos_tensor.to(device) init_state = torch.cat( [dec_init_state.clone().detach(), pos_tensor], dim=-1) # reset beam pool beam_pool.reset(init_state, rhythms, rhyme, self.filter.get_rhyme_cids(rhyme), self.filter.get_repetitive_ids()) candidates, costs, states = self.beam_search( beam_pool, length, context) lines = [self.tool.idxes2line(idxes) for idxes in candidates] lines, costs, states = self.filter.filter_illformed( lines, costs, states, step) if len(lines) == 0: return "", "generation failed!" which = 0 if manu: for i, (line, cost) in enumerate(zip(lines, costs)): print("%d, %s, %.2f" % (i, line, cost)) which = int(input("select sentence>")) line = lines[which] poem.append(line) # the first line determin the rhythm pattern of the poem if step == 0: self.filter.set_pattern(line) # when the first line doesn't rhyme, we can determin the rhyme # in terms of the second line if step == 1 and self.filter.get_rhyme() == -1: self.filter.set_rhyme(line) # set repetitive chars self.filter.add_repetitive(self.tool.line2idxes(line)) # update the context vector context = self.update_context(context, states[which], beam_size, length) return poem, "ok" # ------------------------------------ def beam_search(self, beam_pool, trg_len, ori_context): # current size of beam candidates in the beam pool n_samples = beam_pool.uncompleted_num() # (1, context_size) -> (B, context_size) context = ori_context.repeat(n_samples, 1) for k in range(0, trg_len + 10): #print ("beam search position %d" % (k)) inps, states = beam_pool.get_beam_tails() logits, new_states = self.do_dec_step(inps, states, context[:n_samples, :]) beam_pool.advance(logits, new_states, k) n_samples = beam_pool.uncompleted_num() if n_samples == 0: break candidates, costs, dec_states = beam_pool.get_search_results() return candidates, costs, dec_states # --------------------------- def get_key_state(self, keyword): #print (keyword) key_tensor = self.tool.keys2tensor([keyword]) #print (key_tensor) return self.model.compute_key_state(key_tensor.to(device)) def get_inferred_factor(self, key_state, factor_id): assert factor_id == 1 or factor_id == 2 return self.model.compute_inferred_label(key_state, factor_id) def get_dec_init_state(self, key_state, label1, label2, length): state = self.model.compute_dec_init_state(key_state, label1, label2) # length tensor len_tensor = self.tool.lengths2tensor([length]).to(device) dec_init_state = torch.cat([state, len_tensor], dim=-1) # (1, H) return dec_init_state def do_dec_step(self, inps, states, context): logits, new_state = self.model.dec_step(inps, states, context) return logits, new_state def update_context(self, old_context, ori_states, beam_size, length): # update the context vector # old_context: (1, context_size) # states: (1, H) * L H = ori_states[0].size(1) states = ori_states[1:length + 2] + [ torch.zeros_like(ori_states[0], device=device) ] * (7 - length) states = [state.view(H, 1) for state in states] states = torch.cat(states, dim=1).unsqueeze(0) # (1, H, L) context = self.model.layers['context'](old_context, states) return context