| | 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 |
| | } |
| |
|
| |
|
