import torch


def get_hidden_states(encoded, token_ids_word, model, layers):
     """Push input IDs through model. Stack and sum `layers` (last four by default).
        Select only those subword token outputs that belong to our word of interest
        and average them."""
     with torch.no_grad():
        output = model(**encoded)
 
     # Get all hidden states
     states = output.hidden_states
     last_hidden_states = states[-1]
     #print("last hidden states shape", last_hidden_states.shape)
     #print("last hidden states value: ", last_hidden_states)
     # Stack and sum all requested layers
     output = torch.stack([states[i] for i in layers]).sum(0).squeeze()
     # Only select the tokens that constitute the requested word
     word_tokens_output = output[token_ids_word]
 
     return word_tokens_output.mean(dim=0)