title: How to preprocess a dataset for BERT model implemented in Tensorflow 2.x?
body: <h2>Overview</h2>
  <p>I have a dataset made for classification problem. There are two columns one is <code>sentences</code> and the other is <code>labels</code> (total: 10 labels). I'm trying to convert this dataset to implement it in a BERT model made for classification and that is implemented in Tensorflow 2.x. However, I can't preprocess correctly the dataset to make a <code>PrefetchDataset</code> used as input.</p>
  <h2>What I did?</h2>
  <ul>
  <li>Dataframe is balanced and shuffled (every label have 18708 data)</li>
  <li>Dataframe shape: (187080, 2)</li>
  <li><code>from sklearn.model_selection import train_test_split</code> was used to split the dataframe</li>
  <li>80% train data, 20% test data</li>
  </ul>
  <h3>Training data:</h3>
  <p><strong>X_train</strong></p>
  <pre><code>array(['i hate megavideo  stupid time limits',
         'wow this class got wild quick  functions are a butt',
         'got in trouble no cell phone or computer for a you later twitter',
         ...,
         'we lied down around am rose a few hours later party still going lt',
         'i wanna miley cyrus on brazil  i love u my diva miley rocks',
         'i know i hate it i want my dj danger bck'], dtype=object)
  </code></pre>
  <p><strong>y_train</strong></p>
  <pre class="lang-py prettyprint-override"><code>array(['unfriendly', 'unfriendly', 'unfriendly', ..., 'pos_hp',
         'friendly', 'friendly'], dtype=object)
  </code></pre>
  <p><strong>BERT preprocessing Xy_dataset</strong></p>
  <pre class="lang-py prettyprint-override"><code>AUTOTUNE = tf.data.AUTOTUNE # autotune the buffer_size: optional = 1
  
  train_Xy_slices = tf.data.Dataset.from_tensor_slices(tensors=(X_train, y_train))
  dataset_train_Xy = train_Xy_slices.batch(batch_size=32)
  </code></pre>
  <p><strong>output</strong></p>
  <pre class="lang-py prettyprint-override"><code>dataset_train_Xy
  &lt;PrefetchDataset shapes: ((None,), (None,)), types: (tf.string, tf.string)&gt;
  
  
  for i in dataset_train_Xy:
      print(i)
  (
  &lt;tf.Tensor: shape=(32,), dtype=string, numpy=
  array([b'some of us had to work al day',
         ...
         b'feels claudia cazacus free falling feat audrey gallagher amp thomas bronzwaers look ahead are the best trance offerings this summer'], dtype=object)&gt;,
   
  &lt;tf.Tensor: shape=(32,), dtype=string, numpy=
  array([b'interested', b'uninterested', b'happy', b'friendly', b'neg_hp',
         ...
         b'friendly', b'insecure', b'pos_hp', b'interested', b'happy'],
        dtype=object)&gt;
  )
  </code></pre>
  <h2>Expected output (example)</h2>
  <pre class="lang-py prettyprint-override"><code>dataset_train_Xy
  &lt;PrefetchDataset shapes: ({input_word_ids: (None, 128), input_mask: (None, 128), input_type_ids: (None, 128)}, (None,)), types: ({input_word_ids: tf.int32, input_mask: tf.int32, input_type_ids: tf.int32}, tf.int64)&gt;
  </code></pre>
  <p><a href="https://i.stack.imgur.com/iqmu6.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/iqmu6.png" alt="enter image description here" /></a></p>
  <h2>Observations/problem:</h2>
  <p>I know I need to tokenize <code>X_train</code> and <code>y_train</code>, but when I tried to tokenize had an error:</p>
  <pre class="lang-py prettyprint-override"><code>AUTOTUNE = tf.data.AUTOTUNE # autotune the buffer_size: optional = 1
  
  train_Xy_slices = tf.data.Dataset.from_tensor_slices(tensors=(X_train, y_train))
  dataset_train_Xy = train_Xy_slices.batch(batch_size=batch_size) # 32
  
  print(type(dataset_train_Xy))
  
  # Tokenize the text to word pieces.
  bert_preprocess = hub.load(tfhub_handle_preprocess)
  tokenizer = hub.KerasLayer(bert_preprocess.tokenize, name='tokenizer')
  
  dataset_train_Xy = dataset_train_Xy.map(lambda ex: (tokenizer(ex), ex[1])) #    print(i[1]) # correspond to labels
  dataset_train_Xy = dataset_train_Xy.prefetch(buffer_size=AUTOTUNE)
  </code></pre>
  <h3>Traceback</h3>
  <pre class="lang-py prettyprint-override"><code>&lt;class 'tensorflow.python.data.ops.dataset_ops.BatchDataset'&gt;
  ---------------------------------------------------------------------------
  TypeError                                 Traceback (most recent call last)
  &lt;ipython-input-69-8e486f7b671b&gt; in &lt;module&gt;()
       14 tokenizer = hub.KerasLayer(bert_preprocess.tokenize, name='tokenizer')
       15 
  ---&gt; 16 dataset_train_Xy = dataset_train_Xy.map(lambda ex: (tokenizer(ex), ex[1])) #    print(i[1]) #labels
       17 dataset_train_Xy = dataset_train_Xy.prefetch(buffer_size=AUTOTUNE)
  
  10 frames
  /usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/impl/api.py in wrapper(*args, **kwargs)
      668       except Exception as e:  # pylint:disable=broad-except
      669         if hasattr(e, 'ag_error_metadata'):
  --&gt; 670           raise e.ag_error_metadata.to_exception(e)
      671         else:
      672           raise
  
  TypeError: in user code:
  
  
      TypeError: &lt;lambda&gt;() takes 1 positional argument but 2 were given
  </code></pre>