SQL 生成评估

此仓库包含 Defog 用于评估生成的 SQL 的代码。它基于 Spider 的模式，但使用了一组全新的人工精选问题和查询，并按查询类别分组。如需深入了解我们创建这一评估方法的过程，请参阅这篇博文。

简介

我们的测试流程包括以下步骤。对于每一对问题和查询：

1. 我们生成一条 SQL 查询（可能来自大语言模型）。
2. 我们在各自的数据库上分别运行“黄金”查询和生成的查询，以获得两个包含结果的数据框。
3. 我们使用“精确匹配”和“子集匹配”来比较这两个数据框。TODO 添加博客文章链接。
4. 我们将这些结果以及其他感兴趣的指标（例如使用的 token 数、延迟）记录下来，并汇总结果以生成报告。

入门指南

这是一份全面的说明文档，假定您对命令行、Docker、在数据库上执行 SQL 查询以及常用的 Python 数据处理库（例如 pandas）有一定的了解。

安装依赖

首先，克隆我们存储数据库数据和模式的仓库。安装 requirements.txt 文件中列出的所有 Python 库。如果您使用 NER 启发式方法来实现我们的元数据剪枝方法（即 c 参数设置为大于 0 的任意值，详见下文），则还需要下载一个 spaCy 模型。最后，安装该库。

git clone https://github.com/defog-ai/defog-data.git
cd defog-data
pip install -r requirements.txt
pip install -e .

启动 Postgres 实例

接下来，您需要设置用于执行查询的数据库。我们在此使用 Postgres，因为它是在生产环境中分布最广、使用最广泛的开源数据库。此外，我们建议使用 Docker 来完成此操作，因为这是最简单的入门方式。您可以在这里安装 Docker。

安装 Docker 后，您可以使用以下命令创建 Docker 容器并启动 Postgres 数据库。我们建议将卷挂载到 data/postgres 目录以持久化数据，同时将卷挂载到 data/export 目录以便于导入数据。要创建容器，请运行：

mkdir data/postgres data/export
docker create --name postgres-sql-eval -e POSTGRES_PASSWORD=postgres -p 5432:5432 -v $(pwd)/data/postgres:/var/lib/postgresql/data -v $(pwd)/data/export:/export postgres:16-alpine

要启动容器，请运行：

docker start postgres-sql-eval

如果您想重置 Postgres 服务器实例的状态（例如由于临时连接导致的内存泄漏），可以将其关闭（并在之后重新启动）：

docker stop postgres-sql-eval
# 查看容器是否仍在：
docker container list -a

注意事项：

• 在运行上述命令之前，您需要先停止其他正在监听 5432 端口的 Postgres 实例。
• 您只需运行一次 docker create ... 命令来创建镜像，之后只需运行 docker start/stop postgres-sql-eval。
• 数据会持久化在 data/postgres 目录中，因此关闭容器并不关键。然而，如果删除 data/postgres 文件夹，所有数据都将丢失 T.T。
• 虽然我们将使用 Docker 来部署 Postgres 并进行初始化，但您也可以根据自己的本地安装情况修改脚本或说明。

将数据导入 Postgres

待导入的数据位于我们先前克隆的 defog-data 仓库中。每个文件夹包含对应单个数据库的元数据和数据（例如 academic 文件夹包含重新加载 “academic” 数据库所需的所有数据）。我们假设您已在本地安装了 psql 客户端。我们将使用以下命令在我们的 Postgres 实例中为 7 个 SQL 数据库中的每一个创建一个新的数据库：

# 设置以下环境变量
cd defog-data # 如果您尚未进入 defog-data 目录
export DBPASSWORD="postgres"
export DBUSER="postgres"
export DBHOST="localhost"
export DBPORT=5432
./setup.sh

将数据导入 Snowflake

如果您希望将数据导入 Snowflake，相关设置说明也位于 defog-data 仓库中。在安装 Snowflake CLI 后，按照文档配置您的凭据，并将其设置为如下所示的环境变量，然后运行设置命令。

export SFDBPASSWORD="your_password"
export SFDBUSER="your_username"
export SFDBACCOUNT="your_account"
export SFDBWAREHOUSE="your_warehouse"
./setup_snowflake.sh

请注意，在评估过程中，您需要使用 /data 目录下的 _snowflake 问题文件。这些查询已修改为适用于 Snowflake 数据库。

将数据导入 BigQuery、MySQL、SQLite、SQL Server

这些数据库管理系统的设置说明位于 defog-data 仓库中。请相应地配置您的凭据，设置好环境变量，然后使用以下命令转换并导入评估数据库：

python translate_ddl_dialect.py

在评估过程中，您需要设置正确的 --db_type 标志，并使用 /data 目录下对应的 {dialect} 问题文件。

使用私有数据（可选）

如果您拥有不想公开的私有数据集，但仍希望复用此处的代码进行评估，可以按照以下步骤操作。

• 首先为您的私有数据创建一个独立的 Git 仓库，并在其中包含 setup.py 文件，类似于 defog-data。
• 创建元数据和数据文件，并将它们导入到您的数据库中。这样做是为了让我们的评估框架能够使用实际数据运行生成的查询。您可以参考 defog-data 的 元数据对象 来了解表结构，以及 setup.sh 作为如何将数据导入数据库的示例。我们不规定具体的文件夹结构，您可以根据自己的需求自由组织数据，只要能够方便地将其导入数据库即可。
• 要使用我们的元数据剪枝工具，您需要定义以下内容：
  • 一种用于定义表之间可连接列的方法。在我们的案例中，我们使用一个字典 columns_join，其键为数据库名称，值则是嵌套字典，表示表元组与列名元组之间的对应关系。您可以参考原始数据来了解我们是如何生成这个字典的。

完成上述三个步骤后，您需要：

• 将您的数据库作为依赖项安装，运行 pip install -e .（-e 参数可使代码更改自动生效，无需重新安装）。
• 将 prune_metadata_str 中的相关函数调用和变量替换为您自己导入的函数和变量。请注意，您可能不会将包/模块命名为 defog_data_private.supplementary，因此请根据实际情况进行修改。

需要注意的几点：

• 如果您没有向数据库中填充数据（即仅创建表而未插入数据），大多数情况下会返回空的数据框（无论生成的查询是否符合预期），这会导致所有结果都匹配，并产生大量假阳性。因此，建议您向数据库中填充一些有意义的数据，以便在查询不符合预期时能够返回不同的结果。
• 如果您使用私有数据进行测试，还需要将问题文件指向您自己的问题文件（根据您的数据库模式定制）。

运行器

运行器负责处理每个问题/查询对的工作配置（例如并行化、批处理、所选模型等）。

我们提供了几种常见的运行器：runners/openai_runner.py 用于调用 OpenAI 的 API（支持并行化）、runners/anthropic_runner 用于调用 Anthropic 的 API、runners/hf_runner.py 用于调用本地 Hugging Face 模型，最后，runners/api_runner.py 允许您使用自定义 API 进行评估。

运行测试

OpenAI

如果您计划调用 OpenAI 或 Anthropic 等 LLM API，请务必在运行测试前将您的 API 密钥（OPENAI_API_KEY 或 ANTHROPIC_API_KEY）设置为环境变量。

要仅使用 10 个问题（而不是全部 200 个）进行测试，并以 5 个任务并行执行：

python main.py \
  -db postgres \
  -q "data/questions_gen_postgres.csv" "data/instruct_basic_postgres.csv" "data/instruct_advanced_postgres.csv" \
  -o results/openai_classic.csv results/openai_basic.csv results/openai_advanced.csv \
  -g oa \
  -f prompts/prompt_openai.json \
  -m o3-mini \
  -p 5 \
  -c 0

如果使用最新的 o1-* 模型进行测试（这些模型不支持系统提示词），则应使用不同的提示词文件，减少并行请求数量并增加超时时间：

python main.py \
  -db postgres \
  -q "data/questions_gen_postgres.csv" "data/instruct_basic_postgres.csv" "data/instruct_advanced_postgres.csv" \
  -o results/openai_o1mini_classic.csv results/openai_o1mini_basic.csv results/openai_o1mini_advanced.csv \
  -g oa \
  -f prompts/prompt_openai_o1.json \
  -m o1-mini \
  -p 1 \
  -t 120 \
  -c 0

Anthropic

要对 Claude-3 的完整问题集进行测试：

python main.py \
  -db postgres \
  -q "data/questions_gen_postgres.csv" "data/instruct_basic_postgres.csv" "data/instruct_advanced_postgres.csv" \
  -o results/claude3_classic.csv results/claude3_basic.csv results/claude3_advanced.csv \
  -g anthropic \
  -f prompts/prompt_anthropic.md \
  -m claude-3-opus-20240229 \
  -p 5 \
  -c 0

Hugging Face

要使用我们微调的 SQL 模型仅测试 10 个问题（而不是全部 200 个）：

# 使用 -W 选项忽略关于连续使用 transformers pipeline 的警告
python -W ignore main.py \
  -db postgres \
  -q "data/questions_gen_postgres.csv" "data/instruct_basic_postgres.csv" "data/instruct_advanced_postgres.csv" \
  -o results/hf_classic.csv results/hf_basic.csv results/hf_advanced.csv \
  -g hf \
  -f prompts/prompt.md \
  -m defog/llama-3-sqlcoder-8b \
  -c 0

我们还支持通过 -a 标志加载 PEFT 适配器。请注意，加载适配器会使模型的加载时间比平时稍长。

vLLM

我们还有一个 vllm 运行器，它使用 vLLM 引擎将推理过程作为一个批次整体运行。这种方式速度更快，尤其是在 num_beams 大于 1 时。您需要传入合并后的模型权重、LoRA 适配器路径（如果适用），并且模型架构必须受 vLLM 支持。以下是一个示例命令：

python -W ignore main.py \
  -db postgres \
  -q "data/questions_gen_postgres.csv" "data/instruct_basic_postgres.csv" "data/instruct_advanced_postgres.csv" \
  -o results/vllm_classic.csv results/vllm_basic.csv results/vllm_advanced.csv \
  -g vllm \
  -f "prompts/prompt.md" \
  -m defog/llama-3-sqlcoder-8b \
  -a path/to_adapter \
  -c 0

可选地，如果您正在评估使用 AWQ 量化的模型，可以添加 -qz 或 --quantized 参数。此选项仅适用于 vLLM 运行器。

使用 API 服务器运行

如果需要以不同设置运行测试，可以搭建一个 API 服务器，避免每次测试时都重新加载配置，然后依次运行测试。我们支持搭建两种类型的 API 服务器，即 vLLM API 服务器和 TGI 服务器。

我们还提供对 vLLM API 服务器的自定义修改版本，该版本仅返回生成的输出。

VLLM API 服务器

# 设置 vLLM 服务器
python -m vllm.entrypoints.api_server \
    --model defog/defog-llama-3-sqlcoder-8b \
    --tensor-parallel-size 4 \
    --dtype float16

# 用于设置支持 LoRA 适配器的 vLLM 服务器
python -m vllm.entrypoints.api_server \
    --model defog/llama-3-sqlcoder-8b \
    --tensor-parallel-size 1 \
    --dtype float16 \
    --max-model-len 4096 \
    --enable-lora \
    --max-lora-rank 64

# 使用我们修改后的 API 服务器
python utils/api_server.py \
    --model defog/llama-3-sqlcoder-8b \
    --tensor-parallel-size 4 \
    --dtype float16 \
    --max-model-len 4096 \
    --enable-lora \
    --max-lora-rank 64

# 使用 vLLM 的 OpenAI 兼容 API 服务器
export MODEL_NAME="Qwen/Qwen3-4B"
vllm serve "$MODEL_NAME" --port 8000

# 使用 API 运行器执行 SQL 评估——根据 GPU 的性能，可以将 p 和 b 的值调高
python main.py \
  -db postgres \
  -q "data/questions_gen_postgres.csv" \
  -o results/api.csv \
  -g api \
  -b 1 \
  -f prompts/prompt.md \
  --api_url "http://localhost:8000/generate" \
  --api_type "vllm" \
  -a path/to_adapter_if_applicable \
  -p 8

TGI API 服务器

您可以查阅 TGI 文档 以获取有关如何设置 TGI 服务器的更多信息。以下是一个使用预设 Docker 镜像设置 TGI 服务器并使用 API 运行器运行评估的示例命令。请注意，您需要根据可用的 GPU 数量和所选模型相应地调整分片数和模型 ID。

# 设置 TGI 服务器
model="defog/llama-3-sqlcoder-8b"
docker run --gpus all \
  --shm-size 1g \
  -p 8000:80 \
  -v /models:/models ghcr.io/huggingface/text-generation-inference:2.0 \
  --model-id "${model}" \
  --max-best-of 4 \
  --max-input-tokens 3072 \
  --sharded true \
  --num-shard 4 \
  --hostname 0.0.0.0 \
  --port 80

# 使用 API 运行器执行 SQL 评估——根据 GPU 的性能，可以将 p 和 b 的值调高。请注意，TGI 中的 CUDA 图默认针对 2 的幂次方批量大小进行了优化。
python main.py \
  -db postgres \
  -q "data/questions_gen_postgres.csv" \
  -o results/api.csv \
  -g api \
  -b 1 \
  -f prompts/prompt.md \
  --api_url "http://localhost:8000/generate" \
  --api_type "vllm" \
  -p 8

多个提示

如果您希望在一次运行中测试多个提示（以节省每次运行开始时将模型加载到 GPU 上所需的几分钟时间），可以在 --prompt_file 中指定提示文件列表（例如 -f prompts/prompt-1.md prompts/prompt-2.md prompts/prompt-3.md），并在 --output_file 中指定相应的输出文件列表（例如 -o results/results-1.csv results/results-2.csv results/results-3.csv）。提示和输出文件的数量必须相同。以下是一个示例命令：

python -W ignore main.py \
  -db postgres \
  -q "data/questions_gen_postgres.csv" \
  -o results/results_1.csv results/results_2.csv \
  -g vllm \
  -f prompts/prompt_1.md prompts/prompt_2.md \
  -m defog/sqlcoder2

虽然其他运行器也可以这样做，但在本地加载大型模型与调用始终在线的 API 相比，节省的时间最为显著。

思考型模型

如果您想使用会输出思考标记的模型，可以在运行器中传递 --enable_thinking 标志，这样我们在运行生成的查询之前会从 LLM 输出中移除思考标记。

您可以查看 run_qwen.sh 以了解如何运行思考型模型的示例。

./run_qwen.sh --thinking # 添加 --thinking 以生成思考标记

Bedrock

python -W ignore main.py \
  -db postgres \
  -q "data/questions_gen_postgres.csv" \
  -o results/llama3_70b.csv \
  -g bedrock \
  -f prompts/prompt.md \
  -m meta.llama3-70b-instruct-v1:0

Llama CPP

要使用 Llama CPP 运行评估，可以使用以下代码。在运行之前，您必须按照以下方式安装 llama-cpp-python（适用于 Apple Silicon）：

CMAKE_ARGS="-DLLAMA_METAL=on" pip install llama-cpp-python

请注意，llama-cpp-python 库目前不支持束搜索，因此结果质量会较低。

python -W ignore main.py \
  -q "data/questions_gen_postgres.csv" \
  -db postgres \
  -o "results/llama_cpp.csv" \
  -g llama_cpp \
  -f "prompts/prompt.md" \
  -m path/to/model.gguf

MLX

要使用 MLX 运行评估，可以使用以下代码。在运行之前，您必须使用 pip install mlx-lm 安装 mlx-lm 包。

请注意，MLX 目前也不支持束搜索，因此结果质量会较低。

python -W ignore main.py \
  -db postgres \
  -q "data/questions_gen_postgres.csv" \
  -o "results/mlx_llama-3-sqlcoder-8b.csv" \
  -g mlx \
  -f "prompts/prompt.md" \
  -m mlx-community/defog-llama-3-sqlcoder-8b

Gemini

在运行之前，您需要通过 export GEMINI_API_KEY=<your_api_key> 设置您的凭据。然后，使用 pip install google-generative-ai 安装相关包。

python main.py \
  -db postgres \
  -q "data/questions_gen_postgres.csv" "data/instruct_basic_postgres.csv" "data/instruct_advanced_postgres.csv" \
  -o "results/gemini_flash_basic.csv" "results/gemini_flash_basic.csv" "results/gemini_flash_advanced.csv" \
  -g gemini \
  -f "prompts/prompt_gemini.md" "prompts/prompt_gemini.md" "prompts/prompt_gemini.md" \
  -m gemini-2.0-flash-exp \
  -p 10

Mistral

在运行之前，您必须先在 Mistral 上注册账户并获取 API 密钥，然后将其存储为 export MISTRAL_API_KEY=<your_api_key>。之后，使用 pip install mistralai 安装 mistralai。

python -W ignore main.py \
  -db postgres \
  -q "data/questions_gen_postgres.csv" \
  -o "results/results.csv" \
  -g mistral \
  -f "prompts/prompt_mistral.md" \
  -m mistral-medium \
  -p 5 \
  -n 10

Bedrock

在运行之前，您需要导出以下环境变量，以便 boto3 客户端能够正常工作：

• AWS_ACCESS_KEY_ID
• AWS_SECRET_ACCESS_KEY
• AWS_DEFAULT_REGION

python3 main.py \
  -db postgres \
  -q data/instruct_basic_postgres.csv data/instruct_advanced_postgres.csv data/questions_gen_postgres.csv \
  -o results/bedrock_llama_70b_basic.csv results/bedrock_llama_70b_advanced.csv results/bedrock_llama_70b_v1.csv \
  -g bedrock \
  -f prompts/prompt_cot_postgres.md \
  -m meta.llama3-70b-instruct-v1:0 \
  -c 0 \
  -p 10

Deepseek

在运行此命令之前，您必须先在 Deepseek 上注册一个账户，并获取 API 密钥，然后使用 export DEEPSEEK_API_KEY=<your_api_key> 将其存储起来。接着，通过 pip install openai 安装 openai 库。之后即可运行以下命令：

Deepseek Chat

python main.py
-db postgres
-q "data/questions_gen_postgres.csv" "data/instruct_basic_postgres.csv" "data/instruct_advanced_postgres.csv"
-o results/deepseek_classic.csv results/deepseek_basic.csv results/deepseek_advanced.csv
-g deepseek
-f prompts/prompt_openai.json
-m deepseek-chat
-p 5
-c 0

Deepseek Reasoner

python main.py
-db postgres
-q "data/questions_gen_postgres.csv" "data/instruct_basic_postgres.csv" "data/instruct_advanced_postgres.csv"
-o results/deepseek_classic.csv results/deepseek_basic.csv results/deepseek_advanced.csv
-g deepseek
-f prompts/prompt_openai_o1.json
-m deepseek-reasoner
-p 5
-c 0

Together

在运行此命令之前，您必须先在 Together.ai 上注册一个账户，并获取 API 密钥，然后使用 export TOGETHER_API_KEY=<your_api_key> 将其存储起来。接着，通过 pip install together 安装 together 库。之后即可运行以下命令：

python3 main.py \
  -db postgres \
  -q data/instruct_basic_postgres.csv data/instruct_advanced_postgres.csv data/questions_gen_postgres.csv \
  -o results/together_llama_70b_basic.csv results/together_llama_70b_advanced.csv results/together_llama_70b_v1.csv \
  -g together \
  -f prompts/prompt_together.json \
  -m "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo" \
  -c 0 \
  -p 10

CLI 标志

您可以在命令行中使用以下标志来更改评估运行的配置。

数据相关参数

CLI 标志	描述
-q, --questions_file	包含测试问题和正确查询的 CSV 文件。如果未设置，则默认使用相应的 questions_gen_<db_type>.csv 文件。建议始终将您的问题文件名以 *<db_type>.csv 结尾，以确保查询与所选数据库类型兼容。
-n, --num_questions	使用此选项可限制要测试的问题总数。
-db, --db_type	要在其上运行查询的数据库类型。目前支持的类型包括 postgres 和 snowflake。
-d, --use_private_data	使用此选项从您自己的私有数据集读取数据。
-dp, --decimal_points	使用此选项指定结果应四舍五入的小数位数。默认值为 None。

模型相关参数

CLI 标志	描述
-g, --model_type	所使用的模型类型。请确保此选项与实际使用的模型匹配。当前在 main.py 中定义的选项包括：oa 表示 OpenAI 模型，anthropic 表示 Anthropic 模型，hf 表示 Hugging Face 模型，vllm 表示 vllm 运行器，api 表示 API 端点，llama_cpp 表示 llama cpp，mlx 表示 mlx，bedrock 表示 AWS Bedrock API，together 表示 together.ai 的 API
-m, --model	将被测试并用于生成查询的模型。OpenAI 模型的一些选项包括聊天模型 gpt-4o 和 o3-mini。Anthropic 的选项则包括最新的 Claude-3 系列模型（例如 claude-3-opus-20240229）。对于 Hugging Face 和 VLLM 模型，只需使用您所选模型的路径即可（例如 defog/sqlcoder）。
-a, --adapter	您正在使用的相关适配器模型的路径。仅适用于 hf_runner。
--api_url	您希望将提示发送到的自定义 API 的 URL。仅当 model_type 为 api 时使用。
-qz, --quantized	表示该模型是否为 AWQ 量化模型。仅适用于 vllm_runner。

推理技术相关参数

CLI 标志	描述
-f, --prompt_file	包含用于生成查询的提示的 Markdown 文件。您可以传入一个提示列表，按顺序进行测试，而无需重新加载脚本。
-b, --num_beams	指定推理时要使用的束搜索束数。仅适用于 hf_runner、vllm_runner 和 api_runner。
-c, --num_columns	列数，默认为 20。若不希望剪枝列，请将其设置为 0。
-s, --shuffle_metadata	打乱元数据，默认为 False。此选项会打乱模式中表的顺序以及每张表内列的顺序，但不会在表之间移动列（以保持数据库结构）。
-k, --k_shot	当您希望在提示中包含 k-shot 示例时使用。请确保您的 questions_file 中存在 k_shot_prompt 列。
--cot_table_alias	（实验性）当您希望在实际 SQL 生成之前包含思维链指令时使用。允许值为 instruct。如果使用 instruct，请确保您的提示文件中存在占位符 {cot_instructions}。instruct 将使模型生成思维链表别名。

执行相关参数

CLI 标志	描述
-o, --output_file	用于存储结果的输出 CSV 文件。您需要传递与提示文件数量相同的输出文件路径。
-p, --parallel_threads	用于生成和处理查询的并行工作线程数
-t, --timeout_gen	查询生成超时前的秒数。默认值为 30.0 秒。
-u, --timeout_exec	数据库上查询执行超时前的秒数。默认值为 10.0 秒。
-v, --verbose	在命令行中打印详细信息。
--upload_url	（可选）您希望报告结果的 URL。提供此 URL 的服务器必须具备与 utils/webserver.py 中示例服务器类似的功能。
--run_name	（可选）本次运行的名称，用于日志记录

检查结果

上传 URL

如果您希望启动 Google Cloud Function 来接收结果，可以使用 --upload_url 标志指定您希望报告结果的 URL。在使用该标志运行评估代码之前，您需要创建一个在提供的 URL 上提供服务的服务器。我们已在 results_fn_bigquery 和 results_fn_postgres 文件夹中提供了两个示例 Cloud Function 端点，分别用于写入 BigQuery 或 PostgreSQL。您也可以实现自己的服务器来接收类似的参数。在部署任何 Cloud Function 之前，您需要通过复制 .env.yaml.template 并将其重命名为 .env.yaml，然后填写相关字段来设置环境变量。对于 BigQuery Cloud Function，您还需要将服务账户的 key.json 文件放在同一文件夹中，并将文件名填入 .env.yaml 文件中的 CREDENTIALS_PATH 字段。

完成上述步骤后，您可以部署 Google Cloud Function：

# 用于上传到 BigQuery
gcloud functions deploy results_bigquery \
  --source results_fn_bigquery \
  --entry-point bigquery \
  --env-vars-file results_fn_bigquery/.env.yaml \
  --runtime python311 \
  --memory 512MB \
  --trigger-http \
  --allow-unauthenticated \
  --gen2

# 用于上传到 PostgreSQL
gcloud functions deploy results_postgres \
  --source results_fn_postgres \
  --entry-point postgres \
  --env-vars-file results_fn_postgres/.env.yaml \
  --runtime python311 \
  --memory 512MB \
  --trigger-http \
  --allow-unauthenticated \
  --gen2

Cloud Function 的名称是 gcloud functions deploy 后面的部分（本例中为 results_bigquery），您可以通过运行 gcloud functions logs read results_bigquery 来查看该函数的日志。

随后，您可以使用 --upload_url 标志运行评估代码，将结果报告给 Cloud Function。Cloud Function 那么会将结果写入相应的数据库。

python main.py \
  -db postgres \
  -o results/test.csv \
  -g oa \
  -f prompts/prompt_openai.json \
  -m gpt-4o-mini \
  -n 1 \
  --upload_url <您的 Cloud Function URL>

如果您希望始终将结果报告到 upload_url，即使未显式提供，也可以在环境变量中将其设置为 SQL_EVAL_UPLOAD_URL。

在本地测试函数

如果您想修改这些函数并在本地进行测试，可以运行以下示例命令，在本地部署函数并触发 OpenAI 运行器：

functions-framework --target bigquery --source results_fn_bigquery --debug
python main.py \
  -db postgres \
  -o results/test.csv \
  -g oa \
  -f prompts/prompt_openai.json \
  -m gpt-4o-mini \
  -n 1 \
  --upload_url http://127.0.0.1:8080/

其他

我们欢迎对本项目做出贡献，特别是：

• 数据集
  • 添加新的数据库模式/数据
• 框架代码
  • 改进现有生成器/运行器（例如添加新指标）

更多信息请参阅 CONTRIBUTING.md。

sql-eval 快速上手指南

sql-eval 是 Defog 开源的 SQL 生成评估工具，用于测试大语言模型（LLM）生成 SQL 查询的准确性。它基于 Spider 数据集模式，但使用了新的手工筛选问题集，支持多种数据库后端（Postgres, Snowflake, BigQuery 等）和多种模型接口（OpenAI, Anthropic, Hugging Face, vLLM 等）。

环境准备

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

• 操作系统：Linux 或 macOS（Windows 用户建议使用 WSL2 或 Docker Desktop）。
• 基础工具：
  • Git
  • Python 3.8+
  • Docker（推荐用于快速部署数据库）
  • psql 客户端（用于 Postgres 数据导入）
• 前置知识：熟悉命令行操作、Docker 基本用法及 Python 数据库（如 pandas）。
• API Key：如果测试商业模型，需准备好 OPENAI_API_KEY 或 ANTHROPIC_API_KEY 环境变量。

提示：国内用户在安装 Python 依赖时，建议指定清华或阿里镜像源以加速下载： pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

1. 克隆数据仓库并安装依赖

首先克隆存放数据库模式和数据的 defog-data 仓库，并安装相关 Python 库。

git clone https://github.com/defog-ai/defog-data.git
cd defog-data
pip install -r requirements.txt
pip install -e .

(可选) 如果使用元数据剪枝功能（-c 参数大于 0），还需安装 spacy 模型：

python -m spacy download en_core_web_sm

2. 启动 Postgres 数据库实例

使用 Docker 快速启动一个持久化的 Postgres 实例。

创建容器（只需执行一次）：

mkdir data/postgres data/export
docker create --name postgres-sql-eval -e POSTGRES_PASSWORD=postgres -p 5432:5432 -v $(pwd)/data/postgres:/var/lib/postgresql/data -v $(pwd)/data/export:/export postgres:16-alpine

启动容器：

docker start postgres-sql-eval

注意：请确保本地没有其他进程占用 5432 端口。数据将持久化存储在 data/postgres 目录中。

3. 导入评估数据

将测试所需的元数据和样本数据导入到刚启动的 Postgres 实例中。

# 确保在 defog-data 目录下
export DBPASSWORD="postgres"
export DBUSER="postgres"
export DBHOST="localhost"
export DBPORT=5432
./setup.sh

基本使用

以下示例展示如何使用 OpenAI 模型对 10 个问题进行并行评估。

前提：设置环境变量

export OPENAI_API_KEY="your_api_key_here"

运行测试命令： 该命令将针对 Postgres 数据库，使用 o3-mini 模型，并行度为 5，测试部分问题集。

python main.py \
  -db postgres \
  -q "data/questions_gen_postgres.csv" "data/instruct_basic_postgres.csv" "data/instruct_advanced_postgres.csv" \
  -o results/openai_classic.csv results/openai_basic.csv results/openai_advanced.csv \
  -g oa \
  -f prompts/prompt_openai.json \
  -m o3-mini \
  -p 5 \
  -c 0

参数说明：

• -db: 数据库类型 (postgres, snowflake 等)。
• -q: 输入的问题文件路径（支持多个）。
• -o: 输出结果文件路径。
• -g: 运行器类型 (oa 代表 OpenAI, anthropic, hf, vllm 等)。
• -f: Prompt 模板文件。
• -m: 模型名称。
• -p: 并行请求数量。
• -c: 元数据剪枝阈值（0 表示禁用）。

其他模型示例：

• Anthropic (Claude 3):

  export ANTHROPIC_API_KEY="your_key"
  python main.py -db postgres -q "data/questions_gen_postgres.csv" ... -g anthropic -m claude-3-opus-20240229 -p 5 -c 0
  # 注意：实际使用时请补全 -o 和 -f 参数，参考上文完整命令
  

• 本地 Hugging Face 模型:

  python -W ignore main.py \
    -db postgres \
    -q "data/questions_gen_postgres.csv" "data/instruct_basic_postgres.csv" "data/instruct_advanced_postgres.csv" \
    -o results/hf_classic.csv results/hf_basic.csv results/hf_advanced.csv \
    -g hf \
    -f prompts/prompt.md \
    -m defog/llama-3-sqlcoder-8b \
    -c 0
  

运行结束后，评估结果（包括精确匹配率、子集匹配率、Token 消耗等）将保存至指定的 -o CSV 文件中。