

TITLE
Trying to save a model with TFT5ForConditionalGeneration

COMMENTS
9

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
### System Info

transformers-cli env ouput:
- `transformers` version: 4.28.0.dev0
- Platform: macOS-13.2.1-x86_64-i386-64bit
- Python version: 3.9.6
- Huggingface_hub version: 0.12.0
- PyTorch version (GPU?): 1.13.1 (False)
- Tensorflow version (GPU?): 2.11.0 (False)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: no


--------------

I've compiled a TensorFlow graph that uses a pre-trained [flan-t5-large](https://huggingface.co/google/flan-t5-large), which means one of the layers uses `TFT5ForConditionalGeneration` but there are more layers before and after and my goal is the export the graph for TF serving framework. 

When I'm trying to `.save` the model I get the following error from Tensorflow:
```
Traceback (most recent call last):


Traceback (most recent call last):
  File "/Users/serlich/Library/Application Support/JetBrains/Toolbox/apps/PyCharm-P/ch-0/223.8214.51/PyCharm.app/Contents/plugins/python/helpers/pydev/_pydevd_bundle/pydevd_exec2.py", line 3, in Exec
    exec(exp, global_vars, local_vars)
  File "<string>", line 1, in <module>
  File "venv/lib/python3.9/site-packages/tensorflow/python/framework/ops.py", line 955, in __bool__
    self._disallow_bool_casting()
  File "venv/lib/python3.9/site-packages/tensorflow/python/framework/ops.py", line 554, in _disallow_bool_casting
    self._disallow_when_autograph_enabled(
  File "venv/lib/python3.9/site-packages/tensorflow/python/framework/ops.py", line 537, in _disallow_when_autograph_enabled
    raise errors.OperatorNotAllowedInGraphError(
tensorflow.python.framework.errors_impl.OperatorNotAllowedInGraphError: Using a symbolic `tf.Tensor` as a Python `bool` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.
Traceback (most recent call last):
  File "venv/lib/python3.9/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.9/lib/python3.9/contextlib.py", line 124, in __exit__
    next(self.gen)
  File "/Users/serlich/Documents/case-wrap-up/t5_tensorflow_code.py", line 58, in call
    outputs = self.model.generate(input_ids=input_ids, attention_mask=attention_mask)
  File "venv/lib/python3.9/site-packages/transformers/generation/tf_utils.py", line 925, in generate
    return self.greedy_search(
  File "venv/lib/python3.9/site-packages/transformers/generation/tf_utils.py", line 1728, in greedy_search
    if greedy_search_cond_fn(generated, finished_sequences, cur_len, model_kwargs):
tensorflow.python.framework.errors_impl.OperatorNotAllowedInGraphError: Exception encountered when calling layer 'complete_sentence_transformer' (type CompleteSentenceTransformer).

Using a symbolic `tf.Tensor` as a Python `bool` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

Call arguments received by layer 'complete_sentence_transformer' (type CompleteSentenceTransformer):
  • args=('tf.Tensor(shape=(None, 1), dtype=string)',)
  • kwargs={'training': 'False'}

Process finished with exit code 1

```

the error resonates  from the following[if statement](https://github.com/huggingface/transformers/blob/60d51ef5123d949fd8c59cd4d3254e711541d278/src/transformers/generation/tf_utils.py#L1728): 
```
        # 1st generation step has to be run before to initialize `past_key_values`
        generated, finished_sequences, cur_len, model_kwargs = greedy_search_body_fn(
            generated, finished_sequences, cur_len, model_kwargs
        )

        # 2-to-n generation steps can then be run in autoregressive fashion
        # only in case 1st generation step does NOT yield EOS token though
        if greedy_search_cond_fn(generated, finished_sequences, cur_len, model_kwargs):
            maximum_iterations = max_length - cur_len
            generated, _, cur_len, _ = tf.while_loop(
                greedy_search_cond_fn,
                greedy_search_body_fn,
                (generated, finished_sequences, cur_len, model_kwargs),
                maximum_iterations=maximum_iterations,
            )

```
during saving `finished_sequences` is a symbolic tensor and so, TensorFlow prevents evaluating an if statement of a symbolic tensor.


Commenting out 
`if greedy_search_cond_fn(generated, finished_sequences, cur_len, model_kwargs)
` allow me to save the model and load it later, however, remove the safeguard if the model predicts EOS in the first generation (which is very not likely).

@ArthurZucker @younesbelkada @gante   


### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

```
from transformers import TFT5ForConditionalGeneration
import tensorflow as tf
import tensorflow_text as text
from tensorflow.python.platform import gfile

save_dir = ''
class CompleteSentenceTransformer(tf.keras.layers.Layer):

    def __init__(self):
        super().__init__()
        self._pad_token = 1
        self.tokenizer = text.SentencepieceTokenizer(model=gfile.GFile('test/flan-t5-large/spiece.model', 'rb').read())
        self.model = TFT5ForConditionalGeneration.from_pretrained('test/flan-t5-large', from_pt=True)

    def call(self, inputs, *args, **kwargs):
        tokens = self.tokenizer.tokenize(inputs)
        input_ids, attention_mask = text.pad_model_inputs(tokens, max_seq_length=self._max_seq_length, pad_value=self._pad_token)
        outputs = self.model.generate(input_ids=input_ids, attention_mask=attention_mask)
        return self.tokenizer.detokenize(outputs)


complete_model = CompleteSentenceTransformer()
inputs = tf.keras.layers.Input(shape=(1,), dtype=tf.string, name="inputs")
outputs = complete_model(inputs)
keras_model = tf.keras.Model(inputs, outputs)
keras_model.save(save_dir)
```
Python 3.9.6
tensorflow                    2.11.0
tensorflow-text                  2.11.0
transformers                  4.28.0.dev0 (from master)

### Expected behavior

save model