def batch_denormalize(): data_dir = sys.argv[1] schema_graphs = load_schema_graphs_spider(data_dir, 'spider') train_data = load_data_split_spider(data_dir, 'train', schema_graphs) # random.shuffle(train_data) for i, example in enumerate(train_data): if not example.program in complex_queries: continue schema_graph = schema_graphs.get_schema(example.db_id) if DEBUG: denormalizer_unit_test(example.program, schema_graph) else: try: denormalizer_unit_test(example.program, schema_graph) except KeyError as e: example.pretty_print(schema_graph) print(str(e)) import pdb pdb.set_trace() except AssertionError as e: example.pretty_print(schema_graph) print(str(e)) import pdb pdb.set_trace() except pyparsing.ParseException as e: example.pretty_print(schema_graph) print(str(e)) import pdb pdb.set_trace() if i > 0 and i % 500 == 0: print('{} examples processed'.format(i))
def test_restore_clause_order(): in_json = sys.argv[1] data_dir = sys.argv[2] schema_graphs = load_schema_graphs_spider(data_dir) with open(in_json) as f: content = json.load(f) num_errors = 0 for i, example in enumerate(content): sql = example['query'] print('Orig SQL: {}'.format(sql)) db_name = example['db_id'] schema = schema_graphs[db_name] # sn_sql = shallow_normalize(sql, schema) dn_sql, _ = denormalize(sql, schema) print('DN SQL:\t\t{}'.format(dn_sql)) eo_sql = convert_to_execution_order(dn_sql, schema) restored_sql = restore_clause_order(eo_sql, schema) print('Restored SQL:\t{}'.format(restored_sql)) # print('EO Pred SQL: {}'.format(eo_sql)) # print('SN SQL:\t\t{}'.format(sn_sql)) print() if dn_sql != restored_sql: num_errors += 1 import pdb pdb.set_trace() print('{}/{} errors detected'.format(num_errors, len(content)))
def check_foreign_keys_in_queries(): data_dir = sys.argv[1] dataset = sys.argv[2] schema_graphs = load_schema_graphs_spider(data_dir) train_data = load_data_split_spider(data_dir, 'train', schema_graphs) dev_data = load_data_split_spider(data_dir, 'dev', schema_graphs) tables = load_spider_tables(os.path.join(data_dir, 'tables.json')) table_dict = dict() for table in tables: table_dict[table['db_id']] = table in_json = os.path.join(data_dir, '{}.parsed.json'.format(dataset)) with open(in_json) as f: parsed_sqls = json.load(f) for i, example in enumerate(train_data + dev_data): schema_graph = schema_graphs.get_schema(example.db_id) ast, _ = get_ast(example.program, parsed_sqls, denormalize_sql=True, schema_graph=schema_graph) foreign_keys_readable, foreign_keys = extract_foreign_keys( ast, schema_graph) for f_key in foreign_keys: if not tuple(sorted(f_key)) in schema_graph.foreign_keys: print(example.program) print(json.dumps(ast, indent=4)) print('Missing foreign key detected:') print('- {}'.format(schema_graph.get_field_signature( f_key[0]))) print('- {}'.format(schema_graph.get_field_signature( f_key[1]))) import pdb pdb.set_trace()
def batch_tokenize(data_dir, db_name, denormalization=False): print('Tokenizing {}'.format(db_name)) num_tokens_list = [] if db_name == 'spider': schema_graphs = load_schema_graphs_spider(data_dir, 'spider') in_json = os.path.join(data_dir, 'dev.json') with open(in_json) as f: content = json.load(f) for i, example in enumerate(content): sql = example['query'] if sql.endswith(';'): sql = sql[:-1] db_name = example['db_id'] schema = schema_graphs[db_name] ast = parse(sql) denormalized_ast, _ = denormalize(ast, schema, return_parse_tree=True) tokens, token_types, constants = \ sql_tokenize(denormalized_ast, bu.tokenizer.tokenize, return_token_types=True, schema=schema, keep_singleton_fields=True, atomic_value=True, num_token=' <NUM> ', str_token=' <STRING> ') num_tokens_list.append(len(tokens)) else: return NotImplementedError print('{}: avg # tokens = {}'.format(db_name, np.mean(num_tokens_list)))
def test_schema_consistency(): data_dir = sys.argv[1] db_name = 'flight_2' schema_graphs = load_schema_graphs_spider(data_dir, 'spider') schema = schema_graphs[db_name] schema.pretty_print() in_sql = 'SELECT singer.Name FROM concert JOIN singer_in_concert ON singer_in_concert.Singer_ID = singer.Singer_ID WHERE concert.Year = 2014' in_sql = 'SELECT singer.concert FROM singer WHERE singer.age > (SELECT avg(singer.age) FROM singer)' in_sql = 'SELECT singer.Name, singer.Country FROM singer INTERSECT SELECT singer.Name, singer.Country, singer.Age FROM singer WHERE singer.Age = "?" ORDER BY singer.Age DESC' in_sql = 'SELECT COUNT(*) FROM singer' in_sql = 'SELECT concert.concert_Name, concert.Theme, COUNT(*) FROM concert GROUP BY concert.Theme, concert.Theme' in_sql = 'SELECT T2.name FROM singer_in_concert AS T1 JOIN singer AS T2 ON T1.singer_id = T2.singer_id JOIN concert AS T3 ON T1.concert_id = T3.concert_id WHERE T3.year = 2014' in_sql = 'SELECT AIRPORTS.AirportCode FROM AIRPORTS JOIN FLIGHTS ON AIRPORTS.AirportCode = FLIGHTS.DestAirport OR AIRPORTS.AirportCode = FLIGHTS.SourceAirport GROUP BY AIRPORTS.AirportCode ORDER BY COUNT(*) DESC LIMIT 1' # in_sql = 'from singer select singer.Name , singer.Country union from singer where singer.Age = "age" select singer.Name , singer.Country , singer.Age order by singer.Age desc' # in_sql = 'from stadium join concert on stadium.Stadium_ID = concert.Stadium_ID where stadium.Capacity = (from stadium select max (stadium.Capacity)) select count (*)' # in_sql = 'from singer join singer_in_concert on singer.Singer_ID = singer_in_concert.Singer_ID join concert on singer_in_concert.Singer_ID = concert.concert_ID where concert.Year = 2014 select singer.Name' # in_sql = 'from Students select Students.other_student_details order by Students.other_student_details desc limit <UNK>' # in_sql = 'from Sections join Sections on Addresses.address_id = * where Sections.section_name = "h" select Sections.section_name' # in_sql = 'from stadium join concert on stadium.Stadium_ID = concert.Stadium_ID where concert.Year = 2014 select stadium.Name , stadium.Location intersect from concert join stadium on stadium.Stadium_ID = stadium.Stadium_ID where concert.Year = 2015 select stadium.Name , stadium.Location' # in_sql, _ = denormalize(in_sql, schema) ast = parse(in_sql) check_schema_consistency(ast, schema, verbose=True)
def test_execution_order(): in_sql = sys.argv[1] in_sql = "SELECT song_name FROM singer WHERE age > (SELECT avg(age) FROM singer)" data_dir = sys.argv[2] db_name = sys.argv[3] schema_graphs = load_schema_graphs_spider(data_dir, 'spider') in_sqls = complex_queries[:4] db_names = ['flight_4', 'academic', 'baseball_1', 'voter_2'] for db_name, in_sql in zip(db_names, in_sqls): schema = schema_graphs[db_name] ast = parse(in_sql) # print(json.dumps(ast, indent=4)) ast_c = copy.deepcopy(ast) eo_sql = format(ast, schema, in_execution_order=True) eo_tokens = tokenize(in_sql, bu.tokenizer.tokenize, schema=schema, in_execution_order=True) print('in_sql: {}'.format(in_sql)) print('eo_sql: {}'.format(eo_sql)) # print('eo_tokens: {}'.format(eo_tokens)) eo_ast = eo_parse(eo_sql) assert (json.dumps(ast_c, sort_keys=True) == json.dumps(eo_ast, sort_keys=True)) # print(json.dumps(eo_ast, indent=4)) restored_sql = format(eo_ast, schema) # print('restored_sql: {}'.format(restored_sql)) print()
def load_data_spider(args): """ Load the Spider dataset released by Yu et. al. 2018. """ in_dir = args.data_dir dataset = dict() schema_graphs = load_schema_graphs_spider( in_dir, 'spider', augment_with_wikisql=args.augment_with_wikisql, db_dir=args.db_dir) dataset['train'] = load_data_split_spider( in_dir, 'train', schema_graphs, get_data_augmentation_tag(args), augment_with_wikisql=args.augment_with_wikisql) dataset['dev'] = load_data_split_spider( in_dir, 'dev', schema_graphs, augment_with_wikisql=args.augment_with_wikisql) dataset['schema'] = schema_graphs fine_tune_set = load_data_split_spider( in_dir, 'fine-tune', schema_graphs, augment_with_wikisql=args.augment_with_wikisql) if fine_tune_set: dataset['fine-tune'] = fine_tune_set return dataset
def test_value_extractor(): in_sql = 'SELECT singer.Name, singer.Country FROM singer INTERSECT SELECT singer.Name, singer.Country, singer.Age FROM singer WHERE singer.Age = "?" ORDER BY singer.Age DESC' in_sql = 'SELECT avg(age) FROM Student WHERE StuID IN ( SELECT T1.StuID FROM Has_allergy AS T1 JOIN Allergy_Type AS T2 ON T1.Allergy = T2.Allergy WHERE T2.allergytype = "animal" INTERSECT SELECT T1.StuID FROM Has_allergy AS T1 JOIN Allergy_Type AS T2 ON T1.Allergy = T2.Allergy WHERE T2.allergytype = "animal")' in_sql = 'SELECT DISTINCT T1.age FROM management AS T2 JOIN head AS T1 ON T1.head_id = T2.head_id WHERE T2.temporary_acting = \'Yes\'' in_sql = 'SELECT t3.title FROM authors AS t1 JOIN authorship AS t2 ON t1.authid = t2.authid JOIN papers AS t3 ON t2.paperid = t3.paperid JOIN inst AS t4 ON t2.instid = t4.instid WHERE t4.country = \"USA\" AND t2.authorder = 2 AND t1.lname = \"Turon\"' data_dir = sys.argv[1] db_name = sys.argv[2] schema_graphs = load_schema_graphs_spider(data_dir, 'spider') schema = schema_graphs[db_name] value_extractor_unit_test(in_sql, schema)
def test_restore_clause_order(): in_sql = 'from (from countrylanguage where countrylanguage.Language = "spanish" select max (countrylanguage.Percentage)) as T0 JOIN countrylanguage ON T0. = countrylanguage.Percentage select count (*)' in_sql = 'from poker_player where poker_player.Earnings = (from poker_player select sum (poker_player.Earnings)) select poker_player.Money_Rank order by poker_player.Earnings desc limit 1' in_sql = 'from Student join Has_Pet on Student.StuID = Has_Pet.StuID join Pets on Has_Pet.PetID = Pets.PetID where Pets.PetType = "cat" select Student.Major , Student.Major' data_dir = sys.argv[1] db_name = sys.argv[2] schema_graphs = load_schema_graphs_spider(data_dir, 'spider') schema = schema_graphs[db_name] print('eo_sql: {}'.format(in_sql)) restored_sql = restore_clause_order(in_sql, schema) print('restored_sql: {}'.format(restored_sql))
def ensemble(): text_tokenize, program_tokenize, post_process, table_utils = tok.get_tokenizers( args) schema_graphs = schema_loader.load_schema_graphs_spider( args.codalab_data_dir, 'spider', db_dir=args.codalab_db_dir) schema_graphs.lexicalize_graphs(tokenize=text_tokenize, normalized=(args.model_id in [utils.BRIDGE])) text_vocab = Vocabulary('text', func_token_index=functional_token_index, tu=table_utils) for v in table_utils.tokenizer.vocab: text_vocab.index_token(v, True, table_utils.tokenizer.convert_tokens_to_ids(v)) program_vocab = sql_reserved_tokens if args.pretrained_transformer else sql_reserved_tokens_revtok vocabs = {'text': text_vocab, 'program': program_vocab} examples = data_loader.load_data_split_spider(args.codalab_data_dir, 'dev', schema_graphs) print('{} {} examples loaded'.format(len(examples), 'dev')) for i, example in enumerate(examples): schema_graph = schema_graphs.get_schema(example.db_id) preprocess_example('dev', example, args, None, text_tokenize, program_tokenize, post_process, table_utils, schema_graph, vocabs) print('{} {} examples processed'.format(len(examples), 'dev')) checkpoint_paths = [ 'ensemble_models/model1.tar', 'ensemble_models/model2.tar', 'ensemble_models/model3.tar' ] sps = [EncoderDecoderLFramework(args) for _ in checkpoint_paths] for i, checkpoint_path in enumerate(checkpoint_paths): sps[i].schema_graphs = schema_graphs sps[i].load_checkpoint(checkpoint_path) sps[i].cuda() sps[i].eval() out_dict = sps[0].inference( examples, restore_clause_order=args.process_sql_in_execution_order, check_schema_consistency_=args.sql_consistency_check, inline_eval=False, model_ensemble=[sp.mdl for sp in sps], verbose=False) assert (sps[0].args.prediction_path is not None) out_txt = sps[0].args.prediction_path with open(out_txt, 'w') as o_f: for pred_sql in out_dict['pred_decoded']: o_f.write('{}\n'.format(pred_sql[0])) print('Model predictions saved to {}'.format(out_txt))
def test_tokenizer(): # for sql in complex_queries: if True: sql = 'SELECT avg(age) FROM Student WHERE StuID IN ( SELECT T1.StuID FROM Has_allergy AS T1 JOIN Allergy_Type AS T2 ON T1.Allergy = T2.Allergy WHERE T2.allergytype = "food" INTERSECT SELECT T1.StuID FROM Has_allergy AS T1 JOIN Allergy_Type AS T2 ON T1.Allergy = T2.Allergy WHERE T2.allergytype = "animal")' sql = 'SELECT T1.Name FROM Tourist_Attractions AS T1 JOIN VISITORS AS T2 JOIN VISITS AS T3 ON T1.Tourist_Attraction_ID = T3.Tourist_Attraction_ID AND T2.Tourist_ID = T3.Tourist_ID WHERE T2.Tourist_Details = "Vincent" INTERSECT SELECT T1.Name FROM Tourist_Attractions AS T1 JOIN VISITORS AS T2 JOIN VISITS AS T3 ON T1.Tourist_Attraction_ID = T3.Tourist_Attraction_ID AND T2.Tourist_ID = T3.Tourist_ID WHERE T2.Tourist_Details = "Marcelle"' sql = "SELECT Perpetrator_ID FROM perpetrator WHERE Year IN ('1995.0', '1994.0', '1982.0')" print(sql) data_dir = sys.argv[1] db_name = sys.argv[2] schema_graphs = load_schema_graphs_spider(data_dir, 'spider') schema = schema_graphs[db_name] tokens = tokenize(sql, bu.tokenizer.tokenize, in_execution_order=True, schema=schema)[0] print(tokens) print()