import numpy as np import random import torch from tqdm import tqdm from transformers import AutoTokenizer from utils.classifier_model import SchemaItemClassifier from transformers.trainer_utils import set_seed def prepare_inputs_and_labels(sample, tokenizer): table_names = [table["table_name"] for table in sample["schema"]["schema_items"]] column_names = [table["column_names"] for table in sample["schema"]["schema_items"]] column_num_in_each_table = [len(table["column_names"]) for table in sample["schema"]["schema_items"]] # `column_name_word_indices` and `table_name_word_indices` record the word indices of each column and table in `input_words`, whose element is an integer column_name_word_indices, table_name_word_indices = [], [] input_words = [sample["text"]] for table_id, table_name in enumerate(table_names): input_words.append("|") input_words.append(table_name) table_name_word_indices.append(len(input_words) - 1) input_words.append(":") for column_name in column_names[table_id]: input_words.append(column_name) column_name_word_indices.append(len(input_words) - 1) input_words.append(",") # remove the last "," input_words = input_words[:-1] tokenized_inputs = tokenizer( input_words, return_tensors="pt", is_split_into_words = True, padding = "max_length", max_length = 512, truncation = True ) # after tokenizing, one table name or column name may be splitted into multiple tokens (i.e., sub-words) # `column_name_token_indices` and `table_name_token_indices` records the token indices of each column and table in `input_ids`, whose element is a list of integer column_name_token_indices, table_name_token_indices = [], [] word_indices = tokenized_inputs.word_ids(batch_index = 0) # obtain token indices of each column in `input_ids` for column_name_word_index in column_name_word_indices: column_name_token_indices.append([token_id for token_id, word_index in enumerate(word_indices) if column_name_word_index == word_index]) # obtain token indices of each table in `input_ids` for table_name_word_index in table_name_word_indices: table_name_token_indices.append([token_id for token_id, word_index in enumerate(word_indices) if table_name_word_index == word_index]) encoder_input_ids = tokenized_inputs["input_ids"] encoder_input_attention_mask = tokenized_inputs["attention_mask"] # print("\n".join(tokenizer.batch_decode(encoder_input_ids, skip_special_tokens = True))) if torch.cuda.is_available(): encoder_input_ids = encoder_input_ids.cuda() encoder_input_attention_mask = encoder_input_attention_mask.cuda() return encoder_input_ids, encoder_input_attention_mask, \ column_name_token_indices, table_name_token_indices, column_num_in_each_table def get_schema(tables_and_columns): schema_items = [] table_names = list(dict.fromkeys([t for t, c in tables_and_columns])) for table_name in table_names: schema_items.append( { "table_name": table_name, "column_names": [c for t, c in tables_and_columns if t == table_name] } ) return {"schema_items": schema_items} def get_sequence_length(text, tables_and_columns, tokenizer): table_names = [t for t, c in tables_and_columns] # duplicate `table_names` while preserving order table_names = list(dict.fromkeys(table_names)) column_names = [] for table_name in table_names: column_names.append([c for t, c in tables_and_columns if t == table_name]) input_words = [text] for table_id, table_name in enumerate(table_names): input_words.append("|") input_words.append(table_name) input_words.append(":") for column_name in column_names[table_id]: input_words.append(column_name) input_words.append(",") # remove the last "," input_words = input_words[:-1] tokenized_inputs = tokenizer(input_words, is_split_into_words = True) return len(tokenized_inputs["input_ids"]) # handle extremely long schema sequences def split_sample(sample, tokenizer): text = sample["text"] table_names = [] column_names = [] for table in sample["schema"]["schema_items"]: table_names.append(table["table_name"] + " ( " + table["table_comment"] + " ) " \ if table["table_comment"] != "" else table["table_name"]) column_names.append([column_name + " ( " + column_comment + " ) " \ if column_comment != "" else column_name \ for column_name, column_comment in zip(table["column_names"], table["column_comments"])]) splitted_samples = [] recorded_tables_and_columns = [] for table_idx, table_name in enumerate(table_names): for column_name in column_names[table_idx]: if get_sequence_length(text, recorded_tables_and_columns + [[table_name, column_name]], tokenizer) < 500: recorded_tables_and_columns.append([table_name, column_name]) else: splitted_samples.append( { "text": text, "schema": get_schema(recorded_tables_and_columns) } ) recorded_tables_and_columns = [[table_name, column_name]] splitted_samples.append( { "text": text, "schema": get_schema(recorded_tables_and_columns) } ) return splitted_samples def merge_pred_results(sample, pred_results): # table_names = [table["table_name"] for table in sample["schema"]["schema_items"]] # column_names = [table["column_names"] for table in sample["schema"]["schema_items"]] table_names = [] column_names = [] for table in sample["schema"]["schema_items"]: table_names.append(table["table_name"] + " ( " + table["table_comment"] + " ) " \ if table["table_comment"] != "" else table["table_name"]) column_names.append([column_name + " ( " + column_comment + " ) " \ if column_comment != "" else column_name \ for column_name, column_comment in zip(table["column_names"], table["column_comments"])]) merged_results = [] for table_id, table_name in enumerate(table_names): table_prob = 0 column_probs = [] for result_dict in pred_results: if table_name in result_dict: if table_prob < result_dict[table_name]["table_prob"]: table_prob = result_dict[table_name]["table_prob"] column_probs += result_dict[table_name]["column_probs"] merged_results.append( { "table_name": table_name, "table_prob": table_prob, "column_names": column_names[table_id], "column_probs": column_probs } ) return merged_results def filter_func(dataset, dataset_type, sic, num_top_k_tables = 5, num_top_k_columns = 5): for data in tqdm(dataset, desc = "filtering schema items for the dataset"): filtered_schema = dict() filtered_schema["schema_items"] = [] table_names = [table["table_name"] for table in data["schema"]["schema_items"]] table_comments = [table["table_comment"] for table in data["schema"]["schema_items"]] column_names = [table["column_names"] for table in data["schema"]["schema_items"]] column_comments = [table["column_comments"] for table in data["schema"]["schema_items"]] if dataset_type == "eval": # predict scores for each tables and columns pred_results = sic.predict(data) # remain top_k1 tables for each database and top_k2 columns for each remained table table_probs = [pred_result["table_prob"] for pred_result in pred_results] table_indices = np.argsort(-np.array(table_probs), kind="stable")[:num_top_k_tables].tolist() elif dataset_type == "train": table_indices = [table_idx for table_idx, table_label in enumerate(data["table_labels"]) if table_label == 1] if len(table_indices) < num_top_k_tables: unused_table_indices = [table_idx for table_idx, table_label in enumerate(data["table_labels"]) if table_label == 0] table_indices += random.sample(unused_table_indices, min(len(unused_table_indices), num_top_k_tables - len(table_indices))) random.shuffle(table_indices) for table_idx in table_indices: if dataset_type == "eval": column_probs = pred_results[table_idx]["column_probs"] column_indices = np.argsort(-np.array(column_probs), kind="stable")[:num_top_k_columns].tolist() elif dataset_type == "train": column_indices = [column_idx for column_idx, column_label in enumerate(data["column_labels"][table_idx]) if column_label == 1] if len(column_indices) < num_top_k_columns: unused_column_indices = [column_idx for column_idx, column_label in enumerate(data["column_labels"][table_idx]) if column_label == 0] column_indices += random.sample(unused_column_indices, min(len(unused_column_indices), num_top_k_columns - len(column_indices))) random.shuffle(column_indices) filtered_schema["schema_items"].append( { "table_name": table_names[table_idx], "table_comment": table_comments[table_idx], "column_names": [column_names[table_idx][column_idx] for column_idx in column_indices], "column_comments": [column_comments[table_idx][column_idx] for column_idx in column_indices] } ) # replace the old schema with the filtered schema data["schema"] = filtered_schema if dataset_type == "train": del data["table_labels"] del data["column_labels"] return dataset def lista_contains_listb(lista, listb): for b in listb: if b not in lista: return 0 return 1 class SchemaItemClassifierInference(): def __init__(self, model_save_path): set_seed(42) # load tokenizer self.tokenizer = AutoTokenizer.from_pretrained(model_save_path, add_prefix_space = True) # initialize model self.model = SchemaItemClassifier(model_save_path, "test") # load fine-tuned params self.model.load_state_dict(torch.load(model_save_path + "/dense_classifier.pt", map_location=torch.device('cpu')), strict=False) if torch.cuda.is_available(): self.model = self.model.cuda() self.model.eval() def predict_one(self, sample): encoder_input_ids, encoder_input_attention_mask, column_name_token_indices,\ table_name_token_indices, column_num_in_each_table = prepare_inputs_and_labels(sample, self.tokenizer) with torch.no_grad(): model_outputs = self.model( encoder_input_ids, encoder_input_attention_mask, [column_name_token_indices], [table_name_token_indices], [column_num_in_each_table] ) table_logits = model_outputs["batch_table_name_cls_logits"][0] table_pred_probs = torch.nn.functional.softmax(table_logits, dim = 1)[:, 1].cpu().tolist() column_logits = model_outputs["batch_column_info_cls_logits"][0] column_pred_probs = torch.nn.functional.softmax(column_logits, dim = 1)[:, 1].cpu().tolist() splitted_column_pred_probs = [] # split predicted column probs into each table for table_id, column_num in enumerate(column_num_in_each_table): splitted_column_pred_probs.append(column_pred_probs[sum(column_num_in_each_table[:table_id]): sum(column_num_in_each_table[:table_id]) + column_num]) column_pred_probs = splitted_column_pred_probs result_dict = dict() for table_idx, table in enumerate(sample["schema"]["schema_items"]): result_dict[table["table_name"]] = { "table_name": table["table_name"], "table_prob": table_pred_probs[table_idx], "column_names": table["column_names"], "column_probs": column_pred_probs[table_idx], } return result_dict def predict(self, test_sample): splitted_samples = split_sample(test_sample, self.tokenizer) pred_results = [] for splitted_sample in splitted_samples: pred_results.append(self.predict_one(splitted_sample)) return merge_pred_results(test_sample, pred_results) def evaluate_coverage(self, dataset): max_k = 100 total_num_for_table_coverage, total_num_for_column_coverage = 0, 0 table_coverage_results = [0]*max_k column_coverage_results = [0]*max_k for data in dataset: indices_of_used_tables = [idx for idx, label in enumerate(data["table_labels"]) if label == 1] pred_results = sic.predict(data) # print(pred_results) table_probs = [res["table_prob"] for res in pred_results] for k in range(max_k): indices_of_top_k_tables = np.argsort(-np.array(table_probs), kind="stable")[:k+1].tolist() if lista_contains_listb(indices_of_top_k_tables, indices_of_used_tables): table_coverage_results[k] += 1 total_num_for_table_coverage += 1 for table_idx in range(len(data["table_labels"])): indices_of_used_columns = [idx for idx, label in enumerate(data["column_labels"][table_idx]) if label == 1] if len(indices_of_used_columns) == 0: continue column_probs = pred_results[table_idx]["column_probs"] for k in range(max_k): indices_of_top_k_columns = np.argsort(-np.array(column_probs), kind="stable")[:k+1].tolist() if lista_contains_listb(indices_of_top_k_columns, indices_of_used_columns): column_coverage_results[k] += 1 total_num_for_column_coverage += 1 indices_of_top_10_columns = np.argsort(-np.array(column_probs), kind="stable")[:10].tolist() if lista_contains_listb(indices_of_top_10_columns, indices_of_used_columns) == 0: print(pred_results[table_idx]) print(data["column_labels"][table_idx]) print(data["question"]) print(total_num_for_table_coverage) print(table_coverage_results) print(total_num_for_column_coverage) print(column_coverage_results) if __name__ == "__main__": dataset_name = "bird_with_evidence" # dataset_name = "bird" # dataset_name = "spider" sic = SchemaItemClassifierInference("sic_ckpts/sic_{}".format(dataset_name)) import json dataset = json.load(open("./data/sft_eval_{}_text2sql.json".format(dataset_name))) sic.evaluate_coverage(dataset)