import torch


def run_inference(before, within, after, position_embeddings,
                  sentence_embeddings, attention_mask, embedding_map, model):
  """Run inference on a given one hot embedding split into three parts.

  Args:
    before: Tensor containing all the one_hot values before the changed part.
    within: The changed one_hot part of the sentence.
    after: Tensor containing all the one_hot values after the changed part.
    position_embeddings: Positional embeddings used by BERT.
    sentence_embeddings: Sentence embeddings used by BERT.
    attention_mask: Used by BERT for inference.
    embedding_map: Containing all the token embeddings for BERT.
    model: BERT model to run inference on.

  Returns:
    layers: The layer activations of the model.
  """
  one_hots = torch.cat([before, within, after], dim=1)
  # Get the word embedding from the smooth one_hots vector
  words_embeddings = torch.matmul(one_hots, embedding_map.embedding_map)
  # Assemble the actual embedding
  embeddings = words_embeddings + position_embeddings + sentence_embeddings
  embeddings = model.embeddings.LayerNorm(embeddings)
  # Get the prediction by the model
  layers = model.encoder(embeddings, attention_mask)
  # Reorder to get more intuitive ordering of batches, layers, words, neurons
  layers = torch.stack(layers).permute(1, 0, 2, 3)
  return layers