稻草堆里的针 - LLM 压力测试

一种简单的“稻草堆里的针”分析，用于测试长上下文 LLM 的上下文内检索能力。

支持的模型提供商：OpenAI、Anthropic、Cohere

观看概述视频，了解幕后详情。

测试内容

1. 将一个随机的事实或陈述（即“针”）放置在长上下文窗口的中间（即“稻草堆”）。
2. 要求模型检索该陈述。
3. 遍历不同的文档深度（即“针”的位置）和上下文长度，以衡量性能。

这段代码支持了这次 OpenAI 分析和Anthropic 分析。

原始测试的结果位于 /original_results 目录中。自这些测试运行以来，脚本已经进行了大量升级，因此数据格式可能与您当前脚本的结果不一致。

快速开始

设置虚拟环境

我们建议设置一个虚拟环境来隔离 Python 依赖项，确保项目专用的包不会与系统范围的安装发生冲突。

python3 -m venv venv
source venv/bin/activate

环境变量

• NIAH_MODEL_API_KEY - 用于与模型交互的 API 密钥。根据提供商的不同，此密钥将与相应的 SDK 正确配合使用。
• NIAH_EVALUATOR_API_KEY - 如果使用 openai 评估策略，则需要此 API 密钥。

安装包

从 PyPi 安装该包：

pip install needlehaystack

运行测试

通过命令行调用入口点 needlehaystack.run_test 即可开始使用该包。

然后，您可以使用以下命令行参数对 OpenAI、Anthropic 或 Cohere 的模型进行分析：

• provider - 模型的提供商，可选值为 openai、anthropic 和 cohere。默认值为 openai。
• evaluator - 评估器，可以是 model 或 LangSmith。有关 LangSmith 的更多信息请见下文。如果使用 model，目前仅支持 openai。默认值为 openai。
• model_name - 提供商可访问的语言模型名称。默认值为 gpt-3.5-turbo-0125。
• evaluator_model_name - 评估器可访问的语言模型名称。默认值为 gpt-3.5-turbo-0125。

此外，LLMNeedleHaystackTester 的参数也可以作为命令行参数传递，但 model_to_test 和 evaluator 除外。

以下是一些示例用法。

以下命令针对 OpenAI 模型 gpt-3.5-turbo-0125 运行测试，上下文长度为 2000，文档深度为 50%。

needlehaystack.run_test --provider openai --model_name "gpt-3.5-turbo-0125" --document_depth_percents "[50]" --context_lengths "[2000]"

以下命令针对 Anthropic 模型 claude-2.1 运行测试，上下文长度为 2000，文档深度为 50%。

needlehaystack.run_test --provider anthropic --model_name "claude-2.1" --document_depth_percents "[50]" --context_lengths "[2000]"

以下命令针对 Cohere 模型 command-r 运行测试，上下文长度为 2000，文档深度为 50%。

needlehaystack.run_test --provider cohere --model_name "command-r" --document_depth_percents "[50]" --context_lengths "[2000]"

贡献者指南

1. 分支并克隆仓库。
2. 按照上述说明创建并激活虚拟环境。
3. 按照上述说明设置环境变量。
4. 在仓库根目录下运行以下命令，以可编辑模式安装该包：

pip install -e .

现在您可以在测试用例中导入 needlehaystack 包。开发、修改并在本地进行测试。

LLMNeedleHaystackTester 参数：

• model_to_test - 要进行“稻草堆里的针”测试的模型。默认值为 None。
• evaluator - 用于评估模型响应的评估器。默认值为 None。
• needle - 将被放置在您的上下文中（即“稻草堆”）的陈述或事实。
• haystack_dir - 包含要加载为背景上下文的文本文件的目录。仅支持文本文件。
• retrieval_question - 用于在背景上下文中检索“针”的问题。
• results_version - 如果您希望对相同的长度/深度组合多次运行测试，请更改版本号。
• num_concurrent_requests - 默认值为 1。如果您希望并行运行更多请求，请将其设置为更高的值。请注意速率限制。
• save_results - 是否要将结果保存到文件。无论是否保存，结果都会暂时存储在对象中。True/False。如果 save_results = True，则此脚本会填充一个 result/ 目录，其中包含评估信息。由于可能存在并发请求，每次新测试都会保存为几个文件。
• save_contexts - 是否要将上下文保存到文件。警告：这些文件会非常大。True/False。
• final_context_length_buffer - 用于扣除每个输入中的系统消息和输出标记的上下文量。目前采用静态值，未来可能会更智能。默认值为 200 个标记。
• context_lengths_min - 您要遍历的上下文长度列表的起点。
• context_lengths_max - 您要遍历的上下文长度列表的终点。
• context_lengths_num_intervals - 您要在最小值和最大值之间遍历的区间数量。
• context_lengths - 自定义的上下文长度集合。如果已设置，则会覆盖 context_lengths_min、max 和 intervals 的值。
• document_depth_percent_min - 您文档深度的起点。应为大于 0 的整数。
• document_depth_percent_max - 您文档深度的终点。应为小于 100 的整数。
• document_depth_percent_intervals - 您要在最小值和最大值之间进行的迭代次数。
• document_depth_percents - 自定义的文档深度长度集合。如果已设置，则会覆盖 document_depth_percent_min、max 和 intervals 的值。
• document_depth_percent_interval_type - 决定要遍历的深度分布。“linear” 或 “sigmoid”。
• seconds_to_sleep_between_completions - 默认值为 None，如果您希望减缓请求速度，请设置秒数。
• print_ongoing_status - 默认值为 True，表示是否在测试完成时打印状态。

LLMMultiNeedleHaystackTester 参数：

• multi_needle - 是否运行多针测试。True 或 False。
• needles - 要插入上下文中的针的列表。

其他参数：

• model_name - 您想要使用的模型名称。应与需要传递给 API 的确切值匹配。例如：对于 OpenAI 的推理和评估模型，应为 gpt-3.5-turbo-0125。

结果可视化

LLMNeedleInHaystackVisualization.ipynb 包含用于生成透视表可视化的代码。随后，该透视表被导入 Google Slides 进行自定义标注和格式化。请参阅 Google Slides 版本。有关此可视化创建过程的概述，请参见 此处。

OpenAI 的 GPT-4-128K（2023年11月8日运行）

Anthropic 的 Claude 2.1（2023年11月21日运行）

多针评估器

为在上下文中插入多根“针”，请使用 --multi_needle True。

此选项会将第一根“针”放置在指定的 depth_percent 深度处，然后在该深度之后的剩余上下文中均匀分布后续的“针”。

为了实现均匀间隔，它会计算 depth_percent_interval 如下：

depth_percent_interval = (100 - depth_percent) / len(self.needles)

因此，第一根“针”位于 depth_percent 百分比深度处，第二根位于 depth_percent + depth_percent_interval，第三根位于 depth_percent + 2 * depth_percent_interval，以此类推。

以下示例展示了在 10 根“针”和 depth_percent 为 40% 的情况下，各根“针”的深度百分比：

depth_percent_interval = (100 - 40) / 10 = 6

针 1: 40
针 2: 40 + 6 = 46
针 3: 40 + 2 * 6 = 52
针 4: 40 + 3 * 6 = 58
针 5: 40 + 4 * 6 = 64
针 6: 40 + 5 * 6 = 70
针 7: 40 + 6 * 6 = 76
针 8: 40 + 7 * 6 = 82
针 9: 40 + 8 * 6 = 88
针 10: 40 + 9 * 6 = 94

LangSmith 评估器

您可以使用 LangSmith 来编排评估并存储结果。

(1) 注册 LangSmith
(2) 按照设置说明配置 LangSmith 的环境变量。
(3) 在“数据集 + 测试”选项卡中，使用“+ 数据集”创建一个新数据集，例如命名为 multi-needle-eval-sf。
(4) 向数据集中添加一个测试问题：

问题：旧金山有哪些最佳游玩体验？  
答案：“旧金山的最佳游玩体验包括：1) 去多洛雷斯公园；2) 在托尼那不勒斯披萨店用餐；3) 参观恶魔岛；4) 徒步登上双峰山；5) 骑自行车穿越金门大桥。”

(5) 使用 --evaluator langsmith 和 --eval_set multi-needle-eval-sf 运行，以针对我们刚刚创建的评估集进行测试。

让我们在一个新的数据集 multi-needle-eval-pizza 上一起看看这些功能如何协同工作。

以下是 multi-needle-eval-pizza 评估集，其中包含一个问题和参考答案。您也可以查看相关的运行结果： https://smith.langchain.com/public/74d2af1c-333d-4a73-87bc-a837f8f0f65c/d

以下是使用多针评估并传递相关“针”的命令：

needlehaystack.run_test --evaluator langsmith --context_lengths_num_intervals 3 --document_depth_percent_intervals 3 --provider openai --model_name "gpt-4-0125-preview" --multi_needle True --eval_set multi-needle-eval-pizza --needles '["无花果是三种最美味的披萨配料之一。", "意大利生火腿是三种最美味的披萨配料之一。", "山羊奶酪是三种最美味的披萨配料之一。"]'

许可证

本项目采用 MIT 许可证授权——详情请参阅 LICENSE 文件。使用本软件需按照许可证中的规定，注明原作者和项目的归属。

LLMTest_NeedleInAHaystack 快速上手指南

LLMTest_NeedleInAHaystack 是一个用于压力测试大语言模型（LLM）长上下文检索能力的开源工具。它通过“大海捞针”（Needle In A Haystack）的方法，将特定事实（针）隐藏在长文本（草堆）的不同位置，测试模型能否准确检索该信息。支持 OpenAI、Anthropic 和 Cohere 等主流模型提供商。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux, macOS 或 Windows (WSL 推荐)
• Python 版本：Python 3.8 或更高版本
• API Key：
  • NIAH_MODEL_API_KEY：目标测试模型的 API Key（如 OpenAI, Anthropic, Cohere）。
  • NIAH_EVALUATOR_API_KEY：如果使用 OpenAI 作为评估器，需提供其 API Key。

建议创建虚拟环境以隔离依赖：

python3 -m venv venv
source venv/bin/activate

设置环境变量（以 zsh/bash 为例）：

export NIAH_MODEL_API_KEY="your_model_api_key"
export NIAH_EVALUATOR_API_KEY="your_evaluator_api_key"

安装步骤

您可以直接从 PyPI 安装该工具包：

pip install needlehaystack

提示：国内开发者若遇到下载速度慢的问题，可使用清华或阿里镜像源加速安装：

pip install needlehaystack -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

安装完成后，可通过命令行入口 needlehaystack.run_test 直接运行测试。

最简单的测试示例

以下命令将对 OpenAI 的 gpt-3.5-turbo-0125 模型进行一次基础测试：

• 上下文长度：2000 tokens
• 针的位置：文档深度的 50% 处

needlehaystack.run_test --provider openai --model_name "gpt-3.5-turbo-0125" --document_depth_percents "[50]" --context_lengths "[2000]"

测试其他模型提供商

只需更改 --provider 和 --model_name 参数即可适配不同厂商的模型。

测试 Anthropic (Claude):

needlehaystack.run_test --provider anthropic --model_name "claude-2.1" --document_depth_percents "[50]" --context_lengths "[2000]"

测试 Cohere:

needlehaystack.run_test --provider cohere --model_name "command-r" --document_depth_percents "[50]" --context_lengths "[2000]"

进阶参数说明

您可以通过组合以下参数进行更复杂的压力测试：

• --context_lengths: 自定义上下文长度列表，例如 "[2000, 4000, 8000]"。
• --document_depth_percents: 自定义针在文档中的位置百分比列表，例如 "[0, 25, 50, 75, 100]"。
• --multi_needle: 设置为 True 可启用多针测试（在上下文中插入多个关键信息）。
• --evaluator: 指定评估器，默认为 openai，也支持 langsmith。

运行结束后，测试结果通常会保存在当前的 results/ 目录中，您可以使用提供的 Jupyter Notebook (LLMNeedleInHaystackVisualization.ipynb) 进行可视化分析。