如何在LlamaIndex Pandas查询管道的LLM中添加Batch_Size或重试点？

我使用LlamaIndex的Pandas Dataframe Query Pipeline处理仅500行的数据集时，只得到部分输出，推测是LLM的上下文限制导致。想解决这个问题，LlamaIndex建议使用Batchsize，但作为新手，不确定要在LLM的何处及如何添加该配置。

我的LLM和Embeddings代码如下：

from llama_index.llms.azure_openai import AzureOpenAI
from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
import logging
import sys

logging.basicConfig(
    stream=sys.stdout, level=logging.INFO
)  # logging.DEBUG for more verbose output
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))


api_key = "xxxx"
azure_endpoint = "xxxx"
api_version = "xxxx"

llm = AzureOpenAI(
    model="gpt-35-turbo",
    deployment_name="gpt-35-turbo",
    api_key=api_key,
    azure_endpoint=azure_endpoint,
    api_version=api_version,
    temperature=0,
)

# 需要部署自己的嵌入模型和聊天完成模型
embed_model = AzureOpenAIEmbedding(
    model="text-embedding-ada-002",
    deployment_name="text-embedding-ada-002",
    api_key=api_key,
    azure_endpoint=azure_endpoint,
    api_version=api_version,
)

from llama_index.core import Settings

Settings.llm = llm
Settings.embed_model = embed_model


import logging
import sys
from IPython.display import Markdown, display

import pandas as pd
from llama_index.core.query_engine import PandasQueryEngine


logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

解决方案 #

Batchsize配置不是在LLM实例中设置，而是在初始化PandasQueryEngine时添加参数，具体步骤如下：

1. 初始化PandasQueryEngine时指定batch_size
   当你创建PandasQueryEngine实例时，直接传入batch_size参数，控制每次传给LLM处理的数据行数，避免超过上下文窗口限制。示例代码：

   # 假设你已经加载数据集到df变量中
   df = pd.read_csv("your_dataset.csv")  # 替换为你的数据加载方式
   
   # 初始化带batch_size的PandasQueryEngine
   query_engine = PandasQueryEngine(
       df=df,
       llm=llm,
       batch_size=50,  # 可根据实际情况调整，比如30-100
       verbose=True  # 开启日志查看批次处理情况
   )
   
   # 执行查询
   response = query_engine.query("你的查询语句")
   display(Markdown(f"<b>{response}</b>"))
   

2. 调整batch_size的参考建议

   • 先从较小的值（如30-50）开始测试，根据输出完整性和响应速度逐步调整
   • 结合你的LLM上下文限制：比如gpt-35-turbo有4k/16k token的上下文，若单行列较多、内容长，就适当减小batch_size；反之可以增大
   • 开启verbose=True后，能看到每一批数据的处理日志，方便排查问题

3. 额外优化手段

   • 预处理数据集：删除不必要的冗余列，减少单批数据的token总量
   • 若仍有问题，可以尝试升级LLM模型到更大上下文版本（如gpt-35-turbo-16k）

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