# Buggy code
def _get_model_features_for_batch(
    self, batch_token_ids: List[List[int]], batch_tokens: List[List[Token]]
) -> Tuple[np.ndarray, np.ndarray]:
    """Compute dense features of each example in the batch.

    We first add the special tokens corresponding to each language model. Next, we
    add appropriate padding and compute a mask for that padding so that it doesn't
    affect the feature computation. The padded batch is next fed to the language
    model and token level embeddings are computed. Using the pre-computed mask,
    embeddings for non-padding tokens are extracted and subsequently sentence
    level embeddings are computed.

    Args:
        batch_token_ids: List of token ids of each example in the batch.

    Returns:
        Sentence and token level dense representations.
    """
    # Let's first add tokenizer specific special tokens to all examples
    batch_token_ids_augmented = self._add_lm_specific_special_tokens(batch_token_ids)

    # Let's first add padding so that whole batch can be fed to the model
    actual_sequence_lengths, padded_token_ids = self._add_padding_to_batch(
        batch_token_ids_augmented
    )

    # Compute attention mask based on actual_sequence_length
    batch_attention_mask = self._compute_attention_mask(actual_sequence_lengths)

    # Get token level features from the model
    sequence_hidden_states = self._compute_batch_sequence_features(
        batch_attention_mask, padded_token_ids
    )

    # Extract features for only non-padding tokens
    sequence_nonpadded_embeddings = self._extract_nonpadded_embeddings(
        sequence_hidden_states, actual_sequence_lengths
    )

    # Extract sentence level and post-processed features
    (
        sentence_embeddings,
        sequence_embeddings,
    ) = self._post_process_sequence_embeddings(sequence_nonpadded_embeddings)

    # shape of matrix for all sequence embeddings
    batch_dim = len(sequence_embeddings)
    seq_dim = max(e.shape[0] for e in sequence_embeddings)
    feature_dim = sequence_embeddings[0].shape[1]
    shape = (batch_dim, seq_dim, feature_dim)

    # align features with tokens so that we have just one vector per token
    # (don't include sub-tokens)
    sequence_embeddings = train_utils.align_token_features(
        batch_tokens, sequence_embeddings, shape
    )

    # sequence_embeddings is a padded numpy array
    # remove the padding, keep just the non-zero vectors
    sequence_final_embeddings = []
    for embeddings, tokens in zip(sequence_embeddings, batch_tokens):
        sequence_final_embeddings.append(embeddings[: len(tokens)])
    sequence_final_embeddings = np.array(sequence_final_embeddings)

    return sentence_embeddings, sequence_final_embeddings


# Traceback (tensorflow.python.framework.errors_impl.InvalidArgumentError)
2020-07-16 13:05:16 INFO     rasa.nlu.model  - Starting to train component HFTransformersNLP
Traceback (most recent call last):
File "/home/ubuntu/anaconda3/bin/rasa", line 8, in <module>
sys.exit(main())
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/__main__.py", line 92, in main
cmdline_arguments.func(cmdline_arguments)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/cli/train.py", line 140, in train_nlu
persist_nlu_training_data=args.persist_nlu_data,
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/train.py", line 414, in train_nlu
persist_nlu_training_data,
File "uvloop/loop.pyx", line 1456, in uvloop.loop.Loop.run_until_complete
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/train.py", line 453, in _train_nlu_async
persist_nlu_training_data=persist_nlu_training_data,
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/train.py", line 482, in _train_nlu_with_validated_data
persist_nlu_training_data=persist_nlu_training_data,
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/nlu/train.py", line 90, in train
interpreter = trainer.train(training_data, **kwargs)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/nlu/model.py", line 191, in train
updates = component.train(working_data, self.config, **context)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/nlu/utils/hugging_face/hf_transformers.py", line 472, in train
batch_docs = self._get_docs_for_batch(batch_messages, attribute)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/nlu/utils/hugging_face/hf_transformers.py", line 423, in _get_docs_for_batch
) = self._get_model_features_for_batch(batch_token_ids)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/nlu/utils/hugging_face/hf_transformers.py", line 387, in _get_model_features_for_batch
batch_attention_mask, padded_token_ids
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/nlu/utils/hugging_face/hf_transformers.py", line 347, in _compute_batch_sequence_features
np.array(padded_token_ids), attention_mask=np.array(batch_attention_mask)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 822, in __call__
outputs = self.call(cast_inputs, *args, **kwargs)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/transformers/modeling_tf_bert.py", line 708, in call
outputs = self.bert(inputs, **kwargs)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 822, in __call__
outputs = self.call(cast_inputs, *args, **kwargs)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/transformers/modeling_tf_bert.py", line 570, in call
embedding_output = self.embeddings([input_ids, position_ids, token_type_ids, inputs_embeds], training=training)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 822, in __call__
outputs = self.call(cast_inputs, *args, **kwargs)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/transformers/modeling_tf_bert.py", line 150, in call
return self._embedding(inputs, training=training)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/transformers/modeling_tf_bert.py", line 173, in _embedding
position_embeddings = self.position_embeddings(position_ids)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 822, in __call__
outputs = self.call(cast_inputs, *args, **kwargs)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/layers/embeddings.py", line 184, in call
out = embedding_ops.embedding_lookup(self.embeddings, inputs)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/embedding_ops.py", line 323, in embedding_lookup
transform_fn=None)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/embedding_ops.py", line 137, in _embedding_lookup_and_transform
array_ops.gather(params[0], ids, name=name), ids, max_norm)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/util/dispatch.py", line 180, in wrapper
return target(*args, **kwargs)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/array_ops.py", line 4106, in gather
return params.sparse_read(indices, name=name)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/resource_variable_ops.py", line 645, in sparse_read
self._handle, indices, dtype=self._dtype, name=name)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_resource_variable_ops.py", line 556, in resource_gather
_ops.raise_from_not_ok_status(e, name)
File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 6606, in raise_from_not_ok_status
six.raise_from(core._status_to_exception(e.code, message), None)
File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError: indices[0,512] = 512 is not in [0, 512) [Op:ResourceGather] name: tf_bert_model/bert/embeddings/position_embeddings/embedding_lookup/

# Fixed code
def _get_model_features_for_batch(
    self,
    batch_token_ids: List[List[int]],
    batch_tokens: List[List[Token]],
    batch_examples: List[Message],
    attribute: Text,
    inference_mode: bool = False,
) -> Tuple[np.ndarray, np.ndarray]:
    """Compute dense features of each example in the batch.

    We first add the special tokens corresponding to each language model. Next, we
    add appropriate padding and compute a mask for that padding so that it doesn't
    affect the feature computation. The padded batch is next fed to the language
    model and token level embeddings are computed. Using the pre-computed mask,
    embeddings for non-padding tokens are extracted and subsequently sentence
    level embeddings are computed.

    Args:
        batch_token_ids: List of token ids of each example in the batch.
        batch_tokens: List of token objects for each example in the batch.
        batch_examples: List of examples in the batch.
        attribute: attribute of the Message object to be processed.
        inference_mode: Whether the call is during training or during inference.

    Returns:
        Sentence and token level dense representations.
    """
    # Let's first add tokenizer specific special tokens to all examples
    batch_token_ids_augmented = self._add_lm_specific_special_tokens(batch_token_ids)

    # Compute sequence lengths for all examples
    (
        actual_sequence_lengths,
        max_input_sequence_length,
    ) = self._extract_sequence_lengths(batch_token_ids_augmented)

    # Validate that all sequences can be processed based on their sequence lengths and
    # the maximum sequence length the model can handle
    self._validate_sequence_lengths(
        actual_sequence_lengths, batch_examples, attribute, inference_mode
    )

    # Add padding so that whole batch can be fed to the model
    padded_token_ids = self._add_padding_to_batch(
        batch_token_ids_augmented, max_input_sequence_length
    )

    # Compute attention mask based on actual_sequence_length
    batch_attention_mask = self._compute_attention_mask(
        actual_sequence_lengths, max_input_sequence_length
    )

    # Get token level features from the model
    sequence_hidden_states = self._compute_batch_sequence_features(
        batch_attention_mask, padded_token_ids
    )

    # Extract features for only non-padding tokens
    sequence_nonpadded_embeddings = self._extract_nonpadded_embeddings(
        sequence_hidden_states, actual_sequence_lengths
    )

    # Extract sentence level and post-processed features
    (
        sentence_embeddings,
        sequence_embeddings,
    ) = self._post_process_sequence_embeddings(sequence_nonpadded_embeddings)

    # Pad zeros for examples which were truncated in inference mode.
    # This is intentionally done after sentence embeddings have been extracted so that they are not affected
    sequence_embeddings = self._add_extra_padding(
        sequence_embeddings, actual_sequence_lengths
    )

    # shape of matrix for all sequence embeddings
    batch_dim = len(sequence_embeddings)
    seq_dim = max(e.shape[0] for e in sequence_embeddings)
    feature_dim = sequence_embeddings[0].shape[1]
    shape = (batch_dim, seq_dim, feature_dim)

    # align features with tokens so that we have just one vector per token
    # (don't include sub-tokens)
    sequence_embeddings = train_utils.align_token_features(
        batch_tokens, sequence_embeddings, shape
    )

    # sequence_embeddings is a padded numpy array
    # remove the padding, keep just the non-zero vectors
    sequence_final_embeddings = []
    for embeddings, tokens in zip(sequence_embeddings, batch_tokens):
        sequence_final_embeddings.append(embeddings[: len(tokens)])
    sequence_final_embeddings = np.array(sequence_final_embeddings)

    return sentence_embeddings, sequence_final_embeddings