Update synthesis script.

README.md

@@ -26,6 +26,7 @@ The official repository which contains the code and pre-trained models for our p
 
 # 🔥 Updates
 
+- [**2022-10-10**]: We released the data generator code for SQL execution data synthesis. You can [check it out](data_generator) and try to synthesize your own pre-training data.
 - [**2022-04-19**]: TAPEX is officially supported by 🤗 transformers! Now you can find the [example fine-tuning script](https://github.com/huggingface/transformers/tree/main/examples/research_projects/tapex) and the [tapex model](https://huggingface.co/models?search=microsoft/tapex) on the huggingface master branch. Have fun!
 - [**2022-03-09**]: We have fixed the issue in `tapex-large`! Now you can view [here](#-run-tapex-using-huggingface) to see how to fine-tune TAPEX using 🤗 transformers and 🤗 datasets! They will be merged into the main library soon!
 - [**2022-02-20**]: Our paper is accepted by ICLR 2022! We also provided a fine-tuning script based on 🤗 transformers, which is not merged now. You can see the preview version [here](https://github.com/SivilTaram/transformers/tree/add_tapex_bis/examples/research_projects/tapex). ⚠️ It is worth noting that `tapex-large` is not well-prepared now. We found there is a strange bug in `bart-large`, which also affects `tapex-large`. Hope it is solved in the near future.
@@ -167,12 +168,34 @@ $ python examples/tableqa/run_model.py predict --resource-dir ./tapex.large.wtq
 First, you should run the following commands to install the latest lib developed for TAPEX.
 
 ```python
-pip install https://github.com/SivilTaram/datasets/archive/add-wtq-dataset.zip --user
-pip install https://github.com/SivilTaram/transformers/archive/add_tapex_bis.zip --user
+pip install datasets
+pip install transformers
 ```
 
 Then, you could find the detailed tutorial on how to reproduce our results on benchmarks at [here](https://github.com/huggingface/transformers/tree/main/examples/research_projects/tapex).
 
+# 🐣 Synthesize your own pre-training data
+
+Go to [this folder](/data_generator) and install the requirements using the provided file:
+
+```bash
+$ pip install -r requirements.txt
+```
+
+Then directly run the script `main.py`! By default, you will obtain a file with nearly 5 million examples consisting of (SQL query, Table, Answer) and the flattened fairseq input and output. More details can be found in the argument parser as below:
+
+```
+  --template_file TEMPLATE_FILE
+                        SQL query file which provides the template for synthesizing more SQL queries
+  --mode {train,dev}    train or dev for pre-training
+  --dev_id_file DEV_ID_FILE
+                        the dev id file to avoid potential data leakage
+  --instance_number INSTANCE_NUMBER
+                        the expected instance number corresponding to each template
+  --max_source_length MAX_SOURCE_LENGTH
+                        the maximum length for the flattened table plus input SQL query
+```
+
 # 💬 Citation
 
 If our work is useful for you, please consider citing our paper:

data_generator/main.py

@@ -0,0 +1,270 @@
+import json
+import os
+import random
+import shutil
+from argparse import ArgumentParser
+
+import numpy
+from tqdm import tqdm
+from transformers import AutoTokenizer
+
+from utils.common import *
+from utils.dbengine import WTQDBEngine
+from utils.template import apply_sql_on_target_table
+
+DATABASE_PATH = "squall/tables/db"
+# the temp database path will not affect the original db
+TEMP_DATABASE_PATH = "squall/tables/temp_db"
+TABLE_PATH = "squall/tables/json"
+# the delimiter to distinguish different answers in the output
+TGT_DEL = ", "
+tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large")
+
+
+def permute_table(_wtq_table_content: Dict):
+    # shuffle header orders
+    header = _wtq_table_content["header"]
+    types = _wtq_table_content["types"]
+    header_content = list(map(list, zip(*_wtq_table_content["rows"])))
+
+    header_num = len(_wtq_table_content["header"])
+    header_range = list(range(header_num))
+    random.shuffle(header_range)
+
+    # map from the original to the new
+    origin_to_shuffle = {i: header_range[i] for i in range(header_num)}
+    shuffle_header = [header[origin_to_shuffle[i]] for i in range(header_num)]
+    shuffle_types = [types[origin_to_shuffle[i]] for i in range(header_num)]
+    shuffle_content = [header_content[origin_to_shuffle[i]] for i in range(header_num)]
+    shuffle_rows = list(map(list, zip(*shuffle_content)))
+
+    return {
+        "header": shuffle_header,
+        "types": shuffle_types,
+        "rows": shuffle_rows,
+        "alias": _wtq_table_content["alias"]
+    }
+
+
+def truncate_table(_normalized_table: Dict, max_table_length):
+    """
+    Given a normalized table and the maximum flatten length, we should truncate the table if necessary
+    """
+    assert "header" in _normalized_table
+    assert "rows" in _normalized_table
+    header_string = "col : " + " | ".join(_normalized_table["header"]) + " "
+    header_tokens = tokenizer.tokenize(header_string, add_special_tokens=False)
+    # split all cell values into tokens and see how many can be adapt
+    used_token_len = len(header_tokens)
+    # remaining length
+    remain_token_len = max_table_length - used_token_len
+
+    value_string = ""
+    for _, row_example in enumerate(_normalized_table["rows"]):
+        # generally we do not want to make
+        value_string += "row " + str(10) + " : "
+        row_cell_values = [str(cell_value) if isinstance(cell_value, int) else cell_value.lower()
+                           for cell_value in row_example]
+        value_string += " | ".join(row_cell_values) + " "
+    value_token_len = len(tokenizer.tokenize(value_string))
+
+    # calc a drop rate
+    drop_rate = 1.0 - remain_token_len / value_token_len
+    current_chunk_remain_size = remain_token_len
+    truncate_table_contents = _normalized_table
+    keep_rows = []
+    for ind, row_example in enumerate(_normalized_table["rows"]):
+        drop_cur_row = random.random() < drop_rate
+        if drop_cur_row:
+            # do not append this row
+            continue
+        value_string = "row " + str(ind) + " : "
+        row_cell_values = [str(cell_value) if isinstance(cell_value, int) else cell_value.lower()
+                           for cell_value in row_example]
+        value_string += " | ".join(row_cell_values)
+        value_token_len = len(tokenizer.tokenize(value_string))
+        # over the size limit, and take action
+        if value_token_len > current_chunk_remain_size:
+            truncate_table_contents = {
+                "header": _normalized_table["header"],
+                "rows": [_normalized_table["rows"][i] for i in keep_rows],
+                "types": _normalized_table["types"],
+                "alias": _normalized_table["alias"]
+            }
+            # return
+            break
+        keep_rows.append(ind)
+        current_chunk_remain_size -= value_token_len
+
+    # the id in database starts from 1
+    removed_rows = [i + 1 for i in range(len(_normalized_table["rows"])) if i not in keep_rows]
+    return truncate_table_contents, removed_rows
+
+
+def process_table_structure(_wtq_table_content: Dict, _add_all_column: bool = False):
+    # remove id and agg column
+    headers = [_.replace("\n", " ").lower() for _ in _wtq_table_content["headers"][2:]]
+    header_map = {}
+    for i in range(len(headers)):
+        header_map["c" + str(i + 1)] = headers[i]
+    header_types = _wtq_table_content["types"][2:]
+
+    all_headers = []
+    all_header_types = []
+    vertical_content = []
+    for column_content in _wtq_table_content["contents"][2:]:
+        # only take the first one
+        if _add_all_column:
+            for i in range(len(column_content)):
+                column_alias = column_content[i]["col"]
+                # do not add the numbered column
+                if "_number" in column_alias:
+                    continue
+                vertical_content.append([str(_).replace("\n", " ").lower() for _ in column_content[i]["data"]])
+                if "_" in column_alias:
+                    first_slash_pos = column_alias.find("_")
+                    column_name = header_map[column_alias[:first_slash_pos]] + " " + \
+                                  column_alias[first_slash_pos + 1:].replace("_", " ")
+                else:
+                    column_name = header_map[column_alias]
+                all_headers.append(column_name)
+                if column_content[i]["type"] == "TEXT":
+                    all_header_types.append("text")
+                else:
+                    all_header_types.append("number")
+        else:
+            vertical_content.append([str(_).replace("\n", " ").lower() for _ in column_content[0]["data"]])
+    row_content = list(map(list, zip(*vertical_content)))
+
+    if _add_all_column:
+        ret_header = all_headers
+        ret_types = all_header_types
+    else:
+        ret_header = headers
+        ret_types = header_types
+    return {
+        "header": ret_header,
+        "rows": row_content,
+        "types": ret_types,
+        "alias": list(_wtq_table_content["is_list"].keys())
+    }
+
+
+if __name__ == '__main__':
+
+    random.seed(42)
+    numpy.random.seed(42)
+
+    parser = ArgumentParser()
+    # TODO: if you want to provide your SQL templates, you could organize your file with the format of SQUALL data
+    #  and you should also prepare the corresponding database files / csv files for tables.
+    parser.add_argument('--template_file', help='SQL query file which provides the template for synthesizing more SQL queries',
+                        type=str, default='squall/data.json')
+    # the mode is to avoid sampling tables in the dev set of squall to avoid potential data leakage
+    parser.add_argument('--mode', help='train or dev for pre-training', type=str, default='train',
+                        choices=["train", "dev"])
+    parser.add_argument('--dev_id_file', help='the dev id file to avoid potential data leakage', type=str,
+                        default='squall/dev-0.ids')
+    parser.add_argument('--instance_number', help='the expected instance number corresponding to each template', type=int,
+                        default=500)
+    parser.add_argument('--max_source_length', help='the maximum length for the flattened table plus input SQL query',
+                        type=int,
+                        default=1024)
+
+    args = parser.parse_args()
+
+    mode = args.mode
+    dev_table_ids = json.load(open(args.dev_id_file, "r", encoding="utf8"))
+
+    output_f = open("{}.json".format(mode), "w", encoding="utf8")
+
+    # source table should not be truncated!
+    src_table_content_map = {}
+    # tgt table should be truncated!
+    tgt_table_content_map = {}
+    table_drop_rows_map = {}
+    db_engine_map = {}
+
+    if not os.path.exists(TEMP_DATABASE_PATH):
+        os.makedirs(TEMP_DATABASE_PATH)
+    for table_json_file in os.listdir(TABLE_PATH):
+        table_id = table_json_file[:-5]
+        check_table_file = open(os.path.join(TABLE_PATH, table_json_file), "r", encoding="utf8")
+        src_table_content = json.load(check_table_file)
+        src_table_content = process_table_structure(src_table_content)
+        # source should not be truncated
+        src_table_content_map[table_id] = json.loads(json.dumps(src_table_content))
+        src_table_content, table_drop_rows = truncate_table(src_table_content,
+                                                            max_table_length=args.max_source_length - 40)
+        tgt_table_content_map[table_id] = src_table_content
+        table_drop_rows_map[table_id] = table_drop_rows
+
+    for table_db_file in os.listdir(DATABASE_PATH):
+        table_id = table_db_file[:-3]
+        # copy table db file into a temp file since we may delete some rows
+        database_path = os.path.join(DATABASE_PATH, table_db_file)
+        temp_database_path = os.path.join(TEMP_DATABASE_PATH, table_db_file)
+        if os.path.exists(temp_database_path):
+            os.remove(temp_database_path)
+        # future operations on the temp db to avoid effecting the original database
+        shutil.copy(database_path, temp_database_path)
+        db_engine_map[table_id] = WTQDBEngine(temp_database_path)
+        if table_id in table_drop_rows_map and len(table_drop_rows_map[table_id]) != 0:
+            table_drop_rows = table_drop_rows_map[table_id]
+            db_engine_map[table_id].delete_rows(table_drop_rows)
+
+    valid_table_ids = list(src_table_content_map.keys() - set(dev_table_ids)) if mode == "train" else dev_table_ids
+    max_source_len = args.max_source_length
+    # count the dataset size
+    dataset_counter = 0
+    with open(args.template_file) as fs:
+        examples = json.load(open(args.template_file, "r", encoding="utf8"))
+        for example in tqdm(examples):
+            table_id = example["tbl"]
+            if mode == "dev" and table_id not in dev_table_ids:
+                continue
+            elif mode == "train" and table_id in dev_table_ids:
+                continue
+            sql_struct = example["sql"]
+            question = " ".join(example["nl"]).lower()
+
+            engine, src_table_content = db_engine_map[table_id], src_table_content_map[table_id]
+
+            # augment data
+            if args.instance_number > 0:
+                instance_number_count = 0
+                instance_number_upper = args.instance_number
+                maximum_try_times = instance_number_upper * 10
+                # if we have tried for more than 10 times and cannot get a reasonable execution result, just ignore it
+                while instance_number_count < instance_number_upper and maximum_try_times >= 0:
+                    random_table_id = random.choice(valid_table_ids)
+                    tgt_table_content = tgt_table_content_map[random_table_id]
+                    tgt_db_engine = db_engine_map[random_table_id]
+                    try:
+                        random_sql, random_answer, exec_sql = apply_sql_on_target_table(sql_struct,
+                                                                                        src_table_content,
+                                                                                        tgt_table_content,
+                                                                                        tgt_db_engine,
+                                                                                        _unexec_prob=0.0)
+                        if len(random_answer) > 0:
+                            flatten_input = flatten_schema(tgt_table_content, random_sql)
+                            flatten_output = TGT_DEL.join([str(case).lower() for case in random_answer])
+                            tokens = tokenizer.tokenize(flatten_input)
+                            instance_number_count += 1
+                            output_f.write(json.dumps({
+                                "template_SQL": sql_struct,
+                                "executable_SQL": exec_sql,
+                                "input_SQL": random_sql,
+                                "output_answer": random_answer,
+                                "table": tgt_table_content,
+                                "fairseq_input": flatten_input,
+                                "fairseq_output": flatten_output
+                            }) + "\n")
+                    except Exception as e:
+                        print(e)
+                    # when exceeding the upper limit
+                    maximum_try_times -= 1
+            else:
+                print("Cannot synthesize meaningful SQL queries from SQL: {}, table_id: {}".format(exec_sql, table_id))
+
+    output_f.close()

data_generator/squall/dev-0.ids

@@ -0,0 +1,338 @@
+[
+  "203_447",
+  "204_899",
+  "204_530",
+  "201_10",
+  "203_16",
+  "204_751",
+  "204_199",
+  "204_832",
+  "203_817",
+  "202_26",
+  "201_7",
+  "204_251",
+  "204_604",
+  "204_488",
+  "204_228",
+  "203_513",
+  "204_45",
+  "203_320",
+  "203_189",
+  "204_319",
+  "204_58",
+  "203_224",
+  "203_755",
+  "204_853",
+  "203_611",
+  "204_291",
+  "203_357",
+  "203_470",
+  "204_822",
+  "204_324",
+  "203_839",
+  "203_60",
+  "202_127",
+  "203_848",
+  "204_716",
+  "204_374",
+  "204_718",
+  "203_364",
+  "204_568",
+  "203_737",
+  "204_307",
+  "203_76",
+  "203_527",
+  "204_855",
+  "204_544",
+  "204_492",
+  "204_548",
+  "204_363",
+  "203_721",
+  "204_628",
+  "203_823",
+  "203_177",
+  "204_178",
+  "203_855",
+  "202_231",
+  "203_184",
+  "203_694",
+  "204_713",
+  "203_749",
+  "203_466",
+  "204_771",
+  "204_238",
+  "203_206",
+  "203_516",
+  "203_467",
+  "203_866",
+  "203_519",
+  "202_167",
+  "203_440",
+  "203_728",
+  "204_876",
+  "204_524",
+  "204_240",
+  "203_726",
+  "203_735",
+  "203_475",
+  "203_5",
+  "203_653",
+  "203_569",
+  "202_85",
+  "202_74",
+  "204_994",
+  "204_636",
+  "204_281",
+  "203_350",
+  "204_883",
+  "204_96",
+  "204_800",
+  "202_287",
+  "203_213",
+  "203_308",
+  "204_616",
+  "200_0",
+  "203_11",
+  "204_49",
+  "203_450",
+  "204_637",
+  "203_344",
+  "203_298",
+  "203_363",
+  "203_40",
+  "204_778",
+  "204_28",
+  "204_292",
+  "203_412",
+  "203_34",
+  "204_223",
+  "204_630",
+  "203_778",
+  "204_337",
+  "204_29",
+  "204_846",
+  "203_873",
+  "203_372",
+  "204_836",
+  "201_18",
+  "204_935",
+  "202_278",
+  "202_191",
+  "203_462",
+  "203_724",
+  "203_83",
+  "204_383",
+  "203_715",
+  "204_654",
+  "203_99",
+  "201_42",
+  "203_269",
+  "204_37",
+  "204_22",
+  "202_12",
+  "204_763",
+  "204_988",
+  "203_601",
+  "203_419",
+  "202_150",
+  "204_141",
+  "203_753",
+  "203_618",
+  "204_987",
+  "202_269",
+  "203_572",
+  "204_877",
+  "204_829",
+  "203_61",
+  "202_122",
+  "204_907",
+  "204_725",
+  "204_912",
+  "203_115",
+  "202_195",
+  "203_234",
+  "203_220",
+  "203_106",
+  "203_551",
+  "204_200",
+  "203_346",
+  "204_796",
+  "204_51",
+  "201_43",
+  "204_805",
+  "204_997",
+  "203_478",
+  "204_766",
+  "204_484",
+  "204_662",
+  "204_649",
+  "203_683",
+  "204_364",
+  "203_671",
+  "204_27",
+  "203_171",
+  "203_563",
+  "204_940",
+  "204_983",
+  "204_833",
+  "203_522",
+  "204_140",
+  "203_15",
+  "203_723",
+  "204_104",
+  "204_581",
+  "202_251",
+  "204_12",
+  "204_547",
+  "204_958",
+  "204_739",
+  "203_746",
+  "204_275",
+  "203_312",
+  "203_282",
+  "204_146",
+  "204_18",
+  "204_918",
+  "203_656",
+  "204_105",
+  "203_594",
+  "203_438",
+  "204_334",
+  "203_251",
+  "203_411",
+  "204_794",
+  "204_787",
+  "203_633",
+  "203_590",
+  "203_13",
+  "204_453",
+  "202_202",
+  "203_17",
+  "203_621",
+  "203_700",
+  "203_293",
+  "203_400",
+  "204_250",
+  "203_808",
+  "203_496",
+  "203_192",
+  "204_567",
+  "204_0",
+  "203_530",
+  "203_874",
+  "203_802",
+  "204_986",
+  "204_540",
+  "203_174",
+  "203_654",
+  "203_581",
+  "203_788",
+  "203_603",
+  "204_431",
+  "204_969",
+  "203_703",
+  "204_558",
+  "204_780",
+  "204_385",
+  "203_295",
+  "202_50",
+  "204_210",
+  "203_358",
+  "204_184",
+  "204_593",
+  "202_239",
+  "203_130",
+  "204_981",
+  "204_358",
+  "203_8",
+  "204_409",
+  "203_523",
+  "203_797",
+  "204_1",
+  "204_608",
+  "203_137",
+  "204_153",
+  "204_917",
+  "203_676",
+  "204_842",
+  "204_413",
+  "204_737",
+  "203_871",
+  "203_757",
+  "203_782",
+  "203_727",
+  "204_872",
+  "203_332",
+  "204_529",
+  "204_749",
+  "204_706",
+  "204_465",
+  "203_851",
+  "204_205",
+  "203_634",
+  "203_787",
+  "203_529",
+  "204_325",
+  "203_392",
+  "204_927",
+  "204_202",
+  "203_276",
+  "203_387",
+  "203_507",
+  "203_862",
+  "204_4",
+  "204_428",
+  "204_617",
+  "203_180",
+  "204_172",
+  "202_233",
+  "202_178",
+  "202_192",
+  "202_224",
+  "203_141",
+  "204_503",
+  "204_209",
+  "202_135",
+  "204_333",
+  "203_600",
+  "203_859",
+  "203_242",
+  "204_676",
+  "204_59",
+  "203_384",
+  "204_941",
+  "204_63",
+  "204_561",
+  "204_121",
+  "204_388",
+  "204_435",
+  "203_417",
+  "200_10",
+  "203_456",
+  "204_136",
+  "204_808",
+  "204_255",
+  "204_971",
+  "204_6",
+  "204_317",
+  "203_756",
+  "204_967",
+  "204_553",
+  "203_844",
+  "204_772",
+  "202_157",
+  "203_414",
+  "204_212",
+  "204_990",
+  "204_573",
+  "204_811",
+  "204_867",
+  "204_850",
+  "203_94",
+  "204_959",
+  "203_791",
+  "203_616",
+  "203_92",
+  "204_651",
+  "204_217"
+]