NEZHA(中文版GPT2)实现Beam-search Tensorflow1.15 (一)

平时一直使用Transformers的包来调用GPT2或者生成模型，其generate函数封装了top-k和top-p以及beam-search函数。最近在使用华为的NEZHA(中文GPT2)，这个模型的生成模型是用Tensorflow1.15写的。它的sample方式是top-k和top-p的sampling，最近需要使用beam-search，因此根据NEZHA的代码来实现beam-search。

在输入到模型之前，需要对输入copy为beam size的大小，相当于是在生成的时候维护beam size大小个句子：

predict = init_input_ids
predict = tf.tile(predict, [beam_width, 1])

首先，将copy后的predict输入到模型中，step为通过NEZHA得到了预测出的输出logits, 保存中间结果(transformer的k,v等值，用于加速decoding的速度)的past，以及当前处于生成的第几步predict_token_idx（初始值为0，每运行一次加一）

predict_logits, past, predict_token_idx = step(predict, past=past,     
                                predict_token_idx=predict_token_idx)

将输出的概率通过log softmax，之所以要加上log，是因为beam-search在选择的时候是通过语言模型的概率累乘，而加上log后把累乘转换成累加，方便计算。

log_probs = tf.nn.log_softmax(predict_logits, axis=-1)

之后需要分情况，如果是第一步生成，由于原先的predict中的beam width个句子都是一样，因此predict_logits的对于第一步来说，只需要选predict_logits（beam width, vocab size）的其中一个predict_logits[0]，从里面挑选出top-k（k为beam width）个选项，将这topk个的log probs作为初始状态的score(sel_sum_logprobs)

sel_sum_logprobs, indices = tf.nn.top_k(tf.expand_dims(log_probs[0], 0), k=beam_width)

第一步得到的beam width个候选，将其赋值给predicted_tokens。

predicted_tokens = tf.expand_dims(tf.identity(predict), 2)

第一步初始化就结束了，之后每一步就根据初始化的结果进行操作，对于接下来每一步，需要将每一步生成的预测的输出概率log_probs和之前的累加的概率sel_sum_logprobs相加：

predict_logits, past, predict_token_idx = step(predict, past=past, 
                                     predict_token_idx=predict_token_idx)
log_probs = tf.nn.log_softmax(predict_logits, axis=-1)
sum_logprobs = (tf.expand_dims(sel_sum_logprobs, axis=2) + (log_probs))

从累加后的概率中选出，top-k个候选，在top-k之前需要将sum_logprobs的大小进行转换（(batch, beam width, vocab_size) -->(batch, beam width * vocab_size) 。其原因是因为要从所有beam的候选词中进行选择

sel_sum_logprobs, indices = tf.nn.top_k(
            tf.reshape(sum_logprobs, [batch_size, vocab_size* beam_width]),
            k=beam_width)

 其得到的indices是beam width所有的indices，需要知道它在哪一个beam，以及在beam内部中的indice

predict = indices % vocab_size
beam_ids = indices // vocab_size

对于每一个预测的词predict，要判断它是属于哪一个beam id，并把它和beam id对应的已经生成的predicted_tokens拼接，同时past也需要根据beam id更新的past，其原因是得到了新的top-k个predict token，其对应的已经生成的predicted tokens不一定相同，因此需要将past进行更新。

predicted_tokens = tf.concat([batch_gather(predicted_tokens, beam_ids), tf.expand_dims(predict, 2)], axis=2)

past = tf.gather(past, tf.squeeze(beam_ids), axis=2)

这样经过若干轮后，得到的predicted_tokens就是得到色sequence。

拿“一时间风云变幻”作为输入，一般来说beam-search的效果会比greedy的方法要好一些。但是感觉重复现象很严重，这里也实现了repetition penality=1.2，其重复现象会缓解很多。repetition penality的实现以及对于在tensorflow1.15中优化代码放在后篇之后再写

方法	句子
greedy	一时间风云变幻，各种各样的新闻事件层出不穷，各种各样的人物也层出不穷，各种各样的人物也层出不穷，各种各样的人物也层出
beam-search	一时间风云变幻，各种各样的传闻纷至沓来。一时间，各种各样的传闻纷至沓来。一时间，各种各样的传闻纷至沓来。一时间，各种
beam-search （r_p=1.2）	一时间风云变幻，各种各样的传闻纷至沓来。在这种情况下，我们不得不面对一个现实，那就是中国的网络游戏市场已经进入了一个