TFlite conversion not working for GPT2 model with new package imports

[zhubarb]

This dated google colab notebook shows GPT-2 finetuning and exporting the resulting to model to tflite. It runs front to back without issues. The specific imports and the versions of packages used in the notebook are:

!pip install -q git+https://github.com/keras-team/keras-nlp.git@google-io-2023 tensorflow-text==2.12
import numpy as np
import keras_nlp
import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_text as tf_text
from tensorflow import keras
from tensorflow.lite.python import interpreter
import time
from google.colab import files
print(tf.__version__)
print(keras.__version__)
print(keras_nlp.__version__)
2.12.1
2.12.0
0.5.0

I have modified the import as below to enable GPU support, but now the the tflite conversion does not work, throwing the error below. Any ideas where the incompatibility is arising from?

!pip install keras_nlp
print(tf.__version__)
print(keras.__version__)
print(keras_nlp.__version__)
2.16.1
3.3.3
0.11.1

The error occurs while trying to Convert the generate() function from GPT2CausalLM:

@tf.function
def generate(prompt, max_length):
    return gpt2_lm.generate(prompt, max_length)

concrete_func = generate.get_concrete_function(tf.TensorSpec([], tf.string), 100)

Error:

NotImplementedError                       Traceback (most recent call last)
[<ipython-input-10-d8866e8eac5d>](https://localhost:8080/#) in <cell line: 5>()
      3     return gpt2_lm.generate(prompt, max_length)
      4 
----> 5 concrete_func = generate.get_concrete_function(tf.TensorSpec([], tf.string), 100)

27 frames
[/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/tensor.py](https://localhost:8080/#) in __array__(***failed resolving arguments***)
    625   def __array__(self, dtype=None):
    626     del dtype
--> 627     raise NotImplementedError(
    628         f"Cannot convert a symbolic tf.Tensor ({self.name}) to a numpy array."
    629         f" This error may indicate that you're trying to pass a Tensor to"

NotImplementedError: in user code:

    File "<ipython-input-10-d8866e8eac5d>", line 3, in generate  *
        return gpt2_lm.generate(prompt, max_length)
    File "/usr/local/lib/python3.10/dist-packages/keras_nlp/src/models/causal_lm.py", line 371, in postprocess  *
        return self.preprocessor.generate_postprocess(x)
    File "/usr/local/lib/python3.10/dist-packages/keras_nlp/src/models/gpt2/gpt2_causal_lm_preprocessor.py", line 178, in generate_postprocess  *
        token_ids = ops.convert_to_numpy(token_ids)
    File "/usr/local/lib/python3.10/dist-packages/keras/src/ops/core.py", line 512, in convert_to_numpy  **
        return backend.convert_to_numpy(x)
    File "/usr/local/lib/python3.10/dist-packages/keras/src/backend/tensorflow/core.py", line 131, in convert_to_numpy
        return np.asarray(x)

    NotImplementedError: Cannot convert a symbolic tf.Tensor (StatefulPartitionedCall:1) to a numpy array. This error may indicate that you're trying to pass a Tensor to a NumPy call, which is not supported.