A Bert-based NL2SQL Model: A Concise Baseline

基于Bert的NL2SQL模型：一个简明的Baseline 苏剑林 · 2019-06-29 · Machine Learning language model competition SQL attention

In a previous post, "When Bert Meets Keras: Perhaps the Easiest Way to Get Started with Bert", I introduced three NLP examples built by fine-tuning Bert, giving us a taste of both Bert's power and Keras' convenience. In this post, I'll add one more example: a Bert-based NL2SQL model.

The "NL" in NL2SQL stands for Natural Language, so NL2SQL means "converting natural language into SQL statements." It has attracted a fair amount of research in recent years and is one of the more practically useful tasks in the field of AI. My motivation for building this model was the First Chinese NL2SQL Challenge held by my company this year:

The First Chinese NL2SQL Challenge uses tabular data from the finance and general domains as its data source, and provides annotated pairs of natural language and matching SQL statements built on top of it. The goal is for participants to train a model that can accurately convert natural language into SQL.

This NL2SQL competition is one of the bigger NLP events this year — a lot of manpower and resources went into promoting it ahead of time, and the prize money is quite generous. The only catch is that NL2SQL itself is a relatively niche research area, so it was never going to be a huge hit. To help lower the barrier to entry, the organizers released a Baseline written in PyTorch.

With the mindset of "how can a baseline be complete without a Keras version," I took some time to build my own version of this competition using Keras. To simplify the model and improve performance, I also loaded a pretrained Bert model, and the result is this post. more

Sample Data

Each data sample looks like this:

{
    "table_id": "a1b2c3d4", # 相应表格的id
    "question": "世茂茂悦府新盘容积率大于1，请问它的套均面积是多少？", # 自然语言问句
    "sql":{ # 真实SQL
        "sel": [7], # SQL选择的列 
        "agg": [0], # 选择的列相应的聚合函数, '0'代表无
        "cond_conn_op": 0, # 条件之间的关系
        "conds": [
            [1, 2, "世茂茂悦府"], # 条件列, 条件类型, 条件值，col_1 == "世茂茂悦府"
            [6, 0, "1"]
        ]
    }
}

# 其中条件运算符、聚合符、连接符分别如下
op_sql_dict = {0:">", 1:"<", 2:"==", 3:"!="}
agg_sql_dict = {0:"", 1:"AVG", 2:"MAX", 3:"MIN", 4:"COUNT", 5:"SUM"}
conn_sql_dict = {0:"", 1:"and", 2:"or"}

Each sample is also paired with a data table, which includes all the column names of that table along with the corresponding data records. In principle, the generated SQL statement should be executable against the corresponding data table and should return a valid result.

As you can see, although it's called NL2SQL, the organizers have actually formatted the SQL statements very cleanly, which greatly simplifies the task. Take the sel field, for example — it's essentially a multi-label classification problem, except that the classes can change at any time, because the classes here actually correspond to the columns of the data table, and the table (and its meaning) differs from sample to sample. So we need to dynamically encode a class vector based on the table's column names. As for agg, it corresponds one-to-one with sel, and its classes are fixed, while cond_conn_op is a single-label classification problem.

The final piece, conds, is relatively more complex. It requires combining sequence tagging with classification, since we need to simultaneously determine which column is the condition, what the relational operator is, and what the condition value is. It's worth noting that the condition value isn't always literally a span of the question — it could be a formatted version of it. For instance, the question might contain "16年" ("year 16"), while the condition value could be the formatted "2016." However, since the organizers guarantee that the generated SQL is executable against the corresponding table and returns a valid result, whenever the condition operator is "==", the condition value is guaranteed to appear among the values of the corresponding column in the data table. For example, in the sample shown above, the first column of the data table is guaranteed to contain the value "世茂茂悦府" ("Shimao Maoyue Mansion"), and we can use this fact to correct our predictions.

Model Architecture

Before diving into the model, readers might want to think for themselves about how they would approach this. Only after thinking it through will you appreciate where the real difficulties lie, and only then will the key ideas behind this model's design tricks make sense.

Here is a diagram of the model described in this post:

Diagram of the NL2SQL model in this post. It mainly consists of 4 different classifiers: a sequence tagger

For a SQL query, the most fundamental thing is determining which columns should be selected. Since the meaning of the columns differs from table to table, we concatenate the question sentence together with all of the table's headers and feed them all into Bert at once for joint encoding, where each header is also treated as a separate sentence, wrapped by [CLS] ... [SEP].

After passing through Bert, we obtain a series of encoded vectors, and it's up to us to decide how to use them. We take the vector corresponding to the first [CLS] as the sentence vector for the whole question, and use it to predict the connective operator for conds. Each subsequent [CLS] vector is treated as the encoding vector for a given header, and we use it to predict whether the column represented by that header should be selected. There's a trick here: since, as mentioned, besides predicting sel we also need to predict the corresponding agg — and agg has 6 classes representing different aggregation operations — we simply add one more class. This 7th class represents "this column is not selected." That way, each column becomes a 7-way classification problem: if it falls into one of the first 6 classes, the column is selected and we simultaneously get the predicted agg; if it falls into the 7th class, the column is not selected.

That leaves the more complicated part, conds, which looks something like where col_1 == value_1 — we need to identify col_1, value_1, and the operator ==. Predicting conds is done in two steps: first we predict the condition value, then the condition column. Predicting the condition value is essentially a sequence-tagging problem. There are 4 possible operators for the condition value, and again we add one more class to make 5 total, where the 5th class means "this character is not tagged" and any other class means it is tagged. This lets us predict both the condition value and its operator. What remains is predicting which column the condition value belongs to: we take the character vectors of the tagged value, compute a similarity score against each header vector, and apply softmax. The similarity computation I use here is about as simple as it gets — I just concatenate the character vector and the header vector, pass them through a dense layer, and then a Dense(1). I kept it this simple partly because the main goal of this post is to provide a workable demo rather than a polished production system (leaving some room for readers to improve on it), and partly because a more elaborate approach would quickly run into GPU memory limits and OOM errors.

By the way, the model in this post is something I cooked up on my own based on the competition task — if readers want to discuss mainstream NL2SQL models with me, I'm afraid I won't be much help there. My apologies.

Experimental Results

The code for this model's post is located at:

https://github.com/bojone/bert_in_keras/blob/master/nl2sql_baseline.py

Note: if running this code throws an error, you may need to modify Keras' backend/tensorflow_backend.py by changing, inside the sparse_categorical_crossentropy function, the line that originally reads

logits = tf.reshape(output, [-1, int(output_shape[-1])])

to

logits = tf.reshape(output, [-1, tf.shape(output)[-1]])

I've already submitted this fix to the official repo, and it has been merged (see here), so future versions should include this feature automatically.

As I keep saying, if you've carefully studied the competition data and thought through the task independently, the model described above should be pretty easy to follow. And the fact that such a simple model can achieve decent results is thanks to Bert's powerful semantic encoding ability. On the offline validation set, the full-match rate of SQL statements generated by this model is around 58%. The official evaluation metric is (full-match rate + execution-match rate) / 2 — meaning that even if you generate an SQL statement different from the annotated answer, as long as executing it produces the same result, you still get partial credit.

So the final score is bound to be higher than 58%; I'd estimate around 65%. Looking at the current leaderboard, a score of 65% would still place you near the top (the current first place is at 70%). Since employees of my company aren't allowed to participate in the leaderboard evaluation, I haven't submitted any entries myself, so I don't know exactly how it would score online. Anyone interested is welcome to try submitting it themselves.

Oh, and to run this script, you'll want a GPU at least as powerful as a 1080 Ti. If you don't have that much GPU memory, try lowering maxlen and the batch size. Also, there are currently two Chinese pretrained Bert weight sets available: the official version and the HIT (Harbin Institute of Technology) version. Their final performance is similar, but the HIT version converges faster.

Looking at the model as a whole, the trickiest part of the implementation is carefully handling the various masks. In the script above, xm, hm, cm are the three mask variables, which exist to remove the effect of padding during training. Note that masking isn't unique to Keras — whether you're using TensorFlow or PyTorch, you'll in principle need to handle masks carefully. If you really can't make sense of the masking parts, feel free to leave a comment and ask, but before asking, please try to answer the following question:

What does the sequence look like before masking? Which positions change after masking, and what do they become?

Answering this question demonstrates that "you already understand what operation the program is performing, you just don't understand why it's performing that operation." But if you can't even tell what the operation itself is doing, I'm afraid it will be hard for us to communicate (surely you can at least tell which part changed...). In that case, it's best to go study Keras or TensorFlow a bit more before diving in — there's no shortcut to mastering this in one go.

Pre- and Post-processing

For the model itself, the tricky part of the implementation is masking. But if you look at the script as a whole, most of the code is actually devoted to reading and preprocessing the data and post-processing the results — building the model itself only takes about twenty lines (once again, hats off to how concise Keras is and how powerful Bert is).

We mentioned earlier that the condition value doesn't necessarily appear in the question — so how do we use character-level tagging on the question to extract it anyway?

My approach: if the condition value doesn't appear in the question, I tokenize the question and enumerate all its 1-grams, 2-grams, and 3-grams, then find the n-gram closest to the condition value to use as the tagged span. At prediction time, if we find an n-gram as the condition value and the operator is "==", we check whether this n-gram actually appears in the database. If it does, we keep it as is; if not, we look for the closest matching value in the database instead.

The corresponding logic is all reflected in the code — feel free to read through it in detail.

Summary

Everyone's welcome to give it a try~

The First Chinese NL2SQL Challenge

https://tianchi.aliyun.com/markets/tianchi/zhuiyi_cn

Wishing everyone great results!