# -*- coding: utf-8 -*- import argparse import copy import csv import json import re import sqlite3 import traceback import os from vllm import LLM, SamplingParams from func_timeout import func_set_timeout import func_timeout import tqdm prompt_temp = """Given the following database schema and question, your task is to write a valid SQL query whose execution results can accurately answer the question. /* Database schema */ {ds} /* Sample rows of each table */ {sr} /* Question */ {qs}{hint} Answer the question by a SQL query only with no explanation: """ prompt_sl_temp_sft = """Given the following database schema and question, your task is to extract the tables and columns relevant to solving the question. /* Database schema */ {ds} /* Sample rows of each table */ {sr} /* Question */ {qs}{hint} Output the tables and columns only with no explanation: """ prompt_sl_temp = """Given the following database schema and question, your task is to extract the tables and columns relevant to solving the question. /* Examples */ Example 1: Database schema: CREATE TABLE department (Department_ID NUMBER, Name TEXT, Creation TEXT, Ranking NUMBER, Budget_in_Billions NUMBER, Num_Employees NUMBER, PRIMARY KEY(Department_ID, Name, Creation, Ranking, Budget_in_Billions, Num_Employees)); CREATE TABLE head (head_ID NUMBER, name TEXT, born_state TEXT, age NUMBER, PRIMARY KEY(head_ID, name, born_state, age)); CREATE TABLE management (department_ID NUMBER, head_ID NUMBER, temporary_acting TEXT, PRIMARY KEY(department_ID, head_ID, temporary_acting), FOREIGN KEY (head_ID) REFERENCES head(head_ID), FOREIGN KEY (department_ID) REFERENCES department(Department_ID)); Question: What are the names of the heads who are born outside the California state? Output: {{"head": ["name", "born_state"]}} Example 2: Databse schema: CREATE TABLE department (Department_ID NUMBER, Name TEXT, Creation TEXT, Ranking NUMBER, Budget_in_Billions NUMBER, Num_Employees NUMBER, PRIMARY KEY(Department_ID, Name, Creation, Ranking, Budget_in_Billions, Num_Employees)); CREATE TABLE head (head_ID NUMBER, name TEXT, born_state TEXT, age NUMBER, PRIMARY KEY(head_ID, name, born_state, age)); CREATE TABLE management (department_ID NUMBER, head_ID NUMBER, temporary_acting TEXT, PRIMARY KEY(department_ID, head_ID, temporary_acting), FOREIGN KEY (head_ID) REFERENCES head(head_ID), FOREIGN KEY (department_ID) REFERENCES department(Department_ID)); Question: How many departments are led by heads who are not mentioned? Output: {{"department": ["Department_ID"], "management": ["department_ID"]}} Example 3: Database schema: CREATE TABLE department (Department_ID NUMBER, Name TEXT, Creation TEXT, Ranking NUMBER, Budget_in_Billions NUMBER, Num_Employees NUMBER, PRIMARY KEY(Department_ID, Name, Creation, Ranking, Budget_in_Billions, Num_Employees)); CREATE TABLE head (head_ID NUMBER, name TEXT, born_state TEXT, age NUMBER, PRIMARY KEY(head_ID, name, born_state, age)); CREATE TABLE management (department_ID NUMBER, head_ID NUMBER, temporary_acting TEXT, PRIMARY KEY(department_ID, head_ID, temporary_acting), FOREIGN KEY (head_ID) REFERENCES head(head_ID), FOREIGN KEY (department_ID) REFERENCES department(Department_ID)); Question: How many heads of the departments are older than 56? Output: {{'head': ['age']}} Now, let’s get started! /* Database schema */ {ds} /* Sample rows of each table */ {sr} /* Question */ {qs}{hint} Output the tables and columns only with no explanation: """ prompt_nc_temp_sft = """Your task is to determine whether the execution results of a SQL query can answer the given question according to the following database schema. If the execution results cannot correctly answer the question, please give me the correct SQL query. /* Database schema */ {ds} /* Sample rows of each table */ {sr} /* Question */ {qs}{hint} /* SQL query */ {sql}{ex_hint} Output: """ prompt_nc_temp = """Your task is to determine whether the execution results of a SQL query can answer the given question according to the following database schema. If the execution results cannot correctly answer the question, please give me the correct SQL query. /* Examples */ Example 1: Question: Average of the last receipt cost of the products whose average lead time is 60 days. SQL query: SELECT SUM(LastReceiptCost) / COUNT(ProductID) FROM ProductVendor; Output: The execution results of the SQL query cannot correctly answer the question. The correct SQL query should be: ```sql SELECT SUM(LastReceiptCost) / COUNT(ProductID) FROM ProductVendor WHERE AverageLeadTime = 60; ``` Example 2: Question: Calculate the average price of products shipped to the UK. SQL query: SELECT AVG(UnitPrice) AS avg FROM Invoices WHERE Country = 'UK'; Output: The execution results of the SQL query can correctly answer the question. Example 3: Question: What is the total cost for all the orders placed on 5/29/2013? SQL query: SELECT SUM(TotalDue) FROM PurchaseOrderHeader WHERE OrderDate LIKE '2013-05-29%'; Output: The SQL query can correctly answer the question. Now, let’s get started! /* Database schema */ {ds} /* Sample rows of each table */ {sr} /* Question */ {qs}{hint} /* SQL query */ {sql}{ex_hint} Output: """ prompt_cw_temp_sft = """Given the following database schema and question, your task is to write an incomplete SQL query into a complete SQL query whose execution results can correctly answer the question. /* Database schema */ {ds} /* Sample rows of each table */ {sr} /* Question */ {qs}{hint} /* The incomplete SQL query */ ```sql {sql} ``` Output: """ prompt_cw_temp = """Given the following database schema and question, your task is to write an incomplete SQL query into a complete SQL query whose execution results can correctly answer the question. /* Examples */ Example 1: Question: How many heads of the departments are older than 56 ? The incomplete SQL query: ```sql\nELECT count(*);\n``` Output: ```sql\nELECT count(*) FROM head WHERE age > 56;\n``` Example 2: Question: What are the distinct creation years of the departments managed by a secretary born in state 'Alabama'? The incomplete SQL query: ```sql\nSELECT DISTINCT T1.creation FROM department AS T1;\n``` Output: ```sql\nSELECT DISTINCT T1.creation FROM department AS T1 JOIN management AS T2 ON T1.department_id = T2.department_id JOIN head AS T3 ON T2.head_id = T3.head_id WHERE T3.born_state = 'Alabama';\n``` Example 3: Question: Show the name and number of employees for the departments managed by heads whose temporary acting value is 'Yes'? The incomplete SQL query: ```sql\nSELECT T1.name, T1.num_employees FROM department AS T1 JOIN management AS T2;\n``` Output: ```sql\nSELECT T1.name, T1.num_employees FROM department AS T1 JOIN management AS T2 ON T1.department_id = T2.department_id WHERE T2.temporary_acting = 'Yes';\n``` Now, let’s get started! /* Database schema */ {ds} /* Sample rows of each table */ {sr} /* Question */ {qs}{hint} /* The incomplete SQL query */ ```sql {sql} ``` Output: """ def read_json_file(file_path): try: with open(file_path, 'r', encoding='utf-8-sig') as file: data = json.load(file) return data except Exception as e: print("="*10,e) return None class LLM_Model(object): def __init__(self, model= ''): self.model = model model = model.lower().replace('_','').replace('-','') if 'qwen2' in model: self.tag ='qwen2' elif 'llama3' in model: self.tag ='llama3' elif 'llama2' in model: self.tag ='llam2' elif 'deepseek' in model: self.tag ='deepseek' elif 'mistral' in model: self.tag ='mistral' elif 'codellama' in model: self.tag = 'codellama' else: raise TypeError(f"Unexpect model: {model}.") self.llm = LLM(model=self.model, seed=123, gpu_memory_utilization=0.9, tensor_parallel_size=args.gpus, trust_remote_code=True, ) self.tokenizer = self.llm.get_tokenizer() def generate_response(self, prompts, max_tokens=1024, temperature=0.01, top_p=0.5): sampling_params = SamplingParams(temperature=temperature, top_p=top_p, max_tokens=max_tokens, skip_special_tokens=True, stop=self.tokenizer.eos_token) if self.tag in ['mistral']: messages_list = [[{"role": "user", "content": p}] for p in prompts] else: messages_list = [[{"role": "system", "content": "You are a helpful SQLite assistant."},{"role": "user", "content": p}] for p in prompts] messages_list = self.tokenizer.apply_chat_template(messages_list, add_generation_prompt=True,tokenize=False) outputs = self.llm.generate(messages_list, sampling_params) return [output.outputs[0].text for output in outputs] class LLM_Online(object): def __init__(self, model= "qwen72b", device = [0]): None def generate_response(self, prompts): rs = [] for prompt in tqdm.tqdm(prompts): res = None # your online LLM rs.append(res) return rs def parse_dataset(data_path, mode = 'dev', dataset = 'bird'): # redirect path data_tuples_path = '' if dataset == 'bird': data_tuples_path = os.path.join(data_path, dataset, mode, f'{mode}.json') elif 'spider_DK' == dataset: data_tuples_path = os.path.join(data_path, 'spider', 'Spider_DK.json') elif 'spider_real' == dataset: data_tuples_path = os.path.join(data_path, 'spider', 'spider-realistic.json') elif 'spider' in dataset: if mode == 'test': data_tuples_path = os.path.join(data_path, 'spider','test_data/dev.json') else: data_tuples_path = os.path.join(data_path, 'spider', f'{mode}.json') else: raise TypeError(f"Unexpect dataset: {dataset}.") data_tuples = read_json_file(data_tuples_path) return data_tuples def convert_fk_index(data): fk_holder = [] table_names_original = [i.lower() for i in data['table_names_original']] # some bug column_names_original = [(i[0], i[1].lower()) for i in data['column_names_original']] for fk in data["foreign_keys"]: tn, col, ref_tn, ref_col = fk[0][0], fk[0][1], fk[1][0], fk[1][1] if type(tn) is str: tn = tn.lower() if type(col) is str: col = col.lower() if type(ref_tn) is str: ref_tn = ref_tn.lower() if type(ref_col) is str: ref_col = ref_col.lower() ref_cid, cid = None, None try: tid =table_names_original.index(tn) ref_tid = table_names_original.index(ref_tn) for i, (tab_id, col_org) in enumerate(column_names_original): if tab_id == ref_tid and ref_col == col_org: ref_cid = i elif tid == tab_id and col == col_org: cid = i if ref_cid and cid: fk_holder.append([cid, ref_cid]) except: traceback.print_exc() print("table_names_original: ", table_names_original) print("finding tab name: ", tn, ref_tn) print(data) # sys.exit() return fk_holder def dump_db_json_schema(db, f): '''read table and column info''' try: conn = sqlite3.connect(db) except: print(db) exit() conn.execute('pragma foreign_keys=ON') cursor = conn.execute("SELECT name FROM sqlite_master WHERE type='table';") data = {'db_id': f, 'table_names_original': [], 'table_names': [], 'column_names_original': [(-1, '*')], 'column_names': [(-1, '*')], 'column_types': ['text'], 'primary_keys': [], 'foreign_keys': []} fk_holder = [] for i, item in enumerate(cursor.fetchall()): table_name = item[0] data['table_names_original'].append(table_name) data['table_names'].append(table_name.lower().replace("_", ' ')) fks = conn.execute("PRAGMA foreign_key_list('{}') ".format(table_name)).fetchall() #print("db:{} table:{} fks:{}".format(f,table_name,fks)) fk_holder.extend([[(table_name, fk[3]), (fk[2], fk[4])] for fk in fks]) cur = conn.execute("PRAGMA table_info('{}') ".format(table_name)) for j, col in enumerate(cur.fetchall()): data['column_names_original'].append((i, col[1])) data['column_names'].append((i, col[1].lower().replace("_", " "))) #varchar, '' -> text, int, numeric -> integer, col_type = col[2].lower() if 'char' in col_type or col_type == '' or 'text' in col_type or 'var' in col_type: data['column_types'].append('text') elif 'int' in col_type or 'numeric' in col_type or 'decimal' in col_type or 'number' in col_type\ or 'id' in col_type or 'real' in col_type or 'double' in col_type or 'float' in col_type: data['column_types'].append('number') elif 'date' in col_type or 'time' in col_type or 'year' in col_type: data['column_types'].append('time') elif 'boolean' in col_type: data['column_types'].append('boolean') else: data['column_types'].append('others') if col[5] == 1: data['primary_keys'].append(len(data['column_names'])-1) data["foreign_keys"] = fk_holder data['foreign_keys'] = convert_fk_index(data) return data def get_schema_dict(db, kk=3): """ Get database's schema, which is a dict with table name as key and list of column names as value :param db: database path :return: schema dict """ data = dump_db_json_schema(db,db.split('/')[-1]) tables = data['table_names_original'] column_types = data['column_types'] primary_keys = data['primary_keys'] foreign_keys = data['foreign_keys'] column_names = data['column_names_original'] schema_dict = { 'tables': {}, 'foreign_keys':[] } for i, table in enumerate(tables): t = {} for j, c in enumerate(column_names): if c[0] == i: if j in primary_keys: t[c[1]] = [column_types[j].upper(), True] else: t[c[1]] = [column_types[j].upper(), True] schema_dict['tables'][table] = t for foreign_key in foreign_keys: t1 = tables[column_names[foreign_key[0]][0]] c1 = column_names[foreign_key[0]][1] t2 = tables[column_names[foreign_key[1]][0]] c2 = column_names[foreign_key[1]][1] schema_dict['foreign_keys'].append([t1,c1,t2,c2]) conn = sqlite3.connect(db) cursor = conn.cursor() # get exapmles for table in schema_dict['tables'].keys(): try: select_query = f'SELECT * FROM `{table}` LIMIT {kk}' cursor.execute(select_query) rows = cursor.fetchall() cursor.execute(f"PRAGMA table_info(`{table}`);") columns = [column[1] for column in cursor.fetchall() ] for i, c in enumerate(columns): cls_valuse = [f"{row[i][0:100]}..." if type(row[i]) is str and len(row[i]) > 100 else row[i] for row in rows] schema_dict['tables'][table][c].append(cls_valuse) except Exception as e: print(e) return schema_dict def get_example_str(schema_dict, k=1): tables = list(schema_dict['tables'].keys()) examples = {} for table in tables: table_dict = schema_dict['tables'][table] example = [] for cls in table_dict.keys(): example.append(table_dict[cls][2]) example_str = [] for i, v in enumerate(example[0]): example_str.append(tuple([e[i] for e in example])) if (i+1) == k: break examples[table] = example_str e_s = '' for key in examples.keys(): e_s += f"{key}: " + str(examples[key])+'\n' return e_s[:-1] def get_schmea_str_and_examples(schema_dict): schmea_str = "" tables = list(schema_dict['tables'].keys()) examples = {} for table in tables: if ' ' in table: table_str = f'CREATE TABLE "{table}" (' else: table_str = f"CREATE TABLE {table} (" table_dict = schema_dict['tables'][table] pk_str = '' example = [] for cls in table_dict.keys(): try: cls_ = f'"{cls}"' if ' ' in cls else cls table_str += f"{cls_} {table_dict[cls][0]}, " if table_dict[cls][1]: pk_str += cls_+', ' example.append(table_dict[cls][2]) except Exception as e: print(e) example_str = [] try: for i, v in enumerate(example[0]): example_str.append(tuple([e[i] for e in example])) except Exception as e: print(e) examples[table] = example_str if pk_str != '': table_str += f"PRIMARY KEY({pk_str[:-2]}), " fk_str = '' for fk in schema_dict['foreign_keys']: if fk[0] == table and fk[2] in tables: if fk[3] in schema_dict['tables'][fk[2]].keys(): fk = [f'"{f}"' if ' ' in f else f for f in fk ] fk_str += f'FOREIGN KEY ({fk[1]}) REFERENCES {fk[2]}({fk[3]}), ' if fk_str != '': table_str += fk_str schmea_str += table_str[:-2] +'); ' schmea_str = schmea_str[:-1] e_s = '' for key in examples.keys(): e_s += f"{key}: " + str(examples[key])+'\n' return schmea_str, e_s[:-1] # parse SQL def parse_sql_from_string(input_string): input_string = input_string.replace('\n', ' ').replace('\t','') rs = '' if '```sql' in input_string: try: sql_pattern = r'```sql(.*?)```' all_sqls = [] for match in re.finditer(sql_pattern, input_string, re.DOTALL): all_sqls.append(match.group(1).strip()) if all_sqls: rs = all_sqls[-1] if 'SELECT' not in rs and len(all_sqls)>1: rs = all_sqls[-2] except: None if 'select' in input_string.lower() and rs=='': rs = input_string[input_string.find('SELECT'):] if ';' in rs: # end rs = rs[:input_string.find(';')+1] if rs == '': rs = 'SELECT xx FROM xx' return replace_multiple_spaces(rs).replace('```','') def replace_multiple_spaces(text): return re.sub(r'\s{2,}', ' ', text) def filter_dict_by_sql(schema_dict, sql): schema_dict_ = copy.deepcopy(schema_dict) keys = list(schema_dict_['tables'].keys()) keys.sort(key=lambda x: - len(x)) # tables for table in keys: if f'from {table.lower()}' not in sql.lower() and f'join {table.lower()}' not in sql.lower(): schema_dict_['tables'].pop(table, None) # columns keys = list(schema_dict_['tables'].keys()) keys.sort(key=lambda x: - len(x)) for table in keys: cls_keys = list(schema_dict_['tables'][table].keys()) cls_keys.sort(key=lambda x: - len(x)) tabel_dict = copy.deepcopy(schema_dict_['tables'][table]) for cls in cls_keys: if cls.lower() not in sql.lower(): schema_dict_['tables'][table].pop(cls, None) if len(schema_dict_['tables'][table].keys()) == 0: # schema_dict_['tables'][table] = tabel_dict # for COUNT(*) for cls in tabel_dict.keys(): if tabel_dict[cls][1] == True: schema_dict_['tables'][table][cls] = tabel_dict[cls] if len(schema_dict_['tables'][table].keys()) == 0: schema_dict_['tables'][table][tabel_dict.keys()[0]] = tabel_dict[tabel_dict.keys()[0]] schema_dict_['tables'][table][tabel_dict.keys()[1]] = tabel_dict[tabel_dict.keys()[1]] # for COUNT(*) return schema_dict_ def filter_dict_by_sl(schema_dict, sql): schema_dict_ = copy.deepcopy(schema_dict) keys = list(schema_dict_['tables'].keys()) keys.sort(key=lambda x: - len(x)) # tables for table in keys: if f'{table.lower()}' not in sql.lower(): schema_dict_['tables'].pop(table, None) # columns keys = list(schema_dict_['tables'].keys()) keys.sort(key=lambda x: - len(x)) for table in keys: cls_keys = list(schema_dict_['tables'][table].keys()) cls_keys.sort(key=lambda x: - len(x)) tabel_dict = copy.deepcopy(schema_dict_['tables'][table]) for cls in cls_keys: if cls.lower() not in sql.lower(): schema_dict_['tables'][table].pop(cls, None) if len(schema_dict_['tables'][table].keys()) == 0: # schema_dict_['tables'][table] = tabel_dict # for COUNT(*) for cls in tabel_dict.keys(): if tabel_dict[cls][1] == True: schema_dict_['tables'][table][cls] = tabel_dict[cls] if len(schema_dict_['tables'][table].keys()) == 0: schema_dict_['tables'][table][tabel_dict.keys()[0]] = tabel_dict[tabel_dict.keys()[0]] schema_dict_['tables'][table][tabel_dict.keys()[1]] = tabel_dict[tabel_dict.keys()[1]] # for COUNT(*) return schema_dict_ @func_set_timeout(5) def execute_query_limit(db_path, query): error = '' result = None conn = sqlite3.connect(db_path, timeout=5.0, check_same_thread=False) cursor = conn.cursor() cursor = conn.cursor() cursor.execute(query) result = cursor.fetchone() cursor.close() conn.close() return result, error def execute_query(db_path, query): try: result, error = execute_query_limit(db_path, query) except func_timeout.exceptions.FunctionTimedOut: error = "SQL execution timeout" print("*"*30, error, query) result = None except Exception as e: error = str(e) print("*"*30, error, query) result = None return result, error def replace_syn(data1, data2): for i in range(len(data1)): if data1[i]['question'] == data2[i]['SpiderQuestion']: data1[i]['question'] = data2[i]['SpiderSynQuestion'] return data1 def eval_all(args): dataset= args.dataset mode=args.mode data_tuples = parse_dataset(args.data_path, mode, dataset) batch_size = args.batch_size if dataset == 'spider_syn': data2 = read_json_file(os.path.join(args.data_path, 'spider', f'{mode}_syn.json')) data_tuples = replace_syn(data_tuples,data2) dataset = 'spider' args.tag += '_syn' if dataset == 'spider_DK': args.tag += '_DK' dataset = 'spider' if dataset == 'spider_real': args.tag += '_real' dataset = 'spider' if dataset == 'bird': kk = 5 else: kk = 10 kkkkk = 1 if dataset=='bird' else 3 if 'online' in args.tag: generator = LLM_Online() else: generator = LLM_Model(args.LLM_model) tag = args.tag flg1 = False flg2 = False flg3 = False flg4 = False old_flgs = args.flags args.flags = args.flags.split('_') if args.flags[0] == '1': flg1 = True if args.flags[1] == '1': flg2 = True if args.flags[2] == '1': flg3 = True if args.flags[3] == '1': flg4 = True # generate SQL if True: sql_results = [] data_header = [["NLQ", "Predict", "GOLD", 'database']] prompts = [] for index, row in enumerate(data_tuples): if 'spider' in dataset: row['SQL'] = row['query'] if 'drspider' in dataset: row['SQL'] = row['query'] question, db_id = row['question'], row['db_id'] if dataset == 'spider': if mode == 'test': db_path = os.path.join(args.data_path, dataset, 'test_database', db_id, f"{db_id}.sqlite") else: db_path = os.path.join(args.data_path, dataset, 'database', db_id, f"{db_id}.sqlite") elif dataset == 'drspider': db_path = os.path.join(args.data_path, db_id, f"{db_id}.sqlite") elif dataset == 'bird': db_path = os.path.join(args.data_path, dataset, f'{mode}/{mode}_databases', db_id, f"{db_id}.sqlite") else: raise TypeError(f"Unexpect dataset: {dataset}.") schema_dict = get_schema_dict(db_path, kk = kk) database_schema, examples = get_schmea_str_and_examples(schema_dict) schema_dict_ = schema_dict if dataset == 'bird': prompt = [question, schema_dict, f"\n\n/* Question hint */\n{row['evidence']}" if row['evidence'] != '' else '', schema_dict_] else: prompt = [question, schema_dict, '', schema_dict_] prompts.append([database_schema, str(examples), question, row['SQL'], db_id, prompt, db_path]) n_samples = len(data_tuples) n_batches = (n_samples - 1)//batch_size + 1 for i in range(n_batches): start = i*batch_size end = n_samples if i== n_batches -1 else (i+1)*batch_size batch_prompts = prompts[start: end] schema_dicts = [] # only keep the tables # schema linking if flg1 or flg2: response_strs = None c_response_strs = None if flg1: if args.eval_sft == 1: c_response_strs = generator.generate_response(prompts=[prompt_sl_temp_sft.format(ds=j[0],sr=get_example_str(j[5][1],kkkkk),qs=j[2],hint=j[5][2]) for j in batch_prompts]) else: c_response_strs = generator.generate_response(prompts=[prompt_sl_temp.format(ds=j[0],sr=get_example_str(j[5][1],kkkkk),qs=j[2],hint=j[5][2]) for j in batch_prompts]) if flg2: response_strs = generator.generate_response(prompts=[prompt_temp.format(ds=j[0],sr=get_example_str(j[5][1],kkkkk),qs=j[2],hint=j[5][2]) for j in batch_prompts]) if c_response_strs is None: c_response_strs = response_strs if response_strs is None: response_strs = c_response_strs for j, response_str in enumerate(c_response_strs): schema_dict = batch_prompts[j][5][1] gt_sql = batch_prompts[j][3] # schema_dict_gt = filter_dict_by_sql(batch_prompts[j][5][1], gt_sql) # sl c_sql_str1 = response_str.replace('"',"'").replace('\'',"") schema_dict_1 = filter_dict_by_sl(batch_prompts[j][5][1], c_sql_str1) # pre-sql c_sql_str2 = parse_sql_from_string(response_strs[j]).replace('"',"'").replace('\'',"") schema_dict_2 = filter_dict_by_sql(batch_prompts[j][5][1], c_sql_str2) schema_dict_old = copy.deepcopy(schema_dict) keys1 = schema_dict_1['tables'].keys() keys2 = schema_dict_2['tables'].keys() all_keys = list(schema_dict_old['tables'].keys()) for key in all_keys: if key not in keys1 and key not in keys2: schema_dict_old['tables'].pop(key, None) else: clss = [] if key in keys1: clss += schema_dict_1['tables'][key].keys() if key in keys2: clss += schema_dict_2['tables'][key].keys() clss = list(set(clss)) for k in list(schema_dict_old['tables'][key].keys()): if k not in clss: schema_dict_old['tables'][key].pop(k,None) if len(schema_dict_old['tables'][key].keys()) == 0: schema_dict_old['tables'].pop(key, None) schema_dict_ = schema_dict_old # schema_dict_ = schema_dict_gt # gt schema_dict_table = copy.deepcopy(schema_dict) for key in schema_dict['tables'].keys(): if key not in schema_dict_['tables'].keys(): schema_dict_table['tables'].pop(key,None) schema_dicts.append(schema_dict_table) if j == 0: print("######", response_str, list(schema_dict_old['tables'].keys()) ) ds, sr = get_schmea_str_and_examples(schema_dict_) batch_prompts[j][0] = ds batch_prompts[j][1] = sr else: for j, v in enumerate(batch_prompts): batch_prompts[j][1] = get_example_str(batch_prompts[j][5][1],kkkkk) # text-to-sql final_prompts=[prompt_temp.format(ds=j[0],sr=j[1],qs=j[2],hint=j[5][2]) for j in batch_prompts] response_strs = generator.generate_response(prompts=final_prompts) def contains_subquery(sql_query, tables): sql = sql_query.lower() select_num = 0 join_num = 0 tmp = sql while 'select' in tmp: tmp = tmp[tmp.find('select')+6:] select_num += 1 tmp = sql while 'join' in tmp: tmp = tmp[tmp.find('select')+6:] join_num += 1 table_num = len([key for key in tables if f"from {key.lower()}" in sql or f"join {key.lower()}" in sql]) if table_num == 1: hard = 1 elif table_num==2: hard = 2 else: hard = 3 return hard nc_idx = [] continue_sqls = [] # noisy correction if flg3: predSQLs = [parse_sql_from_string(response_str) for response_str in response_strs] nc_prompts = [] for j in range(len(response_strs)): v = batch_prompts[j] predSQL = predSQLs[j] ds = get_schmea_str_and_examples(v[5][1])[0] sr = get_example_str(v[5][1],kkkkk) ex_hint = execute_query(batch_prompts[j][6], predSQL)[1] if ex_hint != '': ex_hint = f"\n\n/* Execution exception */\n{ex_hint}" # ex_hint = '' if args.eval_sft == 1: nc_prompts.append(prompt_nc_temp_sft.format(ds=ds ,sr=sr, qs=v[2], ex_hint = ex_hint, hint=v[5][2],sql = predSQL)) else: nc_prompts.append(prompt_nc_temp.format(ds=ds ,sr=sr, qs=v[2], ex_hint = ex_hint, hint=v[5][2],sql = predSQL)) response_strs_ = generator.generate_response(prompts=nc_prompts) for idx, v in enumerate(response_strs_): if idx == 0: print("******", nc_prompts[0], '\n', v, batch_prompts[idx][3]) v_lower = v.lower() v_lower = v_lower[:v_lower.find('select')+6] if 'select' in v_lower else v_lower flag = 'select' in v_lower and ('can correctly answer' not in v_lower and 'can answer correctly' not in v_lower and 'is correct' not in v_lower and 'will answer correctly' not in v_lower and 'will correctly answer' not in v_lower and 'can answer ' not in v_lower and 'can accurately answer' not in v_lower and 'can answer accurately' not in v_lower and 'is correct' not in v_lower and 'will answer accurately' not in v_lower and 'will accurately answer' not in v_lower and 'can answer ' not in v_lower ) pre_sql = parse_sql_from_string(response_strs[idx]) if flag: ex_flg2 = True if execute_query(batch_prompts[idx][6], parse_sql_from_string(v))[1] == '' else False # if ex_flg2: if ex_flg2: response_strs[idx] = v pre_sql = parse_sql_from_string(response_strs[idx]) ex_flg3 = True if execute_query(batch_prompts[idx][6], pre_sql)[1] == '' else False hard = contains_subquery(pre_sql, batch_prompts[idx][5][1]['tables'].keys()) if ex_flg3 == False or hard > 2: common_sql = 'SELECT ' continue_sqls.append(common_sql) nc_idx.append(idx) else: for idx, v in enumerate(response_strs): pre_sql = parse_sql_from_string(response_strs[idx]) ex_flg3 = True if execute_query(batch_prompts[idx][6], pre_sql)[1] == '' else False hard = contains_subquery(pre_sql, batch_prompts[idx][5][1]['tables'].keys()) if ex_flg3 == False or hard > 2: common_sql = 'SELECT ' continue_sqls.append(common_sql) nc_idx.append(idx) # continuation writing if flg4: cl_prompts = [] for j, idx in enumerate(nc_idx): v = batch_prompts[idx] ds = get_schmea_str_and_examples(v[5][1])[0] sr = get_example_str(v[5][1],kkkkk) common_sql = continue_sqls[j] if args.eval_sft == 1: cl_prompts.append(prompt_cw_temp_sft.format(ds=ds, sr=sr, qs=v[2],hint=v[5][2], sql = common_sql)) else: cl_prompts.append(prompt_cw_temp.format(ds=ds, sr=sr, qs=v[2],hint=v[5][2], sql = common_sql)) if len(nc_idx) > 0: response_strs_ = generator.generate_response(prompts=cl_prompts) print("%%%%%%%%%%%%%%%%%%",response_strs_[0]) for idx, v in enumerate(nc_idx): if execute_query(batch_prompts[v][6], parse_sql_from_string(response_strs_[idx]))[0] is not None: response_strs[v] = response_strs_[idx] for j, response_str in enumerate(response_strs): database_schema = batch_prompts[j][0] question = batch_prompts[j][2] gt_sql = replace_multiple_spaces(batch_prompts[j][3]) db_id = batch_prompts[j][4] prompt = final_prompts[j] print(f"=={start+j+1}/{len(data_tuples)}=={db_id}=={tag}==================") try: if dataset == 'spider': if mode == 'test': db_path = os.path.join(args.data_path, dataset, 'test_database', db_id, f"{db_id}.sqlite") else: db_path = os.path.join(args.data_path, dataset, 'database', db_id, f"{db_id}.sqlite") elif dataset == 'bird': db_path = os.path.join(args.data_path, dataset, f'{mode}/{mode}_databases', db_id, f"{db_id}.sqlite") else: raise TypeError(f"Unexpect dataset: {dataset}.") SQL_str = parse_sql_from_string(response_str) except Exception as e: res = f'error: {str(e)}' print(res, response_str) sql_results.append([question, SQL_str, gt_sql, db_id]) print(prompt) print(f"Question: {question}") print(f"Raw Resp: {response_str}") print(f"Answer: {SQL_str}") print(f"Ground: {gt_sql}") if SQL_str== 'None': exit() if not os.path.isdir(os.path.join(args.output_path, f"{tag}_{dataset}")): os.makedirs(os.path.join(args.output_path, f"{tag}_{dataset}")) with open(os.path.join(args.output_path, f"{tag}_{dataset}", f"rs_{old_flgs}.csv"), mode='w', newline='', encoding='utf-8') as file: writer = csv.writer(file) writer.writerows(data_header + sql_results) import os import pynvml pynvml.nvmlInit() def usegpu(need_gpu_count=1): nouse=[] for index in range(pynvml.nvmlDeviceGetCount()): # 这里的0是GPU id handle = pynvml.nvmlDeviceGetHandleByIndex(index) meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle) used= meminfo.used/meminfo.total if used<0.3: nouse.append(index) if len(nouse)>=need_gpu_count: os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, nouse[:need_gpu_count])) # return nouse[:need_gpu_count] print(nouse[:need_gpu_count]) return need_gpu_count elif len(nouse)>0: os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, nouse)) return len(nouse) else: return 0 if __name__ == "__main__": parser = argparse.ArgumentParser(description='SQL') parser.add_argument("--dataset", default='spider', type=str) parser.add_argument("--data_path", default='./dataset', type=str) parser.add_argument("--output_path", default='./dataset', type=str) parser.add_argument("--mode", default='dev', type=str) parser.add_argument("--tag", default='0701', type=str) parser.add_argument("--gpus", default=0, type=int) parser.add_argument("--eval_sft", default=1, type=int) parser.add_argument("--flags", default='1_0_0', type=str) parser.add_argument("--LLM_model", default='/disk2/qinyang/qwen2-1.5b-instruct', type=str) parser.add_argument("--batch_size", default=32, type=int) args = parser.parse_args() usegpu(need_gpu_count=args.gpus) print(args) eval_all(args)