在本文中,我们将向您展示如何使用LlamaIndex设置一个文本到SQL的查询管道。这将使您能够在数据上动态检索相关表,并在SQL提示中嵌入/索引每一行,从而在文本到SQL的提示中动态检索示例行。
加载和摄取数据
我们使用WikiTableQuestions数据集作为我们的测试数据集。以下是加载和摄取数据的步骤:
# 安装必要的包
%pip install llama-index-llms-openai
# 下载和解压数据集
!wget "https://github.com/ppasupat/WikiTableQuestions/releases/download/v1.0.2/WikiTableQuestions-1.0.2-compact.zip" -O data.zip
!unzip data.zip
# 加载数据
import pandas as pd
from pathlib import Path
data_dir = Path("./WikiTableQuestions/csv/200-csv")
csv_files = sorted([f for f in data_dir.glob("*.csv")])
dfs = []
for csv_file in csv_files:
print(f"processing file: {csv_file}")
try:
df = pd.read_csv(csv_file)
dfs.append(df)
except Exception as e:
print(f"Error parsing {csv_file}: {str(e)}")
提取每个表的表名和摘要
我们使用GPT-3.5提取每个表的表名和摘要。
# 创建目录存储表信息
tableinfo_dir = "WikiTableQuestions_TableInfo"
!mkdir {tableinfo_dir}
from llama_index.core.program import LLMTextCompletionProgram
from llama_index.core.bridge.pydantic import BaseModel, Field
from llama_index.llms.openai import OpenAI
class TableInfo(BaseModel):
"""表格信息"""
table_name: str = Field(..., description="表名,必须使用下划线并且不能有空格")
table_summary: str = Field(..., description="表的简短摘要/标题")
prompt_str = """\
给我一个表格摘要,格式如下:
- 表名必须唯一,并且简洁描述表格内容。
- 不要输出通用的表名(例如:table, my_table)。
排除以下表名:{exclude_table_name_list}
表格内容:
{table_str}
摘要:"""
program = LLMTextCompletionProgram.from_defaults(
output_cls=TableInfo,
llm=OpenAI(model="gpt-3.5-turbo"),
prompt_template_str=prompt_str,
)
import json
def _get_tableinfo_with_index(idx: int) -> str:
results_gen = Path(tableinfo_dir).glob(f"{idx}_*")
results_list = list(results_gen)
if len(results_list) == 0:
return None
elif len(results_list) == 1:
path = results_list[0]
return TableInfo.parse_file(path)
else:
raise ValueError(f"More than one file matching index: {list(results_gen)}")
table_names = set()
table_infos = []
for idx, df in enumerate(dfs):
table_info = _get_tableinfo_with_index(idx)
if table_info:
table_infos.append(table_info)
else:
while True:
df_str = df.head(10).to_csv()
table_info = program(
table_str=df_str,
exclude_table_name_list=str(list(table_names)),
)
table_name = table_info.table_name
print(f"Processed table: {table_name}")
if table_name not in table_names:
table_names.add(table_name)
break
else:
print(f"Table name {table_name} already exists, trying again.")
pass
out_file = f"{tableinfo_dir}/{idx}_{table_name}.json"
json.dump(table_info.dict(), open(out_file, "w"))
table_infos.append(table_info)
将数据放入SQL数据库
我们使用SQLAlchemy将所有表加载到一个SQLite数据库中。
from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer
import re
def sanitize_column_name(col_name):
return re.sub(r"\W+", "_", col_name)
def create_table_from_dataframe(df: pd.DataFrame, table_name: str, engine, metadata_obj):
sanitized_columns = {col: sanitize_column_name(col) for col in df.columns}
df = df.rename(columns=sanitized_columns)
columns = [
Column(col, String if dtype == "object" else Integer)
for col, dtype in zip(df.columns, df.dtypes)
]
table = Table(table_name, metadata_obj, *columns)
metadata_obj.create_all(engine)
with engine.connect() as conn:
for _, row in df.iterrows():
insert_stmt = table.insert().values(**row.to_dict())
conn.execute(insert_stmt)
conn.commit()
engine = create_engine("sqlite:///:memory:")
metadata_obj = MetaData()
for idx, df in enumerate(dfs):
tableinfo = _get_tableinfo_with_index(idx)
print(f"Creating table: {tableinfo.table_name}")
create_table_from_dataframe(df, tableinfo.table_name, engine, metadata_obj)
高级功能1:查询时间表检索的文本到SQL
定义模块以设置端到端的文本到SQL与表检索。
from llama_index.core.objects import SQLTableNodeMapping, ObjectIndex, SQLTableSchema
from llama_index.core import SQLDatabase, VectorStoreIndex
sql_database = SQLDatabase(engine)
table_node_mapping = SQLTableNodeMapping(sql_database)
table_schema_objs = [
SQLTableSchema(table_name=t.table_name, context_str=t.table_summary)
for t in table_infos
]
obj_index = ObjectIndex.from_objects(table_schema_objs, table_node_mapping, VectorStoreIndex)
obj_retriever = obj_index.as_retriever(similarity_top_k=3)
运行一些查询
现在我们准备好在整个管道上运行一些查询。
response = qp.run(query="What was the year that The Notorious B.I.G was signed to Bad Boy?")
print(str(response))
可能遇到的错误
- 数据解析错误:在加载CSV文件时,可能会遇到解析错误,确保数据文件格式正确。
- SQL查询错误:生成的SQL查询可能包含不存在的列名或表名,确保提示模板正确处理。
如果你觉得这篇文章对你有帮助,请点赞,关注我的博客,谢谢!