GPT-2 Logits to tokens for beam search (Generate method)

I have a TF GPT-2 LMHead model running on TF Serving and I want to do a beam search(multiple tokens output) with the models’ output logits.

payload = {“inputs”: input_padded}
requests.post(‘http://localhost:8501/v1/models/gpt2-farmacos5:predict’, data=json.dumps(payload))

That request returns me a tf.tensor of logits for next token. But how do I convert these logits to multiple tokens output like the huggingface’s gpt2model.generate() method (but in tf serving) ?

PD I know there is transformers/generation_tf_utils.py at master ¡ huggingface/transformers ¡ GitHub but I need a simpler implementation.

PD2: I know this post is similar to patrick’s post Generation Probabilities: How to compute probabilities of output scores for GPT2 but he’s working on pytorch and Im in tensorflow. Mabye i should migrate to Pytorch ?

Thanks

2 Likes