How do I get access to the "last_hidden_state" for code generation models in huggingface?

Question

I'm trying to obtain the "last_hidden_state" (as explained here) for code generation models over here. I am unable to figure out how to proceed, other than manually downloading each code-generation-model and checking if its key has that attribute using the following code -

import numpy as np
from datasets import load_dataset
from transformers import AutoTokenizer
from transformers import AutoModel, AutoModelForCausalLM
import torch
from sklearn.linear_model import LogisticRegression

from transformers import AutoTokenizer, AutoModelWithLMHead

tokenizer = AutoTokenizer.from_pretrained("codeparrot/codeparrot")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = AutoModelWithLMHead.from_pretrained("codeparrot/codeparrot").to(device)

inputs = tokenizer("def hello_world():", return_tensors="pt")
inputs = {k:v.to(device) for k,v in inputs.items()}
with torch.no_grad():
    outputs = model(**inputs)
print(outputs.keys())

So far, I tried this strategy on CodeParrot and InCoder with no success. Perhaps there is a better way to access the values of the hidden layers?

Answer 1

The hidden_states of output from CodeGenForCausalLM is already the last_hidden_state for the codegen model. See: link

where hidden_states = transformer_outputs[0] is the output of CodeGenModel (link) and the transformer_outputs[0] is the last_hidden_state

        if not return_dict:
            return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)

        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=presents,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
        )