长文实践 | 详述文文本生成任务之营销文本生成

json_path = os.path.join(abs_path, '服饰_50k.json')
with open(json_path, 'r', encoding='utf8') as file:
    jsf = json.load(file)
for jsobj in jsf.values():
    title = jsobj['title'] + ' '  # Get title.
    kb = dict(jsobj['kb']).items()  # Get attributes.
    kb_merged = ''
    for key, val in kb:
        kb_merged += key+' '+val+' '  # Merge attributes.
    ocr = ' '.join(list(jieba.cut(jsobj['ocr'])))  # Get OCR text.
    texts = []
    texts.append(title + ocr + kb_merged)  # Merge them.
    reference = ' '.join(list(jieba.cut(jsobj['reference'])))
    for text in texts:
        sample = text+'<sep>'+reference  # Seperate source and reference.
        samples.add(sample)
write_path = os.path.join(abs_path, '../files/samples.txt')
write_samples(samples, write_path)

class Vocab(object):
    def __init__(self):
        self.word2index = {}
        self.word2count = Counter()
        self.reserved = ['<PAD>', '<SOS>', '<EOS>', '<UNK>']
        self.index2word = self.reserved[:]
        self.embeddings = None
    def add_words(self, words):
        """Add a new token to the vocab and do mapping between word and index.
        """
        for word in words:
            if word not in self.word2index:
                self.word2index[word] = len(self.index2word)
                self.index2word.append(word)
        self.word2count.update(words)

def build_vocab(self, embed_file: str = None) -> Vocab:
        """Build the vocabulary for the data set.
        """
        # word frequency
        word_counts = Counter()
        count_words(word_counts, [src + tgr for src, tgr in self.pairs])
        vocab = Vocab()
        # Filter the vocabulary by keeping only the top k tokens in terms of
        # word frequncy in the data set
        for word, count in word_counts.most_common(config.max_vocab_size):
            vocab.add_words([word])
        if embed_file is not None:
            count = vocab.load_embeddings(embed_file)
            print("%d pre-trained embeddings loaded." % count)
        return vocab


def count_words(counter, text):
    '''Count the number of occurrences of each word in a set of text'''
    for sentence in text:
        for word in sentence:
            counter[word] += 1

class SampleDataset(Dataset):
    """The class represents a sample set for training.
    """
    def __init__(self, data_pair, vocab):
        self.src_sents = [x[0] for x in data_pair]
        self.trg_sents = [x[1] for x in data_pair]
        self.vocab = vocab
        self._len = len(data_pair)
    def __getitem__(self, index):
        x, oov = source2ids(self.src_sents[index], self.vocab)
        return {
            'x': [self.vocab.SOS] + x + [self.vocab.EOS],
            'OOV': oov,
            'len_OOV': len(oov),
            'y': [self.vocab.SOS] +
            abstract2ids(self.trg_sents[index],
                         self.vocab, oov) + [self.vocab.EOS],
            'x_len': len(self.src_sents[index]),
            'y_len': len(self.trg_sents[index])
        }
    def __len__(self):
        return self._len

def source2ids(source_words, vocab):
    ids = []
    oovs = []
    unk_id = vocab.UNK
    for w in source_words:
        i = vocab[w]
        if i == unk_id:  
            if w not in oovs:  
                oovs.append(w)
            # This is 0 for the first source OOV, 1 for the second source OOV
            oov_num = oovs.index(w)
            # This is e.g. 20000 for the first source OOV, 50001 for the second
            ids.append(vocab.size() + oov_num)
        else:
            ids.append(i)
    return ids, oovs

def abstract2ids(abstract_words, vocab, source_oovs):
    """Map tokens in the abstract (reference) to ids.
       OOV tokens in the source will be remained.
    """
    ids = []
    unk_id = vocab.UNK
    for w in abstract_words:
        i = vocab[w]
        if i == unk_id:  # If w is an OOV word
            if w in source_oovs:  # If w is an in-source OOV
                # Map to its temporary source OOV number
                vocab_idx = vocab.size() + source_oovs.index(w)
                ids.append(vocab_idx)
            else:  # If w is an out-of-source OOV
                ids.append(unk_id)  # Map to the UNK token id
        else:
            ids.append(i)
    return ids

train_data = SampleDataset(dataset.pairs, v)
val_data = SampleDataset(val_dataset.pairs, v)
train_dataloader = DataLoader(dataset=train_data,
                                  batch_size=config.batch_size,
                                  shuffle=True,
                                  collate_fn=collate_fn)

def collate_fn(batch):
    """Split data set into batches and do padding for each batch.
    """
    def padding(indice, max_length, pad_idx=0):
        pad_indice = [item + [pad_idx] * max(0, max_length - len(item)) for item in indice]
        return torch.tensor(pad_indice)
    data_batch = sort_batch_by_len(batch)
    x = data_batch["x"]
    x_max_length = max([len(t) for t in x])
    y = data_batch["y"]
    y_max_length = max([len(t) for t in y])
    OOV = data_batch["OOV"]
    len_OOV = torch.tensor(data_batch["len_OOV"])
    x_padded = padding(x, x_max_length)
    y_padded = padding(y, y_max_length)
    x_len = torch.tensor(data_batch["x_len"])
    y_len = torch.tensor(data_batch["y_len"])
    return x_padded, y_padded, x_len, y_len, OOV, len_OOV

class Encoder(nn.Module):
    def __init__(self,vocab_size, embed_size, hidden_size, rnn_drop: float = 0):
        super(Encoder, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.hidden_size = hidden_size
        self.lstm = nn.LSTM(embed_size,
                            hidden_size,
                            bidirectional=True,
                            dropout=rnn_drop,
                            batch_first=True)
    def forward(self, x):
        """Define forward propagation for the endoer.
        """
        embedded = self.embedding(x)
        output, hidden = self.lstm(embedded)
        return output, hidden

class Attention(nn.Module):
    def __init__(self, hidden_units):
        super(Attention, self).__init__()
        # Define feed-forward layers.
        self.Wh = nn.Linear(2*hidden_units, 2*hidden_units, bias=False)
        self.Ws = nn.Linear(2*hidden_units, 2*hidden_units)

    def forward(self,
                decoder_states,
                encoder_output,
                x_padding_masks,
                coverage_vector):
        """Define forward propagation for the attention network.
        """
        # Concatenate h and c to get s_t and expand the dim of s_t.
        h_dec, c_dec = decoder_states
        s_t = torch.cat([h_dec, c_dec], dim=2)
        s_t = s_t.transpose(0, 1)
        s_t = s_t.expand_as(encoder_output).contiguous()
        # calculate attention scores
        encoder_features = self.Wh(encoder_output.contiguous())
        decoder_features = self.Ws(s_t)
        att_inputs = encoder_features + decoder_features
        score = self.v(torch.tanh(att_inputs))
        attention_weights = F.softmax(score, dim=1).squeeze(2)
        attention_weights = attention_weights * x_padding_masks
        # Normalize attention weights after excluding padded positions.
        normalization_factor = attention_weights.sum(1, keepdim=True)
        attention_weights = attention_weights / normalization_factor
        context_vector = torch.bmm(attention_weights.unsqueeze(1),
                                   encoder_output)
        context_vector = context_vector.squeeze(1)
        return context_vector, attention_weights

class Decoder(nn.Module):
    def __init__(self, vocab_size, embed_size, hidden_size, enc_hidden_size=None, is_cuda=True):
        super(Decoder, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.DEVICE = torch.device('cuda') if is_cuda else torch.device('cpu')
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.lstm = nn.LSTM(embed_size, hidden_size, batch_first=True)
        self.W1 = nn.Linear(self.hidden_size * 3, self.hidden_size)
        self.W2 = nn.Linear(self.hidden_size, vocab_size)

    def forward(self, x_t, decoder_states, context_vector):
        """Define forward propagation for the decoder.
        """
        decoder_emb = self.embedding(x_t)
        decoder_output, decoder_states = self.lstm(decoder_emb, decoder_states)
        # concatenate context vector and decoder state
        decoder_output = decoder_output.view(-1, config.hidden_size)
        concat_vector = torch.cat([decoder_output, context_vector], dim=-1)
        # calculate vocabulary distribution
        FF1_out = self.W1(concat_vector)
        FF2_out = self.W2(FF1_out)
        p_vocab = F.softmax(FF2_out, dim=1)
        h_dec, c_dec = decoder_states
        s_t = torch.cat([h_dec, c_dec], dim=2)
        return p_vocab, decoder_states

class RetduceSate(nn.Module):
    def __init__(self):
        super(ReduceState, self).__init__()

    def forward(self, hidden):
        """The forward propagation of reduce state module.
        """
        h, c = hidden
        h_reduced = torch.sum(h, dim=0, keepdim=True)
        c_reduced = torch.sum(c, dim=0, keepdim=True)
        return (h_reduced, c_reduced)

class Seq2seq(nn.Module):
    def __init__(self, v):
        super(Seq2seq, self).__init__()
        self.v = v
        self.DEVICE = config.DEVICE
        self.attention = Attention(config.hidden_size)
        self.encoder = Encoder(len(v),config.embed_size,config.hidden_size,)
        self.decoder = Decoder(len(v),config.embed_size,config.hidden_size,)
        self.reduce_state = ReduceState()

 
    def forward(self, x, x_len, y, len_oovs, batch, num_batches):
        """Define the forward propagation for the model.
        """
        x_copy = replace_oovs(x, self.v)
        x_padding_masks = torch.ne(x, 0).byte().float()
        encoder_output, encoder_states = self.encoder(x_copy)
        # Reduce encoder hidden states.
        decoder_states = self.reduce_state(encoder_states)
        # Calculate loss for every step.
        step_losses = []
        for t in range(y.shape[1]-1):
            # Do teacher forcing.
            x_t = y[:, t]
            x_t = replace_oovs(x_t, self.v)
            y_t = y[:, t+1]
            # Get context vector from the attention network.
            context_vector, attention_weights = self.attention(decoder_states, encoder_output, x_padding_masks, coverage_vector)
            # Get vocab distribution and hidden states from the decoder.
            p_vocab, decoder_states= self.decoder(x_t.unsqueeze(1), decoder_states, context_vector)
            # Get the probabilities predict by the model for target tokens.
            y_t = replace_oovs(y_t, self.v)
            target_probs = torch.gather(p_vocab, 1, y_t.unsqueeze(1))
            target_probs = target_probs.squeeze(1)
            # Apply a mask such that pad zeros do not affect the loss
            mask = torch.ne(y_t, 0).byte()
            # Do smoothing to prevent getting NaN loss because of log(0).
            loss = -torch.log(target_probs + config.eps)
            mask = mask.float()
            loss = loss * mask
            step_losses.append(loss)

        sample_losses = torch.sum(torch.stack(step_losses, 1), 1)
        # get the non-padded length of each sequence in the batch
        seq_len_mask = torch.ne(y, 0).byte().float()
        batch_seq_len = torch.sum(seq_len_mask, dim=1)
        # get batch loss by dividing the loss of each batch
        # by the target sequence length and mean
        batch_loss = torch.mean(sample_losses / batch_seq_len)
        return batch_loss

    def forward(self, x_t, decoder_states, context_vector):
        """Define forward propagation for the decoder.
        """
        decoder_emb = self.embedding(x_t)
        decoder_output, decoder_states = self.lstm(decoder_emb, decoder_states)
        # concatenate context vector and decoder state
        decoder_output = decoder_output.view(-1, config.hidden_size)
        concat_vector = torch.cat(
            [decoder_output,
             context_vector],
            dim=-1)
        # calculate vocabulary distribution
        FF1_out = self.W1(concat_vector)
        FF2_out = self.W2(FF1_out)
        p_vocab = F.softmax(FF2_out, dim=1)
        h_dec, c_dec = decoder_states
        s_t = torch.cat([h_dec, c_dec], dim=2)
        p_gen = None
        if config.pointer:
            # Calculate p_gen.
            x_gen = torch.cat([context_vector,s_t.squeeze(0),decoder_emb.squeeze(1)], dim=-1)
            p_gen = torch.sigmoid(self.w_gen(x_gen))
        return p_vocab, decoder_states, p_gen

    def forward(self, decoder_states, encoder_output, x_padding_masks, coverage_vector):
        """Define forward propagation for the attention network.
        """
        # Concatenate h and c to get s_t and expand the dim of s_t.
        h_dec, c_dec = decoder_states
        s_t = torch.cat([h_dec, c_dec], dim=2)
        s_t = s_t.transpose(0, 1)
        s_t = s_t.expand_as(encoder_output).contiguous()
        # calculate attention scores
        encoder_features = self.Wh(encoder_output.contiguous())
        decoder_features = self.Ws(s_t)
        att_inputs = encoder_features + decoder_features
        # Add coverage feature.
        if config.coverage:
            coverage_features = self.wc(coverage_vector.unsqueeze(2))  # wc c
            att_inputs = att_inputs + coverage_features
        score = self.v(torch.tanh(att_inputs))
        attention_weights = F.softmax(score, dim=1).squeeze(2)
        attention_weights = attention_weights * x_padding_masks

        # Normalize attention weights after excluding padded positions.
        normalization_factor = attention_weights.sum(1, keepdim=True)
        attention_weights = attention_weights / normalization_factor       
        context_vector = torch.bmm(attention_weights.unsqueeze(1), encoder_output)
        context_vector = context_vector.squeeze(1)
        # Update coverage vector.
        if config.coverage:
            coverage_vector = coverage_vector + attention_weights
        return context_vector, attention_weights, coverage_vector

    def get_final_distribution(self, x, p_gen, p_vocab, attention_weights, max_oov):
        """Calculate the final distribution for the model.
        """
        batch_size = x.size()[0]
        # Clip the probabilities.
        p_gen = torch.clamp(p_gen, 0.001, 0.999)
        # Get the weighted probabilities.
        p_vocab_weighted = p_gen * p_vocab
        attention_weighted = (1 - p_gen) * attention_weights

        # Get the extended-vocab probability distribution    
        extension = torch.zeros((batch_size, max_oov)).float().to(self.DEVICE)
        p_vocab_extended = torch.cat([p_vocab_weighted, extension], dim=1)
        # Add the attention weights to the corresponding vocab positions.
        final_distribution =  p_vocab_extended.scatter_add_(dim=1, index=x, src=attention_weighted)
        return final_distribution

def forward(self, x, x_len, y, len_oovs, batch, num_batches):
        """Define the forward propagation for the model.
        """
        x_copy = replace_oovs(x, self.v)
        x_padding_masks = torch.ne(x, 0).byte().float()
        encoder_output, encoder_states = self.encoder(x_copy)
        # Reduce encoder hidden states.
        decoder_states = self.reduce_state(encoder_states)
        # Initialize coverage vector.
        coverage_vector = torch.zeros(x.size()).to(self.DEVICE)
        # Calculate loss for every step.
        step_losses = []
        for t in range(y.shape[1]-1):
            # Do teacher forcing.
            x_t = y[:, t]
            x_t = replace_oovs(x_t, self.v)

            y_t = y[:, t+1]
            # Get context vector from the attention network.
            context_vector, attention_weights, next_coverage_vector = \
                self.attention(decoder_states,
                               encoder_output,
                               x_padding_masks,
                               coverage_vector)
            # Get vocab distribution and hidden states from the decoder.
            p_vocab, decoder_states, p_gen = self.decoder(x_t.unsqueeze(1),
                                                          decoder_states,
                                                          context_vector)

            final_dist = self.get_final_distribution(x,
                                                     p_gen,
                                                     p_vocab,
                                                     attention_weights,
                                                     torch.max(len_oovs))

            # Get the probabilities predict by the model for target tokens.
            target_probs = torch.gather(final_dist, 1, y_t.unsqueeze(1))
            target_probs = target_probs.squeeze(1)

            # Apply a mask such that pad zeros do not affect the loss
            mask = torch.ne(y_t, 0).byte()
            # Do smoothing to prevent getting NaN loss because of log(0).
            loss = -torch.log(target_probs + config.eps)

            # Add coverage loss.
            ct_min = torch.min(attention_weights, coverage_vector)
            cov_loss = torch.sum(ct_min, dim=1)
            loss = loss + config.LAMBDA * cov_loss
            coverage_vector = next_coverage_vector

            mask = mask.float()
            loss = loss * mask
            step_losses.append(loss)

        sample_losses = torch.sum(torch.stack(step_losses, 1), 1)
        # get the non-padded length of each sequence in the batch
        seq_len_mask = torch.ne(y, 0).byte().float()
        batch_seq_len = torch.sum(seq_len_mask, dim=1)

        # get batch loss by dividing the target sequence length and mean
        batch_loss = torch.mean(sample_losses / batch_seq_len)
        return batch_loss

在这里插入代码片

optimizer = None # Choose an optimizer from torch.optim
batch_losses = []
for batch in batches:
    model.train() # Sets the module in training mode.
    optimizer.zero_grad() # Clear gradients.
    batch_loss = model(**params)# Calculate loss for a batch.
    batch_losses.append(loss.item())
    batch_loss.backward() # Backpropagation.
    optimizer.step() # Update weights.
epoch_loss = torch.mean(batch_losses)

    optimizer = optim.Adam(model.parameters(), lr=config.learning_rate)
    writer = SummaryWriter(config.log_path)
    # tqdm: A tool for drawing progress bars during training.
    with tqdm(total=config.epochs) as epoch_progress:
        for epoch in range(start_epoch, config.epochs):
            batch_losses = []  # Get loss of each batch.
            num_batches = len(train_dataloader)
            with tqdm(total=num_batches//100) as batch_progress:               
                for batch, data in enumerate(tqdm(train_dataloader)):
                    x, y, x_len, y_len, oov, len_oovs = data
                    model.train()  # Sets the module in training mode.
                    optimizer.zero_grad()  # Clear gradients.
                    # Calculate loss.
                    loss = model(x, x_len, y, len_oovs, batch=batch, num_batches=num_batches)
                    batch_losses.append(loss.item())
                    loss.backward()  # Backpropagation.
                    # Do gradient clipping to prevent gradient explosion.
                    clip_grad_norm_(model.encoder.parameters(),
                                    config.max_grad_norm)
                    clip_grad_norm_(model.decoder.parameters(),
                                    config.max_grad_norm)
                    clip_grad_norm_(model.attention.parameters(),
                                    config.max_grad_norm)
                    optimizer.step()  # Update weights.
                    # Output and record epoch loss every 100 batches.
                    if (batch % 100) == 0:
                        batch_progress.set_description(f'Epoch {epoch}')
                        batch_progress.set_postfix(Batch=batch,
                                                   Loss=loss.item())
                        batch_progress.update()
                        # Write loss for tensorboard.
                        writer.add_scalar(f'Average loss for epoch {epoch}',
                                          np.mean(batch_losses),
                                          global_step=batch)
            # Calculate average loss over all batches in an epoch.
            epoch_loss = np.mean(batch_losses)
            epoch_progress.set_description(f'Epoch {epoch}')
            epoch_progress.set_postfix(Loss=epoch_loss)
            epoch_progress.update()
            avg_val_loss = evaluate(model, val_data, epoch)
            print('training loss:{}'.format(epoch_loss), 'validation loss:{}'.format(avg_val_loss))

    writer.close()

    def greedy_search(self,
                      x,
                      max_sum_len,
                      len_oovs,
                      x_padding_masks):
        """Function which returns a summary by always picking
        """

        # Get encoder output and states.
        encoder_output, encoder_states = self.model.encoder(
            replace_oovs(x, self.vocab))

        # Initialize decoder's hidden states with encoder's hidden states.
        decoder_states = self.model.reduce_state(encoder_states)

        # Initialize decoder's input at time step 0 with the SOS token.
        x_t = torch.ones(1) * self.vocab.SOS
        x_t = x_t.to(self.DEVICE, dtype=torch.int64)
        summary = [self.vocab.SOS]
        coverage_vector = torch.zeros((1, x.shape[1])).to(self.DEVICE)
        # Generate hypothesis with maximum decode step.
        while int(x_t.item()) != (self.vocab.EOS) \
                and len(summary) < max_sum_len:
            context_vector, attention_weights, coverage_vector = \
                self.model.attention(decoder_states,
                                     encoder_output,
                                     x_padding_masks,
                                     coverage_vector)
            p_vocab, decoder_states, p_gen = \
                self.model.decoder(x_t.unsqueeze(1),
                                   decoder_states,
                                   context_vector)
            final_dist = self.model.get_final_distribution(x,
                                                           p_gen,
                                                           p_vocab,
                                                           attention_weights,
                                                           torch.max(len_oovs))
            # Get next token with maximum probability.
            x_t = torch.argmax(final_dist, dim=1).to(self.DEVICE)
            decoder_word_idx = x_t.item()
            summary.append(decoder_word_idx)
            x_t = replace_oovs(x_t, self.vocab)
        return summary

class Beam(object):
    def __init__(self,
                 tokens,
                 log_probs,
                 decoder_states,
                 coverage_vector):
        self.tokens = tokens
        self.log_probs = log_probs
        self.decoder_states = decoder_states
        self.coverage_vector = coverage_vector

    def extend(self,
               token,
               log_prob,
               decoder_states,
               coverage_vector):
        return Beam(tokens=self.tokens + [token],
                    log_probs=self.log_probs + [log_prob],
                    decoder_states=decoder_states,
                    coverage_vector=coverage_vector)
    def seq_score(self):
        """
        This function calculate the score of the current sequence.
        """
        len_Y = len(self.tokens)
        # Lenth normalization
        ln = (5+len_Y)**config.alpha / (5+1)**config.alpha
        cn = config.beta * torch.sum(  # Coverage normalization
            torch.log(
                config.eps +
                torch.where(
                    self.coverage_vector < 1.0,
                    self.coverage_vector,
                    torch.ones((1, self.coverage_vector.shape[1])).to(torch.device(config.DEVICE))
                )
            )
        )
        score = sum(self.log_probs) / ln + cn
        return score

    def best_k(self, beam, k, encoder_output, x_padding_masks, x, len_oovs):
        """Get best k tokens to extend the current sequence at the current time step.
        """
        # use decoder to generate vocab distribution for the next token
        x_t = torch.tensor(beam.tokens[-1]).reshape(1, 1)
        x_t = x_t.to(self.DEVICE)

        # Get context vector from attention network.
        context_vector, attention_weights, coverage_vector = \
            self.model.attention(beam.decoder_states,
                                 encoder_output,
                                 x_padding_masks,
                                 beam.coverage_vector)
、
        p_vocab, decoder_states, p_gen = \
            self.model.decoder(replace_oovs(x_t, self.vocab),
                               beam.decoder_states,
                               context_vector)

        final_dist = self.model.get_final_distribution(x,
                                                       p_gen,
                                                       p_vocab,
                                                       attention_weights,
                                                       torch.max(len_oovs))
        # Calculate log probabilities.
        log_probs = torch.log(final_dist.squeeze())
        # EOS token penalty. Follow the definition in
        log_probs[self.vocab.EOS] *= \
            config.gamma * x.size()[1] / len(beam.tokens)
        log_probs[self.vocab.UNK] = -float('inf')
        # Get top k tokens and the corresponding logprob.
        topk_probs, topk_idx = torch.topk(log_probs, k)
        best_k = [beam.extend(x,
                  log_probs[x],
                  decoder_states,
                  coverage_vector) for x in topk_idx.tolist()]
        return best_k

    def beam_search(self,
                    x,
                    max_sum_len,
                    beam_width,
                    len_oovs,
                    x_padding_masks):
        """Using beam search to generate summary.
        """
        # run body_sequence input through encoder
        encoder_output, encoder_states = self.model.encoder(
            replace_oovs(x, self.vocab))
        coverage_vector = torch.zeros((1, x.shape[1])).to(self.DEVICE)
        # initialize decoder states with encoder forward states
        decoder_states = self.model.reduce_state(encoder_states)
        # initialize the hypothesis with a class Beam instance.
        init_beam = Beam([self.vocab.SOS],
                         [0],
                         decoder_states,
                         coverage_vector)
        k = beam_width
        curr, completed = [init_beam], []
        # use beam search for max_sum_len (maximum length) steps
        for _ in range(max_sum_len):
            # get k best hypothesis when adding a new token
            topk = []
            for beam in curr:
                # When an EOS token is generated, add the hypo to the completed
                if beam.tokens[-1] == self.vocab.EOS:
                    completed.append(beam)
                    k -= 1
                    continue
                for can in self.best_k(beam,
                                       k,
                                       encoder_output,
                                       x_padding_masks,
                                       x,
                                       torch.max(len_oovs)
                                       ):
                    # Using topk as a heap to keep track of top k candidates.
                    add2heap(topk, (can.seq_score(), id(can), can), k)

            curr = [items[2] for items in topk]
            # stop when there are enough completed hypothesis
            if len(completed) == beam_width:
                break
        completed += curr
        # sort the hypothesis by normalized probability and choose the best one
        result = sorted(completed,
                        key=lambda x: x.seq_score(),
                        reverse=True)[0].tokens
        return result

class RougeEval():
    def __init__(self, path):
        self.path = path
        self.scores = None
        self.rouge = Rouge()
        self.sources = []
        self.hypos = []
        self.refs = []
        self.process()
    def process(self):
        print('Reading from ', self.path)
        with open(self.path, 'r') as test:
            for line in test:
                source, ref = line.strip().split('<sep>')
                ref = ''.join(list(jieba.cut(ref))).replace('。', '.')
                self.sources.append(source)
                self.refs.append(ref)
        print(f'Test set contains {len(self.sources)} samples.')


    def build_hypos(self, predict):
        print('Building hypotheses.') 
        count = 0
        for source in self.sources:
            count += 1
            if count % 100 == 0:
                print(count)
            self.hypos.append(predict.predict(source.split()))

    def get_average(self):
        assert len(self.hypos) > 0, 'Build hypotheses first!'
        print('Calculating average rouge scores.')
        return self.rouge.get_scores(self.hypos, self.refs, avg=True)

    def one_sample(self, hypo, ref):
        return self.rouge.get_scores(hypo, ref)[0]
rouge_eval = RougeEval(config.test_data_path)
predict = Predict()
rouge_eval.build_hypos(predict)
result = rouge_eval.get_average()
print('rouge1: ', result['rouge-1'])
print('rouge2: ', result['rouge-2'])
print('rougeL: ', result['rouge-l'])

    def extract_keywords(self, dct, tfidf, threshold=0.2, topk=5):
        """find high TFIDF socore keywords
        """
        tfidf = sorted(tfidf, key=lambda x: x[1], reverse=True)
        return list(islice(
            [dct[w] for w, score in tfidf if score > threshold], topk
            ))

    def replace(self, token_list, doc):
        """replace token by another token which is similar in wordvector 
        """
        keywords = self.extract_keywords(self.dct, self.tfidf_model[doc])
        num = int(len(token_list) * 0.3)
        new_tokens = token_list.copy()
        while num == int(len(token_list) * 0.3):
            indexes = np.random.choice(len(token_list), num)
            for index in indexes:
                token = token_list[index]
                if isChinese(token) and token not in keywords and token in self.wv:
                    new_tokens[index] = self.wv.most_similar(
                        positive=token, negative=None, topn=1
                        )[0][0]
            num -= 1
        return ' '.join(new_tokens)

    def generate_samples(self, write_path):
        """generate new samples file
        """
        replaced = []
        count = 0
        for sample, token_list, doc in zip(self.samples, self.refs, self.corpus):
            count += 1
            if count % 100 == 0:
                print(count)
                write_samples(replaced, write_path, 'a')
                replaced = []
            replaced.append(
                sample.split('<sep>')[0] + ' <sep> ' + self.replace(token_list, doc)
                )

def translate(q, source, target):
    """translate q from source language to target language
    """
    #  refer to the official documentation   https://api.fanyi.baidu.com/  
    # There are demo on the website ,  register on the web site ,and get AppID, key, python3 demo.
    appid = ''  # Fill in your AppID
    secretKey = ''  # Fill in your key
    httpClient = None
    myurl = '/api/trans/vip/translate'
    fromLang = source  # The original language
    toLang = target  # The target language
    salt = random.randint(32768, 65536)
    sign = appid + q + str(salt) + secretKey
    sign = hashlib.md5(sign.encode()).hexdigest()
    myurl = '/api/trans/vip/translate' + '?appid=' + appid + '&q=' + urllib.parse.quote(
        q) + '&from=' + fromLang + '&to=' + toLang + '&salt=' + str(
        salt) + '&sign=' + sign
    try:
        httpClient = http.client.HTTPConnection('api.fanyi.baidu.com')
        httpClient.request('GET', myurl)
        # response is HTTPResponse object
        response = httpClient.getresponse()
        result_all = response.read().decode("utf-8")
        result = json.loads(result_all)
        return result
    except Exception as e:
        print(e)
    finally:
        if httpClient:
            httpClient.close()

def back_translate(q):
    """back_translate
    """
    en = translate(q, "zh", 'en')['trans_result'][0]['dst']
    time.sleep(1.5)
    target = translate(en, "en", 'zh')['trans_result'][0]['dst']
    time.sleep(1.5)
    return target

def semi_supervised(samples_path, write_path, beam_search):
    """use reference to predict source
    """
    pred = Predict()
    print('vocab_size: ', len(pred.vocab))
    # Randomly pick a sample in test set to predict.
    count = 0
    semi = []
    with open(samples_path, 'r') as f:
        for picked in f:
            count += 1
            source, ref = picked.strip().split('<sep>')
            prediction = pred.predict(ref.split(), beam_search=beam_search)
            semi.append(prediction + ' <sep> ' + ref)
            if count % 100 == 0:
                write_samples(semi, write_path, 'a')
                semi = []

# source 文本
帕莎 太阳镜 男⼠ 太阳镜 偏光 太阳眼镜 墨镜 潮 典雅 灰 颜⾊ 选择 ， 细节 特⾊ 展示 ， ⻆度 展示 ， ⻛ 尚 演 ， ⾦属， 透光 量 ， 不易 变形 ， 坚固耐⽤ ， 清晰 柔和 ， 镜⽚ 材质 ， 尼⻰ ⾼ ， 型号 ， 镜架 材质 ， ⾦属 镜腿 ， 镜布 镜盒 ， ⿐ 间距 ， 轻盈 ⿐托 ， 品牌 刻印 ， ⽆缝 拼接 眼镜 配件类型 镜盒 功 能 偏光 类型 偏光太阳镜 镜⽚材质 树脂 ⻛格 休闲⻛ 上市时间 2016年夏季 镜框形状 ⽅形 适⽤性别 男 镜架材质 合⾦

# reference文本
夏天 到 了 ， 在 刺眼 的 阳光 下 少不了 这 款 时尚 的 男⼠ 太阳镜 ！ 时尚 的 版型 ， 适⽤ 各种 脸型 ， 突出 您 的 型 男 ⻛范 。 其⾼ 镍 ⾦属 材质 的 镜架 ， ⼗分 的 轻盈 ， 带给 您 舒适 的 佩戴 体验 。

# 将reference送⼊模型后⽣成的数据new source：

版型 设计 ， 时尚 百搭 ， 适合 多种 场合 佩戴 。 时尚 ⾦属 拉链 ， 经久耐⽤ 。 带给 你 舒适 的 佩戴 体验 。 ⻛范 与 铰链 的 镜架 ， 舒适 耐磨 ， 不易 变形 。 带给 你 意想不到 的 修身 效果 ， 让 你 的 夏 季 充满 着 不 ⼀样 的 魅⼒ ， 让 这个 夏天 格外 绚烂 ， 让 这个 夏天 格外 绚烂 。 时尚 的 版型 ， 让 你 穿 起来 更 有 男⼈味 ， 让 这个 夏天格外 绚烂 绚烂 ， 让 这个 夏天 格外 绚烂 。 时尚 的 版型 ， 让 你 穿 起来 更 有 男⼈味 ， 让 这个 夏天 格外 绚烂 绚烂 ， 是 秋装 上 上 之选 。

        if config.weight_tying:
            embedded = decoder_embedding(x)
        else:
            embedded = self.embedding(x)

    if config.weight_tying:
            FF2_out = torch.mm(FF1_out, torch.t(self.embedding.weight))
        else:
            FF2_out = self.W2(FF1_out)

class ScheduledSampler():
    def __init__(self, phases):
        self.phases = phases
        self.scheduled_probs = [i / (self.phases - 1) for i in range(self.phases)]

    def teacher_forcing(self, phase):
        """According to a certain probability to choose whether to execute teacher_forcing
        """
        sampling_prob = random.random()
        if sampling_prob >= self.scheduled_probs[phase]:
            return True
        else:
            return False