How to increase the rank (ndim) of input of BERT keras hub layer for learning-to-rank

user_14913631 0
Viewed 265

I am trying to implement a learning-to-rank model using a pre-trained BERT available on tensorflow hub. I am using a variation of ListNet loss function, which requires each training instance to be a list of several ranked documents in relation to a query. I need the model to be able to accept data in a shape (batch_size, list_size, sentence_length), where the model loops over the 'list_size' axis in each training instance, returns the ranks and passes them to the loss function. In a simple model that only consists of dense layers, this is easily done by augmenting the dimensions of the input layer. For example:

from tensorflow.keras.layers import Dense, Input
from tensorflow.keras import Model

input = Input([6,10])
x = Dense(20,activation='relu')(input)
output = Dense(1, activation='sigmoid')(x)
model = Model(inputs=input, outputs=output)

...now the model will perform 6 forward passes over vectors of length 10 before calculating the loss and updating gradients.

I am trying to do the same with the BERT model and its preprocessing layer:

import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_text as text

bert_preprocess_model = hub.KerasLayer('https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-4_H-512_A-8/1')
bert_model = hub.KerasLayer('https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3')
    
text_input = tf.keras.layers.Input(shape=(), dtype=tf.string, name='text')
processed_input = bert_preprocess_model(text_input)
output = bert_model(processed_input)
model = tf.keras.Model(text_input, output)

But when I try to change the shape of 'text_input' to, say, (6), or meddle with it in any way really, it always results in the same type of error:

 ValueError: Could not find matching function to call loaded from the SavedModel. Got:
      Positional arguments (3 total):
        * Tensor("inputs:0", shape=(None, 6), dtype=string)
        * False
        * None
      Keyword arguments: {}
    
    Expected these arguments to match one of the following 4 option(s):
    
    Option 1:
      Positional arguments (3 total):
        * TensorSpec(shape=(None,), dtype=tf.string, name='sentences')
        * False
        * None
      Keyword arguments: {}
     ....

As per https://www.tensorflow.org/hub/api_docs/python/hub/KerasLayer, it seems like you can configure the input shape of hub.KerasLayer via tf.keras.layers.InputSpec. In my case, I guess it would be something like this:

bert_preprocess_model.input_spec = tf.keras.layers.InputSpec(ndim=2)
bert_model.input_spec = tf.keras.layers.InputSpec(ndim=2)

When I run the above code, the attributes indeed get changed, but when trying to build the model, the same exact error appears.

Is there any way to easily resolve this without the necessity to create a custom training loop?

tensorflow-hub bert-language-model keras information-retrieval python
1 Answers

Suppose you have a batch of B examples, each with exactly N text strings, which makes a 2-dimensional Tensor of shape [B, N]. With tf.reshape(), you can turn that into a 1-dimensional tensor of shape [B*N], send it through BERT (which preserves the order of inputs) and then reshape it back to [B,N]. (There's also tf.keras.layers.Reshape, but that hides the batch dimension from you.)

If it's not exactly N text strings each time, you'll have to do some bookkeeping on the side (e.g., store inputs in a tf.RaggedTensor, run BERT on its .values, and construct a new RaggedTensor with the same .row_splits from the result.)

user_70612490
Related