Roberta_model.py

pre_trained_model = TFRobertaForMultipleChoice.from_pretrained('roberta-base')

model_input_ids = Input(shape=(5,128,), name='input_tokens', dtype='int32')
masks_input = Input(shape=(5,128,), name='attention_mask', dtype='int32')

x =  {'input_ids':model_input_ids,
            'attention_mask':masks_input}

x = pre_trained_model(x)['logits']
outputs = Dense(5, activation='softmax')(x)

model = Model(inputs=[model_input_ids, masks_input], outputs=outputs)

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-5),
              loss=tf.keras.losses.CategoricalCrossentropy(),
              metrics=['accuracy'])
valid_y = tf.keras.utils.to_categorical(easy_dev_labels, num_classes=5)
model.fit(x=[easy_train_input_ids, easy_train_attention_mask], y=tf.keras.utils.to_categorical(easy_train_labels),
          batch_size=8, epochs=5,
          validation_data=([easy_dev_input_ids, easy_dev_attention_mask], valid_y),
          validation_batch_size=8)

ok