Trying to save a model with TFT5ForConditionalGeneration

[erlichsefisalesforce]

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, 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 followingif statement:

        # 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

[amyeroberts]

cc @gante

[gante]

Hey @erlichsefisalesforce 👋 looking at the stack trace, we see that inputs's first dimension, the batch size, is unknown (shape = [None, 1]). It is possible that our generate function is not fully serializable with a dynamic batch size, and may need some tweaks.

I'm not sure when I'll be able to fix this problem in particular (it may be complex to solve). However, meanwhile, can you try exporting with a fixed batch size? In other words, define the input as inputs = tf.keras.layers.Input(shape=(1,), dtype=tf.string, name="inputs", batch_size=<some integer>)

[erlichsefisalesforce]

Hi @gante, I'm still getting

  File "/venv/lib/python3.9/site-packages/transformers/generation/tf_utils.py", line 767, in generate
    return self.greedy_search(
  File "venv/lib/python3.9/site-packages/transformers/generation/tf_utils.py", line 1452, 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 'summarizer' (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 'summarizer' (type CompleteSentenceTransformer):
  • args=('tf.Tensor(shape=(3, 1), dtype=string)',)
  • kwargs={'training': 'False'}

[gante]

Hey @erlichsefisalesforce -- in that case, I will need a reproducible example to debug. The example you shared above contains references to local files :)

[erlichsefisalesforce]

@gante, the folder contains flan-t5-large, save_dir is can be populated with any path to your local machine, and I think, that it.

[gante]

It seems like the root issue persists -- text.SentencepieceTokenizer().tokenize() returns a tensor with an unknown batch size, regardless of the input batch size being defined, causing the same problem.

The fix should be straightforward, so I will have a go at it.

---

Script to reproduce it:

# run these commands in advance:
# mkdir /tmp/test
# cd /tmp/test
# git clone https://huggingface.co/google/flan-t5-small

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

save_dir = '/tmp/test/flan-t5-small'
class CompleteSentenceTransformer(tf.keras.layers.Layer):

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

    def call(self, inputs, *args, **kwargs):
        tokens = self.tokenizer.tokenize(inputs)
        breakpoint()
        input_ids, attention_mask = text.pad_model_inputs(tokens, max_seq_length=512, pad_value=self.model.config.pad_token_id)
        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", batch_size=4)
outputs = complete_model(inputs)
keras_model = tf.keras.Model(inputs, outputs)
keras_model.save(save_dir)

[gante]

@erlichsefisalesforce after #22310 gets merged, you should be able to run it on your end :) (you will need to install transformers from main)

[erlichsefisalesforce]

Thank you @gante! will close the issue once I validate the solution on my end. :)

[erlichsefisalesforce]

The solution was validated!