Skip to content

Text-to-SQL task, with BERT for sentence embedding and GAT for syntax info extraction.

Notifications You must be signed in to change notification settings

PrecipiceBlades/BertGAT-for-Spider-Dataset

Repository files navigation

BertGAT-for-Spider-Dataset

Chang Shu, Ruitao Yi, Bo Lun

We use Spider dataset as our main dataset to fulfill the largescale and cross-domain semantic parsing text-to-SQL tasks. For this task, we propose BertGAT, which is a novel approach to the beforementioned task. To build this model, we implement Bidirectional Encoder Representations from Transformers (BERT) to pre-train deep bidirectional representations instead of the traditional Bidirectional recurrent neural networks. Fine-tuning is applied to the pre-trained BERT representations so that we can use just one extra output layer to create state-of-the-art models for wide-ranging text-to-SQL tasks. We use Syntax tree network to employ a treebased SQL generator, and use Graph Attention networks (GATs) to learn the features of syntax-tree.

Environment Setup

  1. The code uses Python 3.6, DGL 0.5.0 and Pytorch 1.4.0 GPU.
  2. Install Python dependency: pip install -r requirements.txt

Download Data, Embeddings, Scripts, and Pretrained Models

  1. Download the dataset from the Spider task website to be updated, and put tables.json, train.json, and dev.json under data/ directory.
  2. Download the pretrained Glove, and put it as glove/glove.%dB.%dd.txt
  3. Download evaluation.py and process_sql.py from the Spider github page
  4. Download preprocessed train/dev datasets and pretrained models from here. It contains: -generated_datasets/
    • generated_data for original Spider training datasets, pretrained models can be found at generated_data/saved_models
    • generated_data_augment for original Spider + augmented training datasets, pretrained models can be found at generated_data_augment/saved_models

Generating Train/dev Data for Modules

You could find preprocessed train/dev data in generated_datasets/.

To generate them by yourself, update dirs under TODO in preprocess_train_dev_data.py, and run the following command to generate training files for each module:

python preprocess_train_dev_data.py train|dev

Folder/File Description

  • data/ contains raw train/dev/test data and table file
  • generated_datasets/ described as above
  • models/ contains the code for each module.
  • evaluation.py is for evaluation. It uses process_sql.py.
  • train.py is the main file for training. Use train_all.sh to train all the modules (see below).
  • test.py is the main file for testing. It uses supermodel.sh to call the trained modules and generate SQL queries. In practice, and use test_gen.sh to generate SQL queries.
  • generate_wikisql_augment.py for cross-domain data augmentation

Training

Run train_all.sh to train all the modules. It looks like:

python train.py \
    --data_root       path/to/generated_data \
    --save_dir        path/to/save/trained/module \
    --history_type    full|no \
    --table_type      std|no \
    --train_component <module_name> \
    --epoch           <num_of_epochs>

Testing

Run test_gen.sh to generate SQL queries. test_gen.sh looks like:

SAVE_PATH=generated_datasets/generated_data/saved_models_hs=full_tbl=std
python test.py \
    --test_data_path  path/to/raw/test/data \
    --models          path/to/trained/module \
    --output_path     path/to/print/generated/SQL \
    --history_type    full|no \
    --table_type      std|no \

Evaluation

Follow the general evaluation process in the Spider github page.

About

Text-to-SQL task, with BERT for sentence embedding and GAT for syntax info extraction.

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published