使用LlamaIndex+自定义OpenAssistant Pythia模型如何避免上下文外作答？

我正在使用LlamaIndex搭配OpenAssistant Pythia自定义LLM，代码如下。data目录下的france.txt文件内容为The captial of France is XYZ，但运行代码查询What is capital of france?时，模型仍回答“Paris”，需要实现仅基于输入文件内容作答。

import os
from llama_index import (
    GPTKeywordTableIndex,
    SimpleDirectoryReader,
    LLMPredictor,
    ServiceContext,
    PromptHelper
)
from langchain import OpenAI

import torch
from langchain.llms.base import LLM
from llama_index import SimpleDirectoryReader, LangchainEmbedding, GPTListIndex
from llama_index import LLMPredictor, ServiceContext
from transformers import pipeline
from typing import Optional, List, Mapping, Any

from transformers import AutoModelForCausalLM, AutoTokenizer

# define prompt helper
# set maximum input size
max_input_size = 2048
# set number of output tokens
num_output = 256
# set maximum chunk overlap
max_chunk_overlap = 20
prompt_helper = PromptHelper(max_input_size, num_output, max_chunk_overlap)

class CustomLLM(LLM):
    model_name="OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5"
    tokenizer = AutoTokenizer.from_pretrained("OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5", padding_side="left")
    model = AutoModelForCausalLM.from_pretrained("OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5", 
                                             load_in_8bit=True,
                                             device_map="auto")
    pipeline = pipeline(
        "text-generation",
        model=model, 
        tokenizer=tokenizer, 
        max_length=512,
        temperature=0.7,
        top_p=0.95,
        repetition_penalty=1.15
    )

    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        prompt_length = len(prompt)
        response = self.pipeline(prompt, max_new_tokens=num_output)[0]["generated_text"]

        # only return newly generated tokens
        return response[prompt_length:]

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        return {"name_of_model": self.model_name}

    @property
    def _llm_type(self) -> str:
        return "custom"

os.environ['OPENAI_API_KEY'] = 'demo'
documents = SimpleDirectoryReader('data').load_data()

# define LLM
llm_predictor = LLMPredictor(llm=CustomLLM())
service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor)

# build index
index = GPTKeywordTableIndex.from_documents(documents, service_context=service_context)

# get response from query
query_engine = index.as_query_engine()
response = query_engine.query("What is capital of france?");

print(response)

解决方案 #

1. 自定义查询Prompt，强制模型基于给定上下文回答 #

LlamaIndex默认Prompt未明确约束模型只能使用提供的上下文，需自定义模板明确规则：

from llama_index.prompts import PromptTemplate

# 定义仅基于上下文回答的Prompt
query_prompt = PromptTemplate(
    "以下是你可以参考的唯一上下文信息:\n{context_str}\n请仅基于上述上下文回答问题: {query_str}\n如果上下文无相关内容，请回答'无法从给定文档中找到答案'"
)

# 使用自定义Prompt创建查询引擎
query_engine = index.as_query_engine(text_qa_template=query_prompt)

2. 调整LLM生成参数，降低模型创造性 #

当前temperature=0.7会让模型倾向于生成创造性内容，降低该值并关闭采样，减少模型依赖预训练知识的概率：
修改CustomLLM中的pipeline初始化代码：

pipeline = pipeline(
    "text-generation",
    model=model, 
    tokenizer=tokenizer, 
    max_length=512,
    temperature=0.1,  # 降低温度，减少随机输出
    top_p=0.95,
    repetition_penalty=1.15,
    do_sample=False  # 关闭采样，选择概率最高的输出
)

3. 验证上下文检索是否正确 #

确保LlamaIndex正确检索到目标文件内容，添加调试代码检查检索结果：

# 创建检索器并查看检索到的上下文
retriever = index.as_retriever(similarity_top_k=1)
nodes = retriever.retrieve("What is capital of france?")
for node in nodes:
    print("检索到的上下文:", node.text)

如果未检索到目标内容，建议更换索引类型为GPTVectorStoreIndex（基于向量检索更精准），或检查文件路径、编码是否正常。

4. 添加Stop序列，终止模型额外输出 #

在CustomLLM的_call方法中添加自定义Stop序列，防止模型超出上下文继续生成内容：

def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
    # 自定义停止序列，避免模型额外输出
    custom_stop = stop or ["\n", "Question:", "Answer:"]
    response = self.pipeline(prompt, max_new_tokens=num_output, stop=custom_stop)[0]["generated_text"]
    prompt_length = len(prompt)
    return response[prompt_length:]

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