關(guān)鍵詞：Text2SQL扁位，BERT
前言
Text2SQL是指將自然語言轉(zhuǎn)化為類SQL查詢語句匆背，使得用戶的查詢文本可以直接實現(xiàn)和數(shù)據(jù)庫交互，本文介紹一種以BERT為基礎(chǔ)模型墨技，通過模板填充來實現(xiàn)的Text2SQL算法和產(chǎn)品化。
內(nèi)容摘要
- Text2SQL任務說明
- 模板填充的思路
- 條件列選擇子模型搭建（where col）
- 條件合并類型子模型搭建（and挎狸，or）
- 條件值匹配子模型搭建（col扣汪，value）
- 排序條件子模型搭建
- limit子模型搭建
- 整體模型pipeline
- 基于Streamlit快速產(chǎn)品化
Text2SQL任務說明
本文的任務是將用戶的查詢語言自動翻譯為ElasticSearch的DSL查詢語句，然后以DSL語句查詢ElasticSearch返回結(jié)果锨匆。業(yè)務場景為將上市公司的標簽存儲在ElasticSearch中崭别，通過自然語言挑選出對應的股票，實現(xiàn)“一句話選股”恐锣，該產(chǎn)品在市面上的APP已經(jīng)廣泛實現(xiàn)茅主，本文是對其的一個簡單功能復現(xiàn)。
注意本文需要轉(zhuǎn)化的自然語言任務僅為挑選出企業(yè)侥蒙，并且展示出所需要的列暗膜，不設(shè)計復雜的聚合和聚合函數(shù)，比如groupby鞭衩，max学搜，sum不在本次的考慮范圍內(nèi)。
模板填充的思路
模板填充生成查詢語句的思路類似于完型填空论衍，不論是SQL還是DSL查詢語句瑞佩，任何查詢都有固定部分，那把不固定的部分遮蔽掉坯台，單獨預測遮蔽掉的關(guān)鍵詞即可補全SQL炬丸，我們以ElasticSearch的DSL舉例，模板填充的形式如下
無底色部分代表固定的DSL形式蜒蕾，有底層部分代表可變的關(guān)鍵詞部分稠炬，不同的顏色代表不同的DSL模塊：
- must：篩選條件的組合形式，must代表且咪啡，should代表或
- term首启， range：篩選條件的比較類型，term是相等匹配撤摸，range是值大小比對
- col：篩選條件用到的列毅桃，以及select需要的列
- gte：比較符，gte代表大于等于
- desc：排序類型准夷，desc代表降序
- size：limit條件钥飞，代表最大輸出多少數(shù)據(jù)
因此模板填充的任務就是把可變的關(guān)鍵詞預測出來，從而拼接出最終的DSL語句衫嵌。
條件列選擇子模型搭建（where col）
第一步是要預測出條件列读宙，只有先預測出條件列，后續(xù)才能進一步完成列的條件匹配和值匹配楔绞。我們采用BERT的句子對的方式结闸，將每一列和用戶輸入的查詢語句拼接起來掖棉，采用[CLS]位置作為表征，預測二分類是否匹配膀估。
舉例如果有100個候選列幔亥，則一條查詢語句需要預測100條樣本，每條樣本是一個二分類察纯，最終可以預測出多個列或者一個列帕棉，也可以沒有列。
條件合并類型子模型搭建（and饼记，or）
條件合并類型預測是預測出有多個條件時的組合類型香伴，且，或具则，甚至沒有即纲，是一個三分類問題。只需要對原始的查詢語句使用BERT的[CLS]預測即可
條件值匹配子模型搭建（col博肋，value）
條件值匹配是一個三元組的預測低斋，涉及元素為col列，value值匪凡，和比較符膊畴，col列是條件列選擇子模型搭建（where col）該任務的輸出，及條件值匹配是條件列選擇的后續(xù)模型病游，由于col列已經(jīng)被上游任務預測出來了唇跨，所以本任務的目標是先抽取出value值，并且搭配比較符做匹配衬衬。
在此我們將col列分為real數(shù)值列和type類型枚舉列买猖，因為這兩種情況對應的value值和比較符都不一樣。
- real數(shù)值列：需要通過配置提前獲得預測出的select列的類型為real滋尉，此時比較符設(shè)置為“>”玉控，“<”，“=”兼砖，“!=”奸远，“>=”既棺，“<=”一共六種情況讽挟，value值一定是數(shù)字，因此通過正則表達式加中文單位工具轉(zhuǎn)換可以抽取出來
一個使用Python腳本抽取輸入中的數(shù)字的案例如下
import re
import cn2an
def parse(text: str):
res = []
tmp = re.findall("[\d一二三四五六七八九十百千萬億零.]+", text)
for t in tmp:
value = ""
try:
value = cn2an.cn2an(t, "strict")
except:
try:
value = cn2an.cn2an(t, "normal")
except:
try:
value = cn2an.cn2an(t, "smart")
except:
pass
if value:
value = int(value) if int(value) == value else value
res.append(value)
return res
if __name__ == '__main__':
print(parse("手里有三個蘋果丸冕，單價0.56元一個耽梅，總共一百2十三萬元"))
# 抽取結(jié)果 [3, 0.56, 1, 1230000]
- type類型列：需要通過配置提前獲得預測出的select列的類型為type，此時value值為數(shù)據(jù)庫中該字段的枚舉值全集胖烛，比如該列為省份眼姐，則需要窮舉出所有的省份做匹配诅迷，比較符為"="，"!="众旗。
基于以上約定需要對樣本進行構(gòu)造罢杉，即窮句出所有情況，以real列為例贡歧，所有select列和所有抽取出的value滩租，再疊加所有可能的比較符做二分類預測，同樣采用BERT的句子對做表征利朵，例如
排序條件子模型搭建
排序模型需要預測出order by的列，并且預測出是升序asc還是倒序desc绍弟，同樣通過配置提前預知需要排序的所有列技即，并且將其和升序和倒序拼接在一起輸入BERT進行句子對預測，例如
同樣以[CLS]位置來預測二分類是否匹配樟遣。
limit子模型搭建
該模塊是排序條件模型的下游任務而叼，將條件值匹配抽取到的數(shù)字和排序條件列做窮舉，同樣采用BERT做句子對分類豹悬，模型示意圖如下
同樣以[CLS]位置來預測二分類是否匹配澈歉。
整體模型pipeline
基于以上方案一共有五個BERT模型，分別是where列預測屿衅，op組合條件預測埃难，select列和值的匹配，是否有排序條件涤久，limit值預測涡尘，五個模型的pipeline流程圖如下
在最后一部需要整個5個模型的輸出改造成對應查詢語句，以ElasticSearch的DSL為例响迂，一個改造實現(xiàn)如下
def parse_to_dsl(self, op_combine, match, order_op, limit):
query = {
"query": {
"bool": {}
},
"_source": ["ent_name", "score"],
"sort": [
]
}
if order_op:
label_name = order_op[:-4]
order_name = order_op[-4:]
label_code = self.name_to_label_code[label_name]
query["sort"].append({f"{label_code}": {"order": "desc" if order_name == "降序排名" else "asc"}})
query["_source"].append(label_code)
else:
query["sort"].append({"score": {"order": "desc"}})
if limit != -1:
query["size"] = limit
op_dsl = "must"
if op_combine == "or":
op_dsl = "should"
query["query"]["bool"][op_dsl] = []
sub_queries = {}
for m in match:
cond, op, val = re.split("(=|!=|>=|<=|>|<)", m)
label_code = self.name_to_label_code[cond]
query["_source"].append(label_code)
c_type = self.cols_type[cond]
# TODO 子查詢考抄，記錄命中數(shù)
sub_query = {"query": {}}
if c_type == "real" and op not in ("=", "!="):
range_op = {"range": {f"{label_code}": {}}}
if op == ">":
range_op["range"][f"{label_code}"]["gt"] = val
elif op == ">=":
range_op["range"][f"{label_code}"]["gte"] = val
elif op == "<=":
range_op["range"][f"{label_code}"]["lte"] = val
elif op == "<":
range_op["range"][f"{label_code}"]["lt"] = val
query["query"]["bool"][op_dsl].append(range_op)
sub_query["query"] = range_op
elif c_type in ("real", "type") and op == "=":
term_op = {"term": {f"{label_code}": {"value": val}}}
query["query"]["bool"][op_dsl].append(term_op)
sub_query["query"] = term_op
elif c_type in ("real", "type") and op == "!=":
bool_op = {"bool": {"must_not": [{"term": {f"{label_code}": {"value": val}}}]}}
query["query"]["bool"][op_dsl].append(bool_op)
sub_query["query"] = bool_op
sub_queries[m] = json.dumps(sub_query, ensure_ascii=False)
return sub_queries, json.dumps(query, ensure_ascii=False)
本質(zhì)上先定義了一個基礎(chǔ)DSL模板，將要填充的槽留空蔗彤，基于模型的輸出解析改造成一個個子查詢川梅，再匯總為最終的DSL語句。
基于Streamlit快速產(chǎn)品化
以Streamlit框架為例然遏，快速不屬于個一句話選股的頁面功能贫途，實現(xiàn)如下
import streamlit as st
import pandas as pd
name = st.text_input('請輸入您的股票查詢條件', max_chars=100, help='最大長度為100字符')
if st.button('查詢'):
parse_cond, df_info, order_op, limit = query(name)
style = """
<style>
.box {
border: 1px solid #ddd;
padding: 6px 10px;
margin-top: 10px;
background-color: pink;
color:#000;
display:inline-block;
margin-right:6px;
}
</style>
"""
st.markdown(style, unsafe_allow_html=True)
conds = []
for k, v in parse_cond.items():
conds.append("<span class='box'>{}</span>".format(k + "(" + str(v) + ")"))
if order_op:
conds.append("<span class='box'>{}</span>".format(order_op + f"前{limit}" if limit != -1 else order_op))
st.markdown("".join(conds), unsafe_allow_html=True)
hit = df_info[0]
df = pd.DataFrame(df_info[1]).rename(columns=name_to_label_code)
df.insert(0, '企業(yè)名稱', df.pop('企業(yè)名稱'))
df.insert(1, '綜合指數(shù)', df.pop('綜合指數(shù)'))
st.write("共計命中:{}條企業(yè)".format(hit))
st.dataframe(df)
其中query中包裝了BERT模型的pipeline，其中streamlit彈出web應用如下
如果檢索條件待侵，比如篩選出注冊資本金大于15億丢早，屬于廣東珠海的上市企業(yè)，按照指數(shù)倒序排列，輸入后檢索結(jié)果如下
Text2SQL不僅輸出了最終符合條件的企業(yè)怨酝，也統(tǒng)計出了相關(guān)的子條件傀缩，以及子條件命中的企業(yè)數(shù)，最終符合要求的企業(yè)數(shù)农猬，從解析的條件來看赡艰，模型結(jié)果完全正確。
全文完畢斤葱。
