(TF) model.generate to tf.function for tf serving
Author: gyin94
Pull request: no
Labels: name: wontfix

# ❓ Questions & Help

How can we wrap the model.generate and export it as a part of savedModel pb file? In this way, we can use beam search or topK during the tf serving or converting it to coremltools model.

## Details

I am trying to find a way to wrap the model in a Keras Model. But apparently model.generate is not tf.function supported. Like for loop is not supported in tf.function. 
```
from transformers import *

class WrapModel(tf.keras.models.Model):
    def __init__(self, transformer):
        super(WrapModel, self).__init__()
        self.transformer = transformer

    @tf.function
    def _internal_generate(self, inputs):
        return self.transformer.generate(inputs, max_length=10, length_penalty=1.0, repetition_penalty=2.5,
                                         early_stopping=True, num_beams=3)

    def call(self, inputs, **kwargs):
        print(inputs.shape)
        res = self._internal_generate(inputs)
        return res

gpt2_model = TFGPT2LMHeadModel.from_pretrained('distilgpt2')
w = WrapModel(gpt2_model)
input_layer = tf.keras.layers.Input(shape=10, dtype=tf.int32, name='input_ids')
prediction_model = w(input_layer)
tf_model = tf.keras.models.Model(inputs=input_layer, outputs=prediction_model)

import coremltools as ct
mlmodel = ct.convert(tf_model)
```


Comments:
Hey @gyin-ai, 

Thanks a lot for the issue! Currently `generate` does not seem to be compatible with `tf.function`. I will open an issue about this and hopefully fix generate so that it will become possible to generate using `tf.function`. You're use case should definitely be possible in the future!

I assume a lot of operations will have to changed in the tf generate function though, so this PR might take a while.

Will try to start on this next week: #5662

@patrickvonplaten perfect! Look forwarding to this feature so that we could use the LM model with various decoding solutions directly in the TF Serving or on-device. 

Yeah, it not going to that easy :D 
I will be on holiday for two weeks now, but we will be starting to put the focus much more on TF soon!

Also pinging @jplu here, just for notification.

Hey @gyin-ai 

I have done some work here https://github.com/huggingface/transformers/pull/5468 for making the models, saved model compliants. Can you try it to see if it might solve your issue?

@patrickvonplaten has also started to do some great work on the LM part of TF, so yes @gyin-ai you can expect to have better TF compliancy soon ;)

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.


> Will try to start on this next week: #5662

@patrickvonplaten Has the problem been solved

@patrickvonplaten Any update on this? Seems `model.generate` is still not compatible with `tf.funtion`.

cc @Rocketknight1 

Hi @yuwon, I'm (one of!) the current TF maintainers. We've experimented with wrapping all of `generate()` in a tf.function, but we generally find that buffers are not freed properly after each token is generated and OOM errors usually result after a few steps. `generate()` is important and so we're planning a complete investigation of this to see if there's any way we could make it work, but it's a sizeable project with a lot of other competing priorities and we don't have a concrete ETA right now.

@Rocketknight1 can you share you you've done on wrapping generate inside a tf function? It might be a start point for us to submit a PR and try to solve it.