import torch


def sentence_emb(inp_token_embeddings, n=4):
    """
    Adds the last n layers of each token and stacks all the input token representations except the first and last ones.
    It considers that the first and last tokens are special BERT tokens ([CLS] and [SEP]) and it drops them within the
    function
    :param inp_token_embeddings: torch tensor with all the token embeddings in each layer
    :param n: number of layers to add up
    :return: torch representation of all tokens stacked
    """
    inp_token_embeddings = inp_token_embeddings[1:-1]
    token_vecs_sum = torch.stack(
        [torch.sum(token[-n:], dim=0) for token in inp_token_embeddings])  # sum across last 4 layers
    return token_vecs_sum