如何在LLMChain中获取模型输出的Logits？

LLMChain本身默认仅返回生成的文本结果，不会直接输出模型的Logits。要获取Logits，你需要调整LLM实例的配置，或者通过自定义回调函数捕获模型的底层输出，以下是具体实现方案：

方案1：调整HuggingFace LLM配置，让模型返回生成分数（Logits） #

如果你使用的是LangChain的HuggingFacePipeline封装Llama2，可以在初始化Pipeline时指定生成参数，让模型返回每个token的Logits：

from transformers import pipeline
from langchain.llms import HuggingFacePipeline
import torch

# 初始化带Logits返回的文本生成Pipeline
pipe = pipeline(
    "text-generation",
    model="meta-llama/Llama-2-7b-chat-hf",
    model_kwargs={"temperature": 0.7, "max_length": 512},
    return_full_text=False,
    # 关键参数：开启生成结果的字典返回和分数输出
    return_dict_in_generate=True,
    output_scores=True
)

llm = HuggingFacePipeline(pipeline=pipe)

# 构建你的Prompt和LLMChain
prompt = PromptTemplate(template=template, input_variables=["text"])
llm_chain = LLMChain(prompt=prompt, llm=llm)

之后调用LLMChain并提取Logits：

# 运行链并获取完整输出
response = llm_chain.generate(inputs={"text": "你的输入文本"})

# 提取每个生成token的Logits和概率
logits_list = []
probs_list = []
for generation in response.generations[0]:
    # scores是每个生成token的Logits张量，形状为(vocab_size,)
    for score in generation.scores:
        logit = score.detach().cpu()
        logits_list.append(logit)
        # 计算对数概率
        prob = torch.log_softmax(logit.float(), dim=-1).detach().cpu()
        probs_list.append(prob)

方案2：自定义回调函数捕获Logits #

通过LangChain的回调机制，在模型生成过程中实时捕获Logits，适合需要在链运行过程中处理Logits的场景：

from langchain.callbacks.base import BaseCallbackHandler
import torch

class LogitCaptureCallback(BaseCallbackHandler):
    def __init__(self):
        self.logits = []
        self.probs = []

    def on_llm_end(self, response, **kwargs):
        # 遍历生成结果，提取每个token的分数
        for generation in response.generations[0]:
            if generation.scores:
                for score in generation.scores:
                    logit = score.detach().cpu()
                    self.logits.append(logit)
                    prob = torch.log_softmax(logit.float(), dim=-1).detach().cpu()
                    self.probs.append(prob)

# 初始化回调实例
logit_callback = LogitCaptureCallback()

# 构建带回调的LLMChain
llm_chain = LLMChain(prompt=prompt, llm=llm, callbacks=[logit_callback])

# 运行链
result = llm_chain.run(text="你的输入文本")

# 访问捕获的Logits和概率
print("捕获的Logits列表:", logit_callback.logits)
print("捕获的对数概率列表:", logit_callback.probs)

方案3：绕开LLMChain，直接调用LLM的generate方法 #

LLMChain本质是封装了Prompt构建和LLM调用流程，你可以直接构建Prompt文本后调用LLM的generate方法，更直接地获取底层输出：

# 手动构建Prompt文本
prompt_text = prompt.format(text="你的输入文本")

# 调用LLM的generate方法获取完整输出
llm_output = llm.generate([prompt_text])

# 提取Logits和概率
logits = []
probs = []
for generation in llm_output.generations[0]:
    if generation.scores:
        for score in generation.scores:
            logits.append(score.detach().cpu())
            probs.append(torch.log_softmax(score.float(), dim=-1).detach().cpu())

注意事项 #

• 不同LangChain的LLM封装（如HuggingFaceLLM和HuggingFacePipeline）的输出结构可能略有差异，需根据实际情况调整提取逻辑。
• 开启output_scores=True会增加内存占用，若生成较长文本需注意资源消耗。

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