title: TFBertForSequenceClassification for multi label classification
body: <p>I am trying to fine-tune a bert model for multi-label classification.</p>
  <p>here is how my data looks like. I have put the entire code <a href="https://colab.research.google.com/drive/1QT2uEp8LtFyKjBixszasRwU8ZYg159r5?usp=sharing" rel="nofollow noreferrer">on this colab notebook</a></p>
  <pre><code>({'input_ids': &lt;tf.Tensor: shape=(128,), dtype=int32, numpy=
  array([    2,  8318,  1379,  7892,  2791, 20630,     1,     4,     0,
              0,     0,     0,     0,     0,     0,     0,     0,     0,
              0,     0,     0,     0,     0,     0,     0,     0,     0,
              0,     0,     0,     0,     0,     0,     0,     0,     0,
              0,     0,     0,     0,     0,     0,     0,     0,     0,
              0,     0,     0,     0,     0,     0,     0,     0,     0,
              0,     0,     0,     0,     0,     0,     0,     0,     0,
              0,     0,     0,     0,     0,     0,     0,     0,     0,
              0,     0,     0,     0,     0,     0,     0,     0,     0,
              0,     0,     0,     0,     0,     0,     0,     0,     0,
              0,     0,     0,     0,     0,     0,     0,     0,     0,
              0,     0,     0,     0,     0,     0,     0,     0,     0,
              0,     0,     0,     0,     0,     0,     0,     0,     0,
              0,     0,     0,     0,     0,     0,     0,     0,     0,
              0,     0], dtype=int32)&gt;, 
  'attention_mask': &lt;tf.Tensor: shape=(128,), dtype=int32, numpy=
  array([1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype=int32)&gt;}, 
  &lt;tf.Tensor: shape=(7,), dtype=int64, numpy=array([1, 0, 0, 0, 0, 0, 0])&gt;)
  </code></pre>
  <p>The first element is the id,
  The second element corresponds to the attention_masks
  the third one are the labels - here I have 7 lables.</p>
  <p><strong>First effort:</strong></p>
  <pre><code>        MODEL_NAME_OR_PATH = 'HooshvareLab/bert-fa-base-uncased'
          NUM_LABELS = 7
          
          from transformers import TFBertForSequenceClassification, BertConfig
          model = TFBertForSequenceClassification.from_pretrained(
              MODEL_NAME_OR_PATH, 
              config=BertConfig.from_pretrained(MODEL_NAME_OR_PATH, num_labels=NUM_LABELS, problem_type=&quot;multi_label_classification&quot;)
              )
          
          loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
          model.compile(optimizer='adam', loss=loss, metrics=['accuracy'])
          history = model.fit(train_dataset, epochs=1, steps_per_epoch=115, validation_data=valid_dataset, validation_steps=7)
  </code></pre>
  <pre><code>which ends up with the following error
  </code></pre>
  <pre><code>    InvalidArgumentError                      Traceback (most recent call last)
      
      &lt;ipython-input-48-4408a1f17fbe&gt; in &lt;module&gt;()
           10 loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
           11 model.compile(optimizer='adam', loss=loss, metrics=['accuracy'])
      ---&gt; 12 history = model.fit(train_dataset, epochs=1, steps_per_epoch=115, validation_data=valid_dataset, validation_steps=7)
           13 
           14 
      
      1 frames
      
      /usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
           53     ctx.ensure_initialized()
           54     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
      ---&gt; 55                                         inputs, attrs, num_outputs)
           56   except core._NotOkStatusException as e:
           57     if name is not None:
      
      InvalidArgumentError: Graph execution error:
      
      Detected at node 'Equal' defined at (most recent call last):
          File &quot;/usr/lib/python3.7/runpy.py&quot;, line 193, in _run_module_as_main
            &quot;__main__&quot;, mod_spec)
          File &quot;/usr/lib/python3.7/runpy.py&quot;, line 85, in _run_code
            exec(code, run_globals)
          File &quot;/usr/local/lib/python3.7/dist-packages/ipykernel_launcher.py&quot;, line 16, in &lt;module&gt;
            app.launch_new_instance()
          File &quot;/usr/local/lib/python3.7/dist-packages/traitlets/config/application.py&quot;, line 846, in launch_instance
            app.start()
          File &quot;/usr/local/lib/python3.7/dist-packages/ipykernel/kernelapp.py&quot;, line 499, in start
            self.io_loop.start()
          File &quot;/usr/local/lib/python3.7/dist-packages/tornado/platform/asyncio.py&quot;, line 132, in start
            self.asyncio_loop.run_forever()
          File &quot;/usr/lib/python3.7/asyncio/base_events.py&quot;, line 541, in run_forever
            self._run_once()
          File &quot;/usr/lib/python3.7/asyncio/base_events.py&quot;, line 1786, in _run_once
            handle._run()
          File &quot;/usr/lib/python3.7/asyncio/events.py&quot;, line 88, in _run
            self._context.run(self._callback, *self._args)
          File &quot;/usr/local/lib/python3.7/dist-packages/tornado/platform/asyncio.py&quot;, line 122, in _handle_events
            handler_func(fileobj, events)
          File &quot;/usr/local/lib/python3.7/dist-packages/tornado/stack_context.py&quot;, line 300, in null_wrapper
            return fn(*args, **kwargs)
          File &quot;/usr/local/lib/python3.7/dist-packages/zmq/eventloop/zmqstream.py&quot;, line 577, in _handle_events
            self._handle_recv()
          File &quot;/usr/local/lib/python3.7/dist-packages/zmq/eventloop/zmqstream.py&quot;, line 606, in _handle_recv
            self._run_callback(callback, msg)
          File &quot;/usr/local/lib/python3.7/dist-packages/zmq/eventloop/zmqstream.py&quot;, line 556, in _run_callback
            callback(*args, **kwargs)
          File &quot;/usr/local/lib/python3.7/dist-packages/tornado/stack_context.py&quot;, line 300, in null_wrapper
            return fn(*args, **kwargs)
          File &quot;/usr/local/lib/python3.7/dist-packages/ipykernel/kernelbase.py&quot;, line 283, in dispatcher
            return self.dispatch_shell(stream, msg)
          File &quot;/usr/local/lib/python3.7/dist-packages/ipykernel/kernelbase.py&quot;, line 233, in dispatch_shell
            handler(stream, idents, msg)
          File &quot;/usr/local/lib/python3.7/dist-packages/ipykernel/kernelbase.py&quot;, line 399, in execute_request
            user_expressions, allow_stdin)
          File &quot;/usr/local/lib/python3.7/dist-packages/ipykernel/ipkernel.py&quot;, line 208, in do_execute
            res = shell.run_cell(code, store_history=store_history, silent=silent)
          File &quot;/usr/local/lib/python3.7/dist-packages/ipykernel/zmqshell.py&quot;, line 537, in run_cell
            return super(ZMQInteractiveShell, self).run_cell(*args, **kwargs)
          File &quot;/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py&quot;, line 2718, in run_cell
            interactivity=interactivity, compiler=compiler, result=result)
          File &quot;/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py&quot;, line 2822, in run_ast_nodes
            if self.run_code(code, result):
          File &quot;/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py&quot;, line 2882, in run_code
            exec(code_obj, self.user_global_ns, self.user_ns)
          File &quot;&lt;ipython-input-48-4408a1f17fbe&gt;&quot;, line 12, in &lt;module&gt;
            history = model.fit(train_dataset, epochs=1, steps_per_epoch=115, validation_data=valid_dataset, validation_steps=7)
          File &quot;/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py&quot;, line 64, in error_handler
            return fn(*args, **kwargs)
          File &quot;/usr/local/lib/python3.7/dist-packages/keras/engine/training.py&quot;, line 1384, in fit
            tmp_logs = self.train_function(iterator)
          File &quot;/usr/local/lib/python3.7/dist-packages/keras/engine/training.py&quot;, line 1021, in train_function
            return step_function(self, iterator)
          File &quot;/usr/local/lib/python3.7/dist-packages/keras/engine/training.py&quot;, line 1010, in step_function
            outputs = model.distribute_strategy.run(run_step, args=(data,))
          File &quot;/usr/local/lib/python3.7/dist-packages/keras/engine/training.py&quot;, line 1000, in run_step
            outputs = model.train_step(data)
          File &quot;/usr/local/lib/python3.7/dist-packages/transformers/modeling_tf_utils.py&quot;, line 1156, in train_step
            self.compiled_metrics.update_state(y, y_pred, sample_weight)
          File &quot;/usr/local/lib/python3.7/dist-packages/keras/engine/compile_utils.py&quot;, line 459, in update_state
            metric_obj.update_state(y_t, y_p, sample_weight=mask)
          File &quot;/usr/local/lib/python3.7/dist-packages/keras/utils/metrics_utils.py&quot;, line 70, in decorated
            update_op = update_state_fn(*args, **kwargs)
          File &quot;/usr/local/lib/python3.7/dist-packages/keras/metrics.py&quot;, line 178, in update_state_fn
            return ag_update_state(*args, **kwargs)
          File &quot;/usr/local/lib/python3.7/dist-packages/keras/metrics.py&quot;, line 729, in update_state
            matches = ag_fn(y_true, y_pred, **self._fn_kwargs)
          File &quot;/usr/local/lib/python3.7/dist-packages/keras/metrics.py&quot;, line 4086, in sparse_categorical_accuracy
            return tf.cast(tf.equal(y_true, y_pred), backend.floatx())
      Node: 'Equal'
      required broadcastable shapes
           [[{{node Equal}}]] [Op:__inference_train_function_187978]
  </code></pre>
  <p><strong>Second Effort</strong> inspired by <a href="https://github.com/NZbryan/NLP_bert/blob/master/tf2.0_bert_emb_en_MultiLabel.py#L61" rel="nofollow noreferrer">this piece of code</a></p>
  <pre><code>    from transformers import TFBertPreTrainedModel
      from transformers import TFBertMainLayer
      
      
      class TFBertForMultilabelClassification(TFBertPreTrainedModel):
      
          def __init__(self, config, *inputs, **kwargs):
              super(TFBertForMultilabelClassification, self).__init__(config, *inputs, **kwargs)
              self.num_labels = config.num_labels
      
              self.bert = TFBertMainLayer(config, name='bert')
              self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob)
              self.classifier = tf.keras.layers.Dense(config.num_labels,
                                                      kernel_initializer='random_normal', #get_initializer(config.initializer_range),
                                                      name='classifier',
                                                      activation='sigmoid')
      
          def call(self, inputs, **kwargs):
              outputs = self.bert(inputs, **kwargs)
      
              pooled_output = outputs[1]
      
              pooled_output = self.dropout(pooled_output, training=kwargs.get('training', False))
              logits = self.classifier(pooled_output)
      
              outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here
      
              return outputs  # logits, (hidden_states), (attentions)
      
      
      MODEL_NAME_OR_PATH = 'HooshvareLab/bert-fa-base-uncased'
      NUM_LABELS = len(y_train[0])
      
      model = TFBertForMultilabelClassification.from_pretrained(MODEL_NAME_OR_PATH, num_labels=NUM_LABELS)
      
      optimizer = tf.keras.optimizers.Adam(learning_rate=0.001,epsilon=1e-08, clipnorm=1)
      # we do not have one-hot vectors, we can use sparce categorical cross entropy and accuracy
      loss = tf.keras.losses.BinaryCrossentropy()
      metric = tf.keras.metrics.CategoricalAccuracy()
      model.compile(optimizer=optimizer, loss=loss, metrics=['accuracy'])
      history = model.fit(train_dataset, epochs=1, validation_data=valid_dataset)
  </code></pre>
  <p>returns the following error</p>
  <pre><code>    InvalidArgumentError                      Traceback (most recent call last)
      
      &lt;ipython-input-49-8aa1173bef76&gt; in &lt;module&gt;()
            4 metric = tf.keras.metrics.CategoricalAccuracy()
            5 model.compile(optimizer=optimizer, loss=loss, metrics=['accuracy'])
      ----&gt; 6 history = model.fit(train_dataset, epochs=1, validation_data=valid_dataset)
      
      1 frames
      
          /usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
               53     ctx.ensure_initialized()
               54     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
          ---&gt; 55                                         inputs, attrs, num_outputs)
               56   except core._NotOkStatusException as e:
               57     if name is not None:
          
          InvalidArgumentError: Graph execution error:
          
          Detected at node 'Equal' defined at (most recent call last):
              File &quot;/usr/lib/python3.7/runpy.py&quot;, line 193, in _run_module_as_main
                &quot;__main__&quot;, mod_spec)
              File &quot;/usr/lib/python3.7/runpy.py&quot;, line 85, in _run_code
                exec(code, run_globals)
              File &quot;/usr/local/lib/python3.7/dist-packages/ipykernel_launcher.py&quot;, line 16, in &lt;module&gt;
                app.launch_new_instance()
              File &quot;/usr/local/lib/python3.7/dist-packages/traitlets/config/application.py&quot;, line 846, in launch_instance
                app.start()
              File &quot;/usr/local/lib/python3.7/dist-packages/ipykernel/kernelapp.py&quot;, line 499, in start
                self.io_loop.start()
              File &quot;/usr/local/lib/python3.7/dist-packages/tornado/platform/asyncio.py&quot;, line 132, in start
                self.asyncio_loop.run_forever()
              File &quot;/usr/lib/python3.7/asyncio/base_events.py&quot;, line 541, in run_forever
                self._run_once()
              File &quot;/usr/lib/python3.7/asyncio/base_events.py&quot;, line 1786, in _run_once
                handle._run()
              File &quot;/usr/lib/python3.7/asyncio/events.py&quot;, line 88, in _run
                self._context.run(self._callback, *self._args)
              File &quot;/usr/local/lib/python3.7/dist-packages/tornado/platform/asyncio.py&quot;, line 122, in _handle_events
                handler_func(fileobj, events)
              File &quot;/usr/local/lib/python3.7/dist-packages/tornado/stack_context.py&quot;, line 300, in null_wrapper
                return fn(*args, **kwargs)
              File &quot;/usr/local/lib/python3.7/dist-packages/zmq/eventloop/zmqstream.py&quot;, line 577, in _handle_events
                self._handle_recv()
              File &quot;/usr/local/lib/python3.7/dist-packages/zmq/eventloop/zmqstream.py&quot;, line 606, in _handle_recv
                self._run_callback(callback, msg)
              File &quot;/usr/local/lib/python3.7/dist-packages/zmq/eventloop/zmqstream.py&quot;, line 556, in _run_callback
                callback(*args, **kwargs)
              File &quot;/usr/local/lib/python3.7/dist-packages/tornado/stack_context.py&quot;, line 300, in null_wrapper
                return fn(*args, **kwargs)
              File &quot;/usr/local/lib/python3.7/dist-packages/ipykernel/kernelbase.py&quot;, line 283, in dispatcher
                return self.dispatch_shell(stream, msg)
              File &quot;/usr/local/lib/python3.7/dist-packages/ipykernel/kernelbase.py&quot;, line 233, in dispatch_shell
                handler(stream, idents, msg)
              File &quot;/usr/local/lib/python3.7/dist-packages/ipykernel/kernelbase.py&quot;, line 399, in execute_request
                user_expressions, allow_stdin)
              File &quot;/usr/local/lib/python3.7/dist-packages/ipykernel/ipkernel.py&quot;, line 208, in do_execute
                res = shell.run_cell(code, store_history=store_history, silent=silent)
              File &quot;/usr/local/lib/python3.7/dist-packages/ipykernel/zmqshell.py&quot;, line 537, in run_cell
                return super(ZMQInteractiveShell, self).run_cell(*args, **kwargs)
              File &quot;/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py&quot;, line 2718, in run_cell
                interactivity=interactivity, compiler=compiler, result=result)
              File &quot;/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py&quot;, line 2822, in run_ast_nodes
                if self.run_code(code, result):
              File &quot;/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py&quot;, line 2882, in run_code
                exec(code_obj, self.user_global_ns, self.user_ns)
              File &quot;&lt;ipython-input-49-8aa1173bef76&gt;&quot;, line 6, in &lt;module&gt;
                history = model.fit(train_dataset, epochs=1, validation_data=valid_dataset)
              File &quot;/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py&quot;, line 64, in error_handler
                return fn(*args, **kwargs)
              File &quot;/usr/local/lib/python3.7/dist-packages/keras/engine/training.py&quot;, line 1384, in fit
                tmp_logs = self.train_function(iterator)
              File &quot;/usr/local/lib/python3.7/dist-packages/keras/engine/training.py&quot;, line 1021, in train_function
                return step_function(self, iterator)
              File &quot;/usr/local/lib/python3.7/dist-packages/keras/engine/training.py&quot;, line 1010, in step_function
                outputs = model.distribute_strategy.run(run_step, args=(data,))
              File &quot;/usr/local/lib/python3.7/dist-packages/keras/engine/training.py&quot;, line 1000, in run_step
                outputs = model.train_step(data)
              File &quot;/usr/local/lib/python3.7/dist-packages/transformers/modeling_tf_utils.py&quot;, line 1156, in train_step
                self.compiled_metrics.update_state(y, y_pred, sample_weight)
              File &quot;/usr/local/lib/python3.7/dist-packages/keras/engine/compile_utils.py&quot;, line 459, in update_state
                metric_obj.update_state(y_t, y_p, sample_weight=mask)
              File &quot;/usr/local/lib/python3.7/dist-packages/keras/utils/metrics_utils.py&quot;, line 70, in decorated
                update_op = update_state_fn(*args, **kwargs)
              File &quot;/usr/local/lib/python3.7/dist-packages/keras/metrics.py&quot;, line 178, in update_state_fn
                return ag_update_state(*args, **kwargs)
              File &quot;/usr/local/lib/python3.7/dist-packages/keras/metrics.py&quot;, line 729, in update_state
                matches = ag_fn(y_true, y_pred, **self._fn_kwargs)
              File &quot;/usr/local/lib/python3.7/dist-packages/keras/metrics.py&quot;, line 4086, in sparse_categorical_accuracy
                return tf.cast(tf.equal(y_true, y_pred), backend.floatx())
          Node: 'Equal'
          required broadcastable shapes
               [[{{node Equal}}]] [Op:__inference_train_function_214932]
  </code></pre>
  <p>I hope I believe given major changes both in tf2 and (TF-based) huggingface transformers</p>
  <p><strong>UPDATE</strong></p>
  <p>Here is the entire code with a dummy dataset; the whole thing is also available on <a href="https://colab.research.google.com/drive/1QT2uEp8LtFyKjBixszasRwU8ZYg159r5?usp=sharing" rel="nofollow noreferrer">this colab notebook</a></p>
  <p>load the libraries</p>
  <pre><code>import os
  
  import pandas as pd
  import numpy as np
  
  from transformers import TFBertPreTrainedModel
  from transformers import TFBertMainLayer
  from keras.preprocessing.sequence import pad_sequences
  from tqdm import tqdm
  from transformers import BertTokenizer
  import tensorflow as tf
  </code></pre>
  <p>make a dummy data</p>
  <pre><code>x_train = ['هان از وقتی که زفتم مدرسه',
   'معاویه برادر شمر',
   'وقتی که از پنجره سرشرو میاره بیرون دالی میکنه',
   'هر دو سحرند این کجا و آن کجا']
  
  y_train = [[1, 0, 0, 0, 0, 0, 0], [0, 1, 0, 0, 0, 0, 1], [1, 0, 0, 0, 0, 0, 0], [1, 0, 0, 0, 0, 0, 0]]
  
  x_test, x_valid = x_train, x_train
  y_test, y_valid = y_train, y_train
  </code></pre>
  <p>add the configs</p>
  <pre><code># general config
  MAX_LEN = 128
  batch_size = 32
  TRAIN_BATCH_SIZE = batch_size
  VALID_BATCH_SIZE = batch_size
  TEST_BATCH_SIZE = batch_size
  
  EPOCHS = 3
  EEVERY_EPOCH = 1000
  LEARNING_RATE = 2e-5
  CLIP = 0.0
  </code></pre>
  <p>make the data huggingface friendly</p>
  <pre><code>MODEL_NAME_OR_PATH = 'HooshvareLab/bert-fa-base-uncased'
  tokenizer = BertTokenizer.from_pretrained(MODEL_NAME_OR_PATH)
  
  
  MAX_LEN = 128
  def tokenize_sentences(sentences, tokenizer, max_seq_len = 128):
      tokenized_sentences = []
      for sentence in tqdm(sentences):
          tokenized_sentence = tokenizer.encode(
                              sentence,                  # Sentence to encode.
                              add_special_tokens = True, # Add '[CLS]' and '[SEP]'
                              max_length = max_seq_len,  # Truncate all sentences.
                      )
          
          tokenized_sentences.append(tokenized_sentence)
      return tokenized_sentences
  
      def create_attention_masks(tokenized_and_padded_sentences):
          attention_masks = []
          for sentence in tokenized_and_padded_sentences:
              att_mask = [int(token_id &gt; 0) for token_id in sentence]
              attention_masks.append(att_mask)
          return np.asarray(attention_masks)
      
      train_ids = tokenize_sentences(x_train, tokenizer, max_seq_len = 128)
      train_ids = pad_sequences(train_ids, maxlen=MAX_LEN, dtype=&quot;long&quot;, value=0, truncating=&quot;post&quot;, padding=&quot;post&quot;)
      train_masks = create_attention_masks(train_ids)
      
      valid_ids = tokenize_sentences(x_valid, tokenizer, max_seq_len = 128)
      valid_ids = pad_sequences(valid_ids, maxlen=MAX_LEN, dtype=&quot;long&quot;, value=0, truncating=&quot;post&quot;, padding=&quot;post&quot;)
      valid_masks = create_attention_masks(valid_ids)
      
      test_ids = tokenize_sentences(x_test, tokenizer, max_seq_len = 128)
      test_ids = pad_sequences(test_ids, maxlen=MAX_LEN, dtype=&quot;long&quot;, value=0, truncating=&quot;post&quot;, padding=&quot;post&quot;)
      test_masks = create_attention_masks(test_ids)
  </code></pre>
  <p>create the datasets</p>
  <pre><code>def create_dataset(ids, masks, labels):
       def gen():
           for i in range(len(ids)):
               yield (
                   {
                       &quot;input_ids&quot;: ids[i],
                       &quot;attention_mask&quot;: masks[i]
                   },
                   labels[i],
               )
       return tf.data.Dataset.from_generator(
           gen,
           ({&quot;input_ids&quot;: tf.int32, &quot;attention_mask&quot;: tf.int32}, tf.int64),
           (
               {
                   &quot;input_ids&quot;: tf.TensorShape([None]),
                   &quot;attention_mask&quot;: tf.TensorShape([None])
               },
               tf.TensorShape([None]),
           ),
       )
  
  train_dataset = create_dataset(train_ids, train_masks, y_train)
  valid_dataset = create_dataset(valid_ids, valid_masks, y_valid)
  test_dataset = create_dataset(test_ids, test_masks, y_test)
  </code></pre>
  <p>that is how the data looks like</p>
  <pre><code>for item in train_dataset.take(1):
    print(item)
  </code></pre>
  <p><a href="https://i.stack.imgur.com/rmArl.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/rmArl.png" alt="enter image description here" /></a></p>
  <p>Approach 1</p>
  <pre><code>class TFBertForMultilabelClassification(TFBertPreTrainedModel):
  
      def __init__(self, config, *inputs, **kwargs):
          super(TFBertForMultilabelClassification, self).__init__(config, *inputs, **kwargs)
          self.num_labels = config.num_labels
  
          self.bert = TFBertMainLayer(config, name='bert')
          self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob)
          self.classifier = tf.keras.layers.Dense(config.num_labels,
                                                  kernel_initializer='random_normal', #get_initializer(config.initializer_range),
                                                  name='classifier',
                                                  activation='sigmoid')
  
      def call(self, inputs, **kwargs):
          outputs = self.bert(inputs, **kwargs)
  
          pooled_output = outputs[1]
  
          pooled_output = self.dropout(pooled_output, training=kwargs.get('training', False))
          logits = self.classifier(pooled_output)
  
          outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here
  
          return outputs  # logits, (hidden_states), (attentions)
  
  NUM_LABELS = len(y_train[0])
  
  model = TFBertForMultilabelClassification.from_pretrained(MODEL_NAME_OR_PATH, num_labels=NUM_LABELS)
  
  
  optimizer = tf.keras.optimizers.Adam(learning_rate=0.001,epsilon=1e-08, clipnorm=1)
  # we do not have one-hot vectors, we can use sparce categorical cross entropy and accuracy
  loss = tf.keras.losses.BinaryCrossentropy()
  metric = tf.keras.metrics.CategoricalAccuracy()
  model.compile(optimizer=optimizer, loss=loss, metrics=['accuracy'])
  history = model.fit(train_dataset, epochs=1, validation_data=valid_dataset)
  </code></pre>
  <p>with an error</p>
  <pre><code>---------------------------------------------------------------------------
  
  AttributeError                            Traceback (most recent call last)
  
  &lt;ipython-input-36-8aa1173bef76&gt; in &lt;module&gt;()
        4 metric = tf.keras.metrics.CategoricalAccuracy()
        5 model.compile(optimizer=optimizer, loss=loss, metrics=['accuracy'])
  ----&gt; 6 history = model.fit(train_dataset, epochs=1, validation_data=valid_dataset)
  
  1 frames
  
  /usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py in autograph_handler(*args, **kwargs)
     1145           except Exception as e:  # pylint:disable=broad-except
     1146             if hasattr(e, &quot;ag_error_metadata&quot;):
  -&gt; 1147               raise e.ag_error_metadata.to_exception(e)
     1148             else:
     1149               raise
  
  AttributeError: in user code:
  
      File &quot;/usr/local/lib/python3.7/dist-packages/keras/engine/training.py&quot;, line 1021, in train_function  *
          return step_function(self, iterator)
      File &quot;/usr/local/lib/python3.7/dist-packages/keras/engine/training.py&quot;, line 1010, in step_function  **
          outputs = model.distribute_strategy.run(run_step, args=(data,))
      File &quot;/usr/local/lib/python3.7/dist-packages/keras/engine/training.py&quot;, line 1000, in run_step  **
          outputs = model.train_step(data)
      File &quot;/usr/local/lib/python3.7/dist-packages/transformers/modeling_tf_utils.py&quot;, line 1145, in train_step
          if list(y_pred.keys())[0] == &quot;loss&quot;:
  
      AttributeError: 'tuple' object has no attribute 'keys'
  </code></pre>
  <p><strong>Approach 2:</strong></p>
  <pre><code>MODEL_NAME_OR_PATH = 'HooshvareLab/bert-fa-base-uncased'
  NUM_LABELS = 7
  
  from transformers import TFBertForSequenceClassification, BertConfig
  model = TFBertForSequenceClassification.from_pretrained(
      MODEL_NAME_OR_PATH, 
      config=BertConfig.from_pretrained(MODEL_NAME_OR_PATH, num_labels=NUM_LABELS, problem_type=&quot;multi_label_classification&quot;)
      )
  
  loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
  model.compile(optimizer='adam', loss=loss, metrics=['accuracy'])
  history = model.fit(train_dataset, epochs=1, steps_per_epoch=115, validation_data=valid_dataset, validation_steps=7)
  </code></pre>
  <p>and the error</p>
  <pre><code>      outputs = model.distribute_strategy.run(run_step, args=(data,))
      File &quot;/usr/local/lib/python3.7/dist-packages/keras/engine/training.py&quot;, line 1000, in run_step
        outputs = model.train_step(data)
      File &quot;/usr/local/lib/python3.7/dist-packages/transformers/modeling_tf_utils.py&quot;, line 1151, in train_step
        loss = self.compiled_loss(y, y_pred, sample_weight, regularization_losses=self.losses)
      File &quot;/usr/local/lib/python3.7/dist-packages/keras/engine/compile_utils.py&quot;, line 201, in __call__
        loss_value = loss_obj(y_t, y_p, sample_weight=sw)
      File &quot;/usr/local/lib/python3.7/dist-packages/keras/losses.py&quot;, line 141, in __call__
        losses = call_fn(y_true, y_pred)
      File &quot;/usr/local/lib/python3.7/dist-packages/keras/losses.py&quot;, line 245, in call
        return ag_fn(y_true, y_pred, **self._fn_kwargs)
      File &quot;/usr/local/lib/python3.7/dist-packages/keras/losses.py&quot;, line 1863, in sparse_categorical_crossentropy
        y_true, y_pred, from_logits=from_logits, axis=axis)
      File &quot;/usr/local/lib/python3.7/dist-packages/keras/backend.py&quot;, line 5203, in sparse_categorical_crossentropy
        labels=target, logits=output)
  Node: 'sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits'
  logits and labels must have the same first dimension, got logits shape [128,7] and labels shape [7]
       [[{{node sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits}}]] [Op:__inference_train_function_67923]
  </code></pre>