import torch


def get_embeddings(tokens_tensor, segments_tensor, model):
  """Obtain embeddigs for word, position, and sequence.

  adapted from:
  https://github.com/huggingface/pytorch-pretrained-BERT/blob/
  2a329c61868b20faee115a78bdcaf660ff74cf41/pytorch_pretrained_bert/
  modeling.py#L264-L277)

  Args:
    tokens_tensor: Tokens for which to get the tokens embedding.
    segments_tensor: Used to generate the segments embedding.
    model: The model used for inference.

  Returns:
    words_embeddings: Word embeddings for the given tokens.
    position_embeddings: Positional embeddings for the tokens.
    sentence embeddings: Sentence embeddings for the tokens.
  """
  seq_length = tokens_tensor.size(1)
  position_ids = torch.arange(seq_length, dtype=torch.long,
                              device=tokens_tensor.device)
  position_ids = position_ids.unsqueeze(0).expand_as(tokens_tensor)
  if segments_tensor is None:
    segments_tensor = torch.zeros_like(tokens_tensor)
  # Get the three token types (words, positions, and sentences individually)
  words_embeddings = model.embeddings.word_embeddings(tokens_tensor)
  position_embeddings = model.embeddings.position_embeddings(position_ids)
  sentence_embeddings = model.embeddings.token_type_embeddings(segments_tensor)
  return words_embeddings, position_embeddings, sentence_embeddings