def call(self, inputs, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, training=False):
    if isinstance(inputs, (tuple, list)):
        input_ids = inputs[0]
        attention_mask = (inputs[1] if (len(inputs) > 1) else attention_mask)
        token_type_ids = (inputs[2] if (len(inputs) > 2) else token_type_ids)
        position_ids = (inputs[3] if (len(inputs) > 3) else position_ids)
        head_mask = (inputs[4] if (len(inputs) > 4) else head_mask)
        inputs_embeds = (inputs[5] if (len(inputs) > 5) else inputs_embeds)
        assert (len(inputs) <= 6), 'Too many inputs.'
    elif isinstance(inputs, dict):
        input_ids = inputs.get('input_ids')
        attention_mask = inputs.get('attention_mask', attention_mask)
        token_type_ids = inputs.get('token_type_ids', token_type_ids)
        position_ids = inputs.get('position_ids', position_ids)
        head_mask = inputs.get('head_mask', head_mask)
        inputs_embeds = inputs.get('inputs_embeds', inputs_embeds)
        assert (len(inputs) <= 6), 'Too many inputs.'
    else:
        input_ids = inputs
    if ((input_ids is not None) and (inputs_embeds is not None)):
        raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time')
    elif (input_ids is not None):
        input_shape = shape_list(input_ids)
        input_ids = tf.reshape(input_ids, [(- 1), input_shape[(- 1)]])
    elif (inputs_embeds is not None):
        input_shape = shape_list(inputs_embeds)[:(- 1)]
    else:
        raise ValueError('You have to specify either input_ids or inputs_embeds')
    if (position_ids is None):
        position_ids = tf.range(input_shape[(- 1)], dtype=tf.int32)[tf.newaxis, :]
    if (attention_mask is not None):
        attention_mask = attention_mask[:, tf.newaxis, tf.newaxis, :]
        attention_mask = tf.cast(attention_mask, tf.float32)
        attention_mask = ((1.0 - attention_mask) * (- 10000.0))
    else:
        attention_mask = None
    if (not (head_mask is None)):
        raise NotImplementedError
    else:
        head_mask = ([None] * self.num_hidden_layers)
    position_ids = tf.reshape(position_ids, [(- 1), shape_list(position_ids)[(- 1)]])
    if (inputs_embeds is None):
        inputs_embeds = self.tokens_embed(input_ids, mode='embedding')
    position_embeds = self.positions_embed(position_ids)
    if (token_type_ids is not None):
        token_type_ids = tf.reshape(token_type_ids, [(- 1), shape_list(token_type_ids)[(- 1)]])
        token_type_embeds = self.tokens_embed(token_type_ids, mode='embedding')
    else:
        token_type_embeds = 0
    hidden_states = ((inputs_embeds + position_embeds) + token_type_embeds)
    hidden_states = self.drop(hidden_states, training=training)
    output_shape = (input_shape + [shape_list(hidden_states)[(- 1)]])
    all_attentions = []
    all_hidden_states = ()
    for (i, block) in enumerate(self.h):
        if self.output_hidden_states:
            all_hidden_states = (all_hidden_states + (tf.reshape(hidden_states, output_shape),))
        outputs = block([hidden_states, attention_mask, head_mask[i]], training=training)
        hidden_states = outputs[0]
        if self.output_attentions:
            all_attentions.append(outputs[1])
    hidden_states = tf.reshape(hidden_states, output_shape)
    if self.output_hidden_states:
        all_hidden_states = (all_hidden_states + (hidden_states,))
    outputs = (hidden_states,)
    if self.output_hidden_states:
        outputs = (outputs + (all_hidden_states,))
    if self.output_attentions:
        attention_output_shape = ((input_shape[:(- 1)] + [(- 1)]) + shape_list(all_attentions[0])[(- 2):])
        all_attentions = tuple((tf.reshape(t, attention_output_shape) for t in all_attentions))
        outputs = (outputs + (all_attentions,))
    return outputs
def call(self, inputs, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, training=False):
    if isinstance(inputs, (tuple, list)):
        input_ids = inputs[0]
        attention_mask = (inputs[1] if (len(inputs) > 1) else attention_mask)
        token_type_ids = (inputs[2] if (len(inputs) > 2) else token_type_ids)
        position_ids = (inputs[3] if (len(inputs) > 3) else position_ids)
        head_mask = (inputs[4] if (len(inputs) > 4) else head_mask)
        inputs_embeds = (inputs[5] if (len(inputs) > 5) else inputs_embeds)
        assert (len(inputs) <= 6), 'Too many inputs.'
    elif isinstance(inputs, dict):
        input_ids = inputs.get('input_ids')
        attention_mask = inputs.get('attention_mask', attention_mask)
        token_type_ids = inputs.get('token_type_ids', token_type_ids)
        position_ids = inputs.get('position_ids', position_ids)
        head_mask = inputs.get('head_mask', head_mask)
        inputs_embeds = inputs.get('inputs_embeds', inputs_embeds)
        assert (len(inputs) <= 6), 'Too many inputs.'
    else:
        input_ids = inputs
    if ((input_ids is not None) and (inputs_embeds is not None)):
        raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time')
    elif (input_ids is not None):
        input_shape = shape_list(input_ids)
        input_ids = tf.reshape(input_ids, [(- 1), input_shape[(- 1)]])
    elif (inputs_embeds is not None):
        input_shape = shape_list(inputs_embeds)[:(- 1)]
    else:
        raise ValueError('You have to specify either input_ids or inputs_embeds')
    if (position_ids is None):
        position_ids = tf.range(input_shape[(- 1)], dtype=tf.int32)[tf.newaxis, :]
    if (attention_mask is not None):
        attention_mask = attention_mask[:, tf.newaxis, tf.newaxis, :]
        attention_mask = tf.cast(attention_mask, tf.float32)
        attention_mask = ((1.0 - attention_mask) * (- 10000.0))
    else:
        attention_mask = None
    if (not (head_mask is None)):
        raise NotImplementedError
    else:
        head_mask = ([None] * self.num_hidden_layers)
    position_ids = tf.reshape(position_ids, [(- 1), shape_list(position_ids)[(- 1)]])
    if (inputs_embeds is None):
        inputs_embeds = self.tokens_embed(input_ids, mode='embedding')
    position_embeds = self.positions_embed(position_ids)
    if (token_type_ids is not None):
        token_type_ids = tf.reshape(token_type_ids, [(- 1), shape_list(token_type_ids)[(- 1)]])
        token_type_embeds = self.tokens_embed(token_type_ids, mode='embedding')
    else:
        token_type_embeds = 0
    hidden_states = ((inputs_embeds + position_embeds) + token_type_embeds)
    hidden_states = self.drop(hidden_states, training=training)
    output_shape = (input_shape + [shape_list(hidden_states)[(- 1)]])
    all_attentions = []
    all_hidden_states = ()
    for (i, block) in enumerate(self.h):
        if self.output_hidden_states:
            all_hidden_states = (all_hidden_states + (tf.reshape(hidden_states, output_shape),))
        outputs = block([hidden_states, attention_mask, head_mask[i]], training=training)
        hidden_states = outputs[0]
        if self.output_attentions:
            all_attentions.append(outputs[1])
    hidden_states = tf.reshape(hidden_states, output_shape)
    if self.output_hidden_states:
        all_hidden_states = (all_hidden_states + (hidden_states,))
    outputs = (hidden_states,)
    if self.output_hidden_states:
        outputs = (outputs + (all_hidden_states,))
    if self.output_attentions:
        attention_output_shape = ((input_shape[:(- 1)] + [(- 1)]) + shape_list(all_attentions[0])[(- 2):])
        all_attentions = tuple((tf.reshape(t, attention_output_shape) for t in all_attentions))
        outputs = (outputs + (all_attentions,))
    return outputs