SQL生成评估

这个仓库包含了Defog用于评估生成SQL的代码。它基于Spider的模式，但有一组新的手动选择的问题和按查询类别分组的查询。要深入了解我们创建这种评估方法的过程，请参阅这里。

简介

我们的测试程序包括以下步骤。对于每个问题/查询对：

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

开始使用

这是一套全面的说明，假设您对命令行、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
python -m spacy download en_core_web_sm
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:14-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作为如何将数据导入数据库的示例。我们不规定任何特定的文件夹结构，由您决定如何组织数据，只要您能轻松地将其导入数据库即可。
• 要使用我们的元数据裁剪工具，您需要定义以下内容：
  • 一种加载嵌入的方法。在我们的情况下，我们调用defog-data的supplementary模块中的load_embeddings函数来加载数据库名称到二维嵌入矩阵（样本数 x 嵌入维度）和每行/样本相关文本元数据的元组的字典。如果您想了解我们如何生成这个元组，可以参考defog-data仓库中的generate_embeddings。
  • 一种加载与各种命名实体相关的列的方法。在我们的情况下，我们调用一个columns_ner字典，将数据库名称映射到一个嵌套字典，该嵌套字典将每种命名实体类型映射到与该命名实体类型相关的列元数据字符串列表。您可以参考原始数据了解我们如何生成这个字典的示例。
  • 一种定义表之间可连接列的方法。在我们的情况下，我们调用一个columns_join字典，将数据库名称映射到表元组到列名元组的嵌套字典。您可以参考原始数据了解我们如何生成这个字典的示例。

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

• 通过运行pip install -e .将您的数据库作为依赖项安装（-e选项可以自动包含编辑而无需重新安装）
• 用您自己导入的函数和变量替换prune_metadata_str中相关的函数调用和变量。注意，您的包/模块名可能不是defog_data_private.supplementary，请相应修改。

需要注意的几点：

• 如果您不用数据填充数据库（即只创建表而不插入数据），大多数时候您会返回空的数据框（无论生成的查询是否是您想要的），这会导致结果总是匹配并产生大量误报。因此，您可能需要考虑用一些有意义的数据填充数据库，这些数据在查询与您想要的不同时会返回不同的结果。
• 如果在您的私有数据上进行测试，您还需要将问题文件更改为指向您自己的问题文件（根据您的数据库模式定制）。

查询生成器

要使用我们的框架测试您自己的查询生成器，您需要扩展Query Generator并实现generate_query方法以返回感兴趣的查询。我们为每个问题/查询对创建一个新类，以在并发运行时将每对的运行时状态与其他对隔离。您也可以参考OpenAIQueryGenerator，它实现了Query Generator并使用简单的提示向OpenAI的API发送消息。您可以根据自己的需求进行扩展。

如果有对所有查询生成器通用的函数，可以将它们放在utils文件夹中。如果需要包含特定的详细模板（例如用于提示测试），可以将它们存储在prompts文件夹中，稍后再导入。能够在一个中心位置对提示进行版本控制对我们团队来说是一个提高生产力的胜利。

运行器

实现查询生成器后，下一个抽象层是运行器。运行器调用查询生成器，负责处理每个问题/查询对的工作配置（例如并行化/批处理/选择的模型等）。

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

在使用现有运行器测试您自己的查询生成器时，您可以用自己的查询生成器类替换运行器代码中的qg_class。

运行测试

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 gpt-4-turbo \
  -p 5 \
  -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

# 使用API运行器运行sql-eval - 根据您的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-eval - 根据您的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则不然。

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

在运行之前,您必须在Google AI创建一个账户,并使用export GOOGLE_APPLICATION_CREDENTIALS=</path/to/service_account.json>设置您的凭证。然后,使用pip install vertexai google-cloud-aiplatform安装这些包。

python -W ignore main.py \
  -db postgres \
  -q "data/questions_gen_postgres.csv" \
  -o "results/gemini_pro.csv" \
  -g gemini \
  -f "prompts/prompt_gemini.md" \
  -m gemini-pro \
  -p 1 \
  -n 5

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

CLI标志

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

数据相关参数

模型相关参数

推理技术相关参数

执行相关参数

检查结果

为了更好地理解您的查询生成器的性能，您可以探索为您关心的各种指标生成和汇总的结果。

上传URL

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

云函数的名称是gcloud functions deploy之后的内容（在本例中为results_bigquery），您可以通过运行gcloud functions logs read results_bigquery来检查函数的日志。

然后，您可以使用--upload_url标志运行评估代码，将结果报告给云函数。云函数随后会将结果写入相关数据库。

python main.py \
  -db postgres \
  -o results/test.csv \
  -g oa \
  -f prompts/prompt_openai.json \
  -m gpt-3.5-turbo-0613 \
  -n 1 \
  --upload_url <your 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-3.5-turbo-0613 \
  -n 1 \
  --upload_url http://127.0.0.1:8080/

其他

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

• 数据集
  • 添加新的数据库模式/数据
• 框架代码
  • 新的查询生成器/运行器（分别在query_generators和eval文件夹中）
  • 改进现有生成器/运行器（例如添加新的指标）

请参阅CONTRIBUTING.md以获取更多信息。