tf.data.Dataset.from_generator : value error "not enough values to unpack (expected 2, got 1)" - tensorflow

Question

I am currently trying to do a classification problem using a pre-trained transformer model. I wrote a custom generator using tf.data.Dataset.from_generator method. The model takes two inputs : input_id and attn_mask. While calling model.fit i am getting value error "not enough values to unpack (expected 2, got 1)" The received arguments list shows it got both input_id and attn_mask. Can anyone help me solve this?
import tensorflow.keras as keras
from tensorflow.keras.models import Model
from transformers import TFBertModel,BertConfig
def _input_fn():
x = (train_data.iloc[:,0:512]).to_numpy()
y = (train_data.iloc[:,512:516]).to_numpy()
attn = np.asarray(np.tile(attn_mask,x.shape[0]).reshape(-1,512))
def generator():
for s1, s2, l in zip(x, attn, y):
yield {"input_id": s1, "attn_mask": s2}, l
dataset = tf.data.Dataset.from_generator(generator, output_types=({"input_id": tf.int32, "attn_mask": tf.int32}, tf.int32))
#dataset = dataset.batch(2)
#dataset = dataset.shuffle
return dataset
train_data is the dataframe containing the training data(16000 x 516). Last four columns are one hot encoded labels. Since i am not using the Autotokenizer function, i am passing the attention mask as attn_mask.
my model
bert = 'bert-base-uncased'
config = BertConfig(dropout=0.2, attention_dropout=0.2)
config.output_hidden_states = False
transformer_model = TFBertModel.from_pretrained(bert, config = config)
input_ids_in = tf.keras.layers.Input(shape=(512), name='input_id', dtype='int32')
input_masks_in = tf.keras.layers.Input(shape=(512), name='attn_mask', dtype='int32')
embedding_layer = transformer_model(input_ids_in, attention_mask=input_masks_in)[0]
#cls_token = embedding_layer[:,0,:]
#X = tf.keras.layers.BatchNormalization()(cls_token)
X = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(50, return_sequences=True, dropout=0.1, recurrent_dropout=0.1))(embedding_layer)
X = tf.keras.layers.GlobalMaxPool1D()(X)
#X = tf.keras.layers.BatchNormalization()(X)
X = tf.keras.layers.Dense(50, activation='relu')(X)
X = tf.keras.layers.Dropout(0.2)(X)
X = tf.keras.layers.Dense(4, activation='softmax')(X)
model = tf.keras.Model(inputs=[input_ids_in, input_masks_in], outputs = X)
for layer in model.layers[:3]:
layer.trainable = False
optimizer = tf.keras.optimizers.Adam(0.001, beta_1=0.9, beta_2=0.98,
epsilon=1e-9)
model.compile(optimizer=optimizer, loss='categorical_crossentropy', metrics=['categorical_accuracy'])
epochs = 1
batch_size =2
history = model.fit(_input_fn(), epochs= epochs, batch_size= batch_size, verbose=2)
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
/tmp/ipykernel_16908/300834086.py in <module>
2 batch_size =2
3 #history = model.fit(trainDataGenerator(batch_size), epochs= epochs, validation_data=valDataGenerator(batch_size), verbose=2) #
----> 4 history = model.fit(_input_fn(), epochs= epochs, batch_size= batch_size, verbose=2) #validation_data=val_ds,
~/.local/lib/python3.8/site-packages/keras/utils/traceback_utils.py in error_handler(*args, **kwargs)
65 except Exception as e: # pylint: disable=broad-except
66 filtered_tb = _process_traceback_frames(e.__traceback__)
---> 67 raise e.with_traceback(filtered_tb) from None
68 finally:
69 del filtered_tb
~/.local/lib/python3.8/site-packages/transformers/models/bert/modeling_tf_bert.py in call(self, input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, past_key_values, use_cache, output_attentions, output_hidden_states, return_dict, training, **kwargs)
1124 kwargs_call=kwargs,
1125 )
-> 1126 outputs = self.bert(
1127 input_ids=inputs["input_ids"],
1128 attention_mask=inputs["attention_mask"],
~/.local/lib/python3.8/site-packages/transformers/models/bert/modeling_tf_bert.py in call(self, input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, past_key_values, use_cache, output_attentions, output_hidden_states, return_dict, training, **kwargs)
771 raise ValueError("You have to specify either input_ids or inputs_embeds")
772
--> 773 batch_size, seq_length = input_shape
774
775 if inputs["past_key_values"] is None:
ValueError: Exception encountered when calling layer "bert" (type TFBertMainLayer).
not enough values to unpack (expected 2, got 1)
Call arguments received:
• input_ids=tf.Tensor(shape=(512,), dtype=int32)
• attention_mask=tf.Tensor(shape=(512,), dtype=int32)
• token_type_ids=None
• position_ids=None
• head_mask=None
• inputs_embeds=None
• encoder_hidden_states=None
• encoder_attention_mask=None
• past_key_values=None
• use_cache=True
• output_attentions=False
• output_hidden_states=False
• return_dict=True
• training=True
• kwargs=<class 'inspect._empty'>
Edit:
Adding the output of calling _input_fn()
<FlatMapDataset shapes: ({input_id: <unknown>, attn_mask: <unknown>}, <unknown>), types: ({input_id: tf.int32, attn_mask: tf.int32}, tf.int32)>

Answer 1

I solved this error by batching my tf.data.Dataset. This gave the TensorSpec in my dataset a shape that had two values to unpack ->
TensorSpec(shape=(16, 200)...
This is what the error refers to.
Solution
print(train_ds) #Before Batching
new_train_ds = train_ds.batch(16, drop_remainder=True)
print(new_train_ds) #After Batching
# Before Batching
<MapDataset element_spec=({'input_ids': TensorSpec(shape=(200,),
dtype=tf.float64, name=None), 'attention_mask': TensorSpec(shape=
(200,), dtype=tf.float64, name=None)}, TensorSpec(shape=(11,),
dtype=tf.float64, name=None))>
# After Batching
<BatchDataset element_spec=({'input_ids': TensorSpec(shape=(16, 200),
dtype=tf.float64, name=None), 'attention_mask': TensorSpec(shape=(16,
200), dtype=tf.float64, name=None)}, TensorSpec(shape=(16, 11),
dtype=tf.float64, name=None))>