如何在HuggingFace中获取代码生成模型的last_hidden_state

核心问题原因 #

你之前用AutoModelWithLMHead失败，是因为这个类已被弃用，且默认不会输出隐藏层状态。针对CodeParrot、InCoder这类因果语言模型，有两种更可靠的方式获取last_hidden_state：

方法一：用AutoModelForCausalLM并开启隐藏层输出 #

这类模型是专门针对因果生成任务设计的，加载时设置output_hidden_states=True，就能拿到包括last_hidden_state在内的所有层隐藏状态：

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("codeparrot/codeparrot")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 关键：开启output_hidden_states参数
model = AutoModelForCausalLM.from_pretrained(
    "codeparrot/codeparrot",
    output_hidden_states=True
).to(device)

inputs = tokenizer("def hello_world():", return_tensors="pt").to(device)
with torch.no_grad():
    outputs = model(**inputs)

# outputs.hidden_states是元组，最后一个元素即为last_hidden_state
last_hidden_state = outputs.hidden_states[-1]
print(f"last_hidden_state形状: {last_hidden_state.shape}")

方法二：用AutoModel直接提取编码器隐藏层 #

如果不需要生成任务的LM头，仅需编码器的最后隐藏状态，可以直接用AutoModel加载，它默认会输出last_hidden_state：

import torch
from transformers import AutoTokenizer, AutoModel

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

inputs = tokenizer("def hello_world():", return_tensors="pt").to(device)
with torch.no_grad():
    outputs = model(**inputs)

# 直接获取last_hidden_state
print(f"last_hidden_state形状: {outputs.last_hidden_state.shape}")

批量检查模型的小技巧 #

不用逐个下载模型验证，可先通过配置文件确认模型是否支持输出隐藏层，还能直接修改配置后加载：

from transformers import AutoConfig, AutoModelForCausalLM

# 加载模型配置
config = AutoConfig.from_pretrained("codeparrot/codeparrot")
# 查看默认是否输出隐藏层（大多默认False）
print(f"默认输出隐藏层: {config.output_hidden_states}")

# 修改配置并加载模型
config.output_hidden_states = True
model = AutoModelForCausalLM.from_pretrained("codeparrot/codeparrot", config=config)

内容的提问来源于stack exchange，提问作者desert_ranger