Issue when saving TF model with a tokenizer as a custom layer

Hi,

I am trying to create a tensorflow model with keras api, when I include the tokenizing process inside the model. It seems to work for the inference locally, but when I am saving the model with tf.saved_model.save , I got an error. I am wondering if there is something wrong in my current code, or if it is currently not possible ?

AssertionError: Tried to export a function which references untracked object Tensor("139395:0", shape=(), dtype=resource).TensorFlow objects (e.g. tf.Variable) captured by functions must be tracked by assigning them to an attribute of a tracked object or assigned to an attribute of the main object directly.

My tokenizer which use the BertTokenizer from tensorflow_text (I take the code from some discussion in this forum and modify it) :

class TokenizerTF(tf.Module):
    def __init__(self, vocab_file_path, sequence_length=512, lower_case=True, pad_id=1, cls_id=2, sep_id=3):
        self.cls_token_id = tf.constant(cls_id, dtype=tf.int32)
        self.sep_token_id = tf.constant(sep_id, dtype=tf.int32)
        self.pad_token_id = tf.constant(pad_id, dtype=tf.int32)

        self.sequence_length = tf.constant(sequence_length)



        # These two lines are basically what makes it work
        # assigning the vocab to a tf.Module and then later assigning the
        # intantiated Module to e.g. a Keras Model
        self.bert_tokenizer = tf_text.BertTokenizer(
            vocab_file_path,
            lower_case=lower_case,
        )

    @tf.function
    def __call__(self, text: tf.Tensor) -> tf.Tensor:
        """
        Perform the BERT preprocessing from text -> input token ids
        """
        # Convert text into token ids
        tokens = self.bert_tokenizer.tokenize(text)

        # Flatten the ragged tensors
        tokens =  tf.cast(tokens.merge_dims(1, 2), tf.int32)

        # Add start and end token ids to the id sequence
        start_tokens = tf.fill([tf.shape(text)[0], 1], self.cls_token_id)
        end_tokens = tf.fill([tf.shape(text)[0], 1], self.sep_token_id)
        tokens = tf.concat([start_tokens, tokens, end_tokens], axis=1)

        # Truncate to sequence length
        tokens = tokens[:, : self.sequence_length]

        # Convert ragged tensor to tensor and pad with PAD_ID
        tokens = tokens.to_tensor(default_value=self.pad_token_id)

        # Pad to sequence length
        pad = self.sequence_length - tf.shape(tokens)[1]
        tokens = tf.pad(tokens, [[0, 0], [0, pad]], constant_values=self.pad_token_id)

        return tf.reshape(tokens, [-1, self.sequence_length])  

My current model :

def get_model(backbone, max_len, tokenizer):
    """
        backbone = transformer model
    """
    padding_idx = tokenizer.pad_token_id
    input_str = tf.keras.layers.Input(shape=(), dtype=tf.string, name = "input_str")
    input_ids = tf.keras.layers.Lambda(lambda x: tokenizer(x))(input_str)
    
    
    #attention_mask = tf.keras.layers.Input(shape=(max_len,), dtype=tf.int32, name = "attention_mask")
    attention_mask = tf.math.not_equal(input_ids, padding_idx)
    predictions = backbone(input_ids, attention_mask=attention_mask)
    outputs = tf.keras.layers.Activation("sigmoid", name="outputs_proba")(predictions)

    model = tf.keras.Model(inputs=input_str, outputs=outputs)
    model.compile(tf.keras.optimizers.Adam(1e-5), loss="binary_crossentropy")
    return model

PS : I am using TF 2.3.1

Start by reproducing the failure in TokenizerTF and get_model with TensorFlow 2.3.1, using tf.saved_model.save as the entry point. Inspect how the tf_text.BertTokenizer resource is tracked through the Lambda layer and Keras model. Done means establishing whether this composition can be saved and documenting or fixing the observed behavior.