使用 Hugging Face 数据集查看器 API 和 Motherduck DuckDB-NSQL-7B 实现 Text2SQL

如今，集成 AI 驱动的功能，尤其是利用大语言模型 (LLM) 的功能，在文本生成、分类、图文转换、图像到图像转换等各种任务中变得越来越普遍。

开发者们越来越认识到这些应用的潜在好处，尤其是在增强核心任务方面，例如编写脚本、Web 开发，以及现在与数据交互。从历史上看，为数据分析编写富有洞察力的 SQL 查询主要是数据分析师、SQL 开发人员、数据工程师或相关领域专业人士的领域，他们都需要处理 SQL 方言语法的细微差别。然而，随着 AI 驱动解决方案的出现，情况正在发生变化。这些先进的模型为与数据交互提供了新的途径，有可能简化流程，并以更高的效率和深度揭示洞见。

如果你无需深入编码就能从数据集中发掘出有趣的洞见，那会怎么样？要收集有价值的信息，需要编写一个专门的 SELECT 语句，考虑要显示哪些列、源表、筛选所选行的条件、聚合方法以及排序偏好。这种传统方法涉及一系列命令：SELECT、FROM、WHERE、GROUP 和 ORDER。

但如果你不是经验丰富的开发者，但仍然想利用数据的力量呢？在这种情况下，寻求 SQL 专家的帮助就变得必要，这突显了在可访问性和可用性方面的差距。

这就是 AI 和 LLM 技术的突破性进展介入以弥合这一鸿沟的地方。想象一下，毫不费力地与你的数据对话，只需用简单的语言陈述你的信息需求，然后让模型将你的请求翻译成查询。

最近几个月，这一领域取得了重大进展。MotherDuck 和 Numbers Station 推出了他们最新的创新：DuckDB-NSQL-7B，这是一款专为 DuckDB SQL 设计的顶尖 LLM。这个模型的使命是什么？让用户能够毫不费力地从数据中发掘洞见。

DuckDB-NSQL-7B 最初是在 Meta 原版的 Llama-2–7b 模型基础上，使用一个涵盖通用 SQL 查询的广泛数据集进行微调的，之后又使用 DuckDB 的 text-to-SQL 对进行了进一步优化。值得注意的是，其功能超出了编写 SELECT 语句的范畴；它可以生成各种有效的 DuckDB SQL 语句，包括官方文档和扩展，使其成为一个多功能的数据探索和分析工具。

在本文中，我们将学习如何使用 DuckDB-NSQL-7B 模型、用于 Parquet 文件的 Hugging Face 数据集查看器 API 以及用于数据检索的 duckdb 来处理 text2sql 任务。

text2sql 流程

如何使用模型

• 使用 Hugging Face transformers 流水线

from transformers import pipeline

pipe = pipeline("text-generation", model="motherduckdb/DuckDB-NSQL-7B-v0.1")

• 使用 transformers 的分词器和模型

from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("motherduckdb/DuckDB-NSQL-7B-v0.1")
model = AutoModelForCausalLM.from_pretrained("motherduckdb/DuckDB-NSQL-7B-v0.1")

• 使用 llama.cpp 加载 GGUF 格式的模型

from llama_cpp import Llama

llama = Llama(
       model_path="DuckDB-NSQL-7B-v0.1-q8_0.gguf", # Path to local model
       n_gpu_layers=-1,
)

llama.cpp 的主要目标是实现 LLM 推理，只需最少的设置，并在各种硬件上（本地和云端）实现最先进的性能。我们将使用这种方法。

适用于超过 12 万个数据集的 Hugging Face 数据集查看器 API

数据是任何机器学习工作中的关键组成部分。Hugging Face 是一个宝贵的资源，提供了超过 12 万个免费开放的数据集，涵盖了 CSV、Parquet、JSON、音频和图像文件等多种格式。

Hugging Face 上托管的每个数据集都配备了全面的数据集查看器。该查看器为用户提供了基本功能，如统计洞察、数据大小评估、全文搜索功能和高效的筛选选项。这个功能丰富的界面使用户能够轻松探索和评估数据集，从而在整个机器学习工作流程中促进明智的决策。

对于这个演示，我们将使用 world-cities-geo 数据集。

world-cities-geo 数据集的数据集查看器

在幕后，Hub 中的每个数据集都由 Hugging Face 数据集查看器 API 处理，该 API 获取有用信息并提供以下功能：

• 列出数据集的拆分、列名和数据类型
• 获取数据集的大小（以行数或字节为单位）
• 下载并查看数据集中任意索引的行
• 在数据集中搜索一个词
• 根据查询字符串筛选行
• 获取关于数据的富有洞察力的统计信息
• 以 parquet 文件的形式访问数据集，以便在你喜欢的处理或分析框架中使用

在这个演示中，我们将使用最后一个功能，即自动转换的 parquet 文件。

从文本指令生成 SQL 查询

首先，下载 DuckDB-NSQL-7B-v0.1 的量化模型版本

下载模型

或者，你可以执行以下代码

huggingface-cli download motherduckdb/DuckDB-NSQL-7B-v0.1-GGUF DuckDB-NSQL-7B-v0.1-q8_0.gguf --local-dir . --local-dir-use-symlinks False

现在，让我们安装所需的依赖项

pip install llama-cpp-python
pip install duckdb

对于 text-to-SQL 模型，我们将使用具有以下结构的提示

   ### Instruction:
   Your task is to generate valid duckdb SQL to answer the following question.
   ### Input:
   Here is the database schema that the SQL query will run on:
   {ddl_create}
  
   ### Question:
   {query_input}
   ### Response (use duckdb shorthand if possible):

• ddl_create 将是作为 SQL CREATE 命令的数据集模式
• query_input 将是用户指令，用自然语言表达

所以，我们需要告诉模型关于 Hugging Face 数据集的模式。为此，我们将获取 jamescalam/world-cities-geo 数据集的第一个 parquet 文件

GET https://huggingface.co/api/datasets/jamescalam/world-cities-geo/parquet

{
   "default":{
      "train":[
         "https://huggingface.co/api/datasets/jamescalam/world-cities-geo/parquet/default/train/0.parquet"
      ]
   }
}

这个 parquet 文件托管在 Hugging Face 查看器的 refs/convert/parquet 修订版下

Parquet 文件

• 从 parquet 文件的第一行模拟一个 DuckDB 表的创建

import duckdb
con = duckdb.connect()
con.execute(f"CREATE TABLE data as SELECT * FROM '{first_parquet_url}' LIMIT 1;")

result = con.sql("SELECT sql FROM duckdb_tables() where table_name ='data';").df()
ddl_create = result.iloc[0,0]
con.close()

CREATE 模式 DDL 是

CREATE TABLE "data"(
    city VARCHAR, 
    country VARCHAR, 
    region VARCHAR,
    continent VARCHAR, 
    latitude DOUBLE, 
    longitude DOUBLE, 
    x DOUBLE, 
    y DOUBLE, 
    z DOUBLE
);

而且，如你所见，它与数据集查看器中的列相匹配

数据集列

• 现在，我们可以用 ddl_create 和 query 输入来构建提示

prompt = """### Instruction:
   Your task is to generate valid duckdb SQL to answer the following question.
   ### Input:
   Here is the database schema that the SQL query will run on:
   {ddl_create}
  
   ### Question:
   {query_input}
   ### Response (use duckdb shorthand if possible):
   """

如果用户想知道来自阿尔巴尼亚国家的城市，提示将如下所示

query = "Cities from Albania country"
prompt = prompt.format(ddl_create=ddl_create, query_input=query)

因此，将发送给 LLM 的扩展提示如下所示

### Instruction:
Your task is to generate valid duckdb SQL to answer the following question.

### Input:
Here is the database schema that the SQL query will run on:
CREATE TABLE "data"(city VARCHAR, country VARCHAR, region VARCHAR, continent VARCHAR, latitude DOUBLE, longitude DOUBLE, x DOUBLE, y DOUBLE, z DOUBLE);
  
### Question:
Cities from Albania country

### Response (use duckdb shorthand if possible):

• 是时候将提示发送给模型了

from llama_cpp import Llama

llm = Llama(
       model_path="DuckDB-NSQL-7B-v0.1-q8_0.gguf",
       n_ctx=2048,
       n_gpu_layers=50
   )
pred = llm(prompt, temperature=0.1, max_tokens=1000)
sql_output = pred["choices"][0]["text"]

输出的 SQL 命令将指向一个 data 表，但由于我们没有一个真正的表，只有一个对 parquet 文件的引用，我们将用 first_parquet_url 替换所有 data 的出现

sql_output = sql_output.replace("FROM data", f"FROM '{first_parquet_url}'")

最终的输出将是

SELECT city FROM 'https://huggingface.co/api/datasets/jamescalam/world-cities-geo/parquet/default/train/0.parquet' WHERE country = 'Albania'

• 现在，是时候最终在数据集上直接执行我们生成的 SQL 了，所以，让我们再次利用 DuckDB 的力量

con = duckdb.connect()
try:
   query_result = con.sql(sql_output).df()
except Exception as error:
   print(f"❌ Could not execute SQL query {error=}")
finally:
   con.close()

这里我们得到了结果（100 行）

执行结果（100 行）

让我们将这个结果与使用“搜索功能”搜索阿尔巴尼亚国家的数据集查看器进行比较，结果应该是一样的

阿尔巴尼亚国家的搜索结果

你也可以直接调用搜索或筛选 API 来获得相同的结果

• 使用 /search API

import requests
API_URL = "https://datasets-server.huggingface.co/search?dataset=jamescalam/world-cities-geo&config=default&split=train&query=Albania"
def query():
    response = requests.get(API_URL)
    return response.json()
data = query()

• 使用 filter API

import requests
API_URL = "https://datasets-server.huggingface.co/filter?dataset=jamescalam/world-cities-geo&config=default&split=train&where=country='Albania'"
def query():
    response = requests.get(API_URL)
    return response.json()
data = query()

我们的最终演示将是一个看起来像这样的 Hugging Face space

你可以在这里查看带代码的 notebook。

以及这里的 Hugging Face Space