|
|
import logging |
|
|
import os |
|
|
import pickle |
|
|
import time |
|
|
import json |
|
|
from dataclasses import dataclass |
|
|
from typing import Any, Dict, List, NewType, Tuple |
|
|
from tqdm import tqdm |
|
|
|
|
|
|
|
|
import torch |
|
|
from torch.utils.data.dataset import Dataset |
|
|
from transformers.tokenization_utils import PreTrainedTokenizer |
|
|
from transformers.data.data_collator import DataCollator |
|
|
from transformers.tokenization_bart import BartTokenizer |
|
|
from transformers.tokenization_roberta import RobertaTokenizer |
|
|
from relogic.pretrainkit.datasets.utils import pad_and_tensorize_sequence |
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
label_mapping = json.load(open("data/preprocessed_data/bart_parser_label_mapping_2.json")) |
|
|
|
|
|
class QuerySchemaRelation2SQLDataset(Dataset): |
|
|
""" |
|
|
Dataset for relation-aware text-to-SQL: query + schema + relation -> SQL |
|
|
""" |
|
|
def __init__(self, tokenizer: PreTrainedTokenizer, file_path: str, block_size: int, local_rank=-1): |
|
|
self.examples = [] |
|
|
self.keywords = label_mapping["keyword"] |
|
|
self.label_eos_id = self.keywords.index(label_mapping["label_eos_token"]) |
|
|
self.label_bos_id = self.keywords.index(label_mapping["label_bos_token"]) |
|
|
add_prefix_space = isinstance(tokenizer, BartTokenizer) or isinstance(tokenizer, RobertaTokenizer) |
|
|
total, valid = 0, 0 |
|
|
with open(file_path, encoding="utf-8") as f: |
|
|
for line in tqdm(f): |
|
|
total += 1 |
|
|
example = json.loads(line) |
|
|
text = example["normalized_question"] |
|
|
columns = example["columns"] |
|
|
tables = example["tables"] |
|
|
columns_text = example["column_text"] |
|
|
tables_text = example["table_text"] |
|
|
sql = example["sql"] |
|
|
|
|
|
token_idx_to_sub_token_start_idx = {} |
|
|
text_tokens = [tokenizer.cls_token] |
|
|
start_idx = 0 |
|
|
for idx, token in enumerate(text.split()): |
|
|
sub_tokens = tokenizer.tokenize(token, add_prefix_space=add_prefix_space) |
|
|
token_idx_to_sub_token_start_idx[idx] = start_idx |
|
|
text_tokens.extend(sub_tokens) |
|
|
start_idx += len(sub_tokens) |
|
|
text_tokens.append(tokenizer.sep_token) |
|
|
question_start, question_end = 1, len(text_tokens) - 1 |
|
|
|
|
|
column_spans = [] |
|
|
start_idx = len(text_tokens) |
|
|
for column_tokens in columns_text: |
|
|
column_str = " ".join(column_tokens) |
|
|
column_tokens = tokenizer.tokenize(column_str, add_prefix_space=add_prefix_space) |
|
|
text_tokens.extend(column_tokens) |
|
|
text_tokens.append(tokenizer.sep_token) |
|
|
end_idx = start_idx + len(column_tokens) |
|
|
column_spans.append((start_idx, end_idx)) |
|
|
start_idx = end_idx + 1 |
|
|
|
|
|
column_start = [column_span[0] for column_span in column_spans] |
|
|
column_end = [column_span[1] for column_span in column_spans] |
|
|
|
|
|
table_spans = [] |
|
|
start_idx = len(text_tokens) |
|
|
for table_tokens in tables_text: |
|
|
table_str = " ".join(table_tokens) |
|
|
table_tokens = tokenizer.tokenize(table_str, add_prefix_space=add_prefix_space) |
|
|
text_tokens.extend(table_tokens) |
|
|
text_tokens.append(tokenizer.sep_token) |
|
|
end_idx = start_idx + len(table_tokens) |
|
|
table_spans.append((start_idx, end_idx)) |
|
|
start_idx = end_idx + 1 |
|
|
|
|
|
table_start = [table_span[0] for table_span in table_spans] |
|
|
table_end = [table_span[1] for table_span in table_spans] |
|
|
|
|
|
input_ids = tokenizer.convert_tokens_to_ids(text_tokens) |
|
|
|
|
|
if len(input_ids) > block_size: |
|
|
continue |
|
|
|
|
|
label_ids = [] |
|
|
try: |
|
|
for token in sql.split(): |
|
|
if token in columns: |
|
|
label_ids.append(columns.index(token) + len(self.keywords)) |
|
|
else: |
|
|
label_ids.append(self.keywords.index(token)) |
|
|
except: |
|
|
continue |
|
|
|
|
|
label_ids = [self.label_bos_id] + label_ids + [self.label_eos_id] |
|
|
|
|
|
primary_key = [int(x) for x in example["sc_struct"]["primary_key"]] |
|
|
foreign_key = {x.split(",")[0]: int(x.split(",")[1]) for x in example["sc_struct"]["foreign_key"]} |
|
|
column_to_table = {"0": None} |
|
|
|
|
|
sc_link = {"q_col_match": {}, "q_tab_match": {}} |
|
|
for k, v in example["sc_link"]["q_col_match"].items(): |
|
|
new_k = str(token_idx_to_sub_token_start_idx[int(k.split(",")[0])]) + "," + k.split(",")[1] |
|
|
sc_link["q_col_match"][new_k] = v |
|
|
|
|
|
for k, v in example["sc_link"]["q_tab_match"].items(): |
|
|
new_k = str(token_idx_to_sub_token_start_idx[int(k.split(",")[0])]) + "," + k.split(",")[1] |
|
|
sc_link["q_tab_match"][new_k] = v |
|
|
|
|
|
cv_link = {"num_date_match": {}, "cell_match": {}} |
|
|
for k, v in example["cv_link"]["num_date_match"].items(): |
|
|
new_k = str(token_idx_to_sub_token_start_idx[int(k.split(",")[0])]) + "," + k.split(",")[1] |
|
|
cv_link["num_date_match"][new_k] = v |
|
|
for k, v in example["cv_link"]["cell_match"].items(): |
|
|
new_k = str(token_idx_to_sub_token_start_idx[int(k.split(",")[0])]) + "," + k.split(",")[1] |
|
|
cv_link["cell_match"][new_k] = v |
|
|
|
|
|
|
|
|
for idx, column in enumerate(columns): |
|
|
if column == "*": |
|
|
continue |
|
|
t = column.split(".")[0] |
|
|
column_to_table[str(idx)] = tables.index(t) |
|
|
|
|
|
foreign_keys_tables = {} |
|
|
for k, v in foreign_key.items(): |
|
|
t_k = str(column_to_table[str(k)]) |
|
|
t_v = str(column_to_table[str(v)]) |
|
|
if t_k not in foreign_keys_tables: |
|
|
foreign_keys_tables[t_k] = [] |
|
|
if int(t_v) not in foreign_keys_tables[t_k]: |
|
|
foreign_keys_tables[t_k].append(int(t_v)) |
|
|
|
|
|
self.examples.append({ |
|
|
"input_ids": input_ids, |
|
|
"example_info": { |
|
|
"normalized_question": text, |
|
|
"columns": columns, |
|
|
"tables": tables, |
|
|
"tokens": text_tokens, |
|
|
"question_start": question_start, |
|
|
"question_end": question_end, |
|
|
"column_start": torch.LongTensor(column_start), |
|
|
"column_end": torch.LongTensor(column_end), |
|
|
"table_start": torch.LongTensor(table_start), |
|
|
"table_end": torch.LongTensor(table_end), |
|
|
"sc_link": sc_link, |
|
|
"cv_link": cv_link, |
|
|
"primary_keys": primary_key, |
|
|
"foreign_keys": foreign_key, |
|
|
"column_to_table": column_to_table, |
|
|
"foreign_keys_tables": foreign_keys_tables |
|
|
}, |
|
|
"column_spans": column_spans, |
|
|
"label_ids": label_ids}) |
|
|
valid += 1 |
|
|
print("Valid Example {}; Invalid Example {}".format(valid, total - valid)) |
|
|
|
|
|
def __len__(self): |
|
|
return len(self.examples) |
|
|
|
|
|
def __getitem__(self, i): |
|
|
return self.examples[i] |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class DataCollatorForQuerySchemaRelation2SQL: |
|
|
""" |
|
|
Data collator used for query + schema -> sql modeling. |
|
|
""" |
|
|
tokenizer: PreTrainedTokenizer |
|
|
label_padding_id = label_mapping["keyword"].index(label_mapping["label_padding_token"]) |
|
|
label_eos_id = label_mapping["keyword"].index(label_mapping["label_eos_token"]) |
|
|
label_bos_id = label_mapping["keyword"].index(label_mapping["label_bos_token"]) |
|
|
def collate_batch(self, examples) -> Dict[str, torch.Tensor]: |
|
|
|
|
|
input_ids_sequences = [example["input_ids"] for example in examples] |
|
|
column_spans_sequences = [example["column_spans"] for example in examples] |
|
|
label_ids_sequences = [example["label_ids"] for example in examples] |
|
|
padded_input_ids_tensor = pad_and_tensorize_sequence( |
|
|
input_ids_sequences, padding_value=self.tokenizer.pad_token_id) |
|
|
padded_column_spans_tensor = pad_and_tensorize_sequence( |
|
|
column_spans_sequences, padding_value=(0, 1)) |
|
|
|
|
|
example_info_list = [] |
|
|
for example in examples: |
|
|
example_info_list.append(example["example_info"]) |
|
|
label_ids_tensor = pad_and_tensorize_sequence( |
|
|
label_ids_sequences, padding_value=self.label_padding_id) |
|
|
return { |
|
|
"input_ids": padded_input_ids_tensor, |
|
|
"column_spans": padded_column_spans_tensor, |
|
|
"labels": label_ids_tensor, |
|
|
"example_info_list": example_info_list, |
|
|
"input_padding_id": self.tokenizer.pad_token_id, |
|
|
"label_padding_id": self.label_padding_id, |
|
|
"label_eos_id": self.label_eos_id, |
|
|
"label_bos_id": self.label_bos_id |
|
|
} |
|
|
|
|
|
|
