| 537 | args.progress_bar_for_each_datapoint) |
| 538 | |
| 539 | def main(self, root_dir, dataset, file, save_file_name, mode="dev", lang_mode="cn", prompt_mode='openai', |
| 540 | sample="True", data_fold="1", test_id=46, insert_value=0, step_name="all"): |
| 541 | assert mode in ["debug", "dev"] |
| 542 | assert dataset in ["spider", "bird"] |
| 543 | assert lang_mode in ["cn", "en"] |
| 544 | assert prompt_mode in ["create", "openai"] |
| 545 | from data_preprocess import generate_db_prompt_spider, generate_db_prompt_bird |
| 546 | import os |
| 547 | from tqdm import tqdm |
| 548 | self.dataset = dataset |
| 549 | if dataset == "bird": |
| 550 | dataset = "dev" |
| 551 | global special_for_bird |
| 552 | special_for_bird = True |
| 553 | with open(os.path.join(root_dir, dataset, file)) as f: |
| 554 | data = json.load(f) |
| 555 | idx = list(range(len(data))) |
| 556 | if dataset == "dev": |
| 557 | dataset = "bird" |
| 558 | if mode == "debug": |
| 559 | data = [data[test_id]] |
| 560 | |
| 561 | """Start using after test failure""" |
| 562 | if args.re_run: |
| 563 | data = data[args.re_run_idx:] |
| 564 | idx = idx[args.re_run_idx:] |
| 565 | |
| 566 | new_data = list(zip(idx, data)) |
| 567 | for cur_idx, items in enumerate(tqdm(new_data)): |
| 568 | if args.re_run: |
| 569 | cur_idx += args.re_run_idx |
| 570 | origin_idx, item = items |
| 571 | db_id, query = item['db_id'], item['question'] |
| 572 | hint, question_id = None, None |
| 573 | if dataset == "spider": |
| 574 | init_table_infos, table_list = generate_db_prompt_spider( |
| 575 | root_dir='data', dataset='spider', db_id=db_id, prompt_db=args.schema_mode, |
| 576 | normalization=args.fk_mode, limit_value=insert_value) |
| 577 | table_info_list = init_table_infos.split("\n\n") |
| 578 | table_info_list = [i for i in table_info_list if i.strip() != ""] |
| 579 | elif dataset == "bird": |
| 580 | init_table_infos, table_list = generate_db_prompt_bird( |
| 581 | root_dir='data', dataset='dev', db_id=db_id, limit_value=0) |
| 582 | table_info_list = init_table_infos.split("\n\n\n") |
| 583 | table_info_list = [i.strip() for i in table_info_list if i.strip() != ""] |
| 584 | hint = str(item["evidence"]) |
| 585 | |
| 586 | input_dict = {} |
| 587 | input_dict["table_list"] = table_list |
| 588 | input_dict["table_info_list"] = table_info_list |
| 589 | # Easy to debug |
| 590 | input_dict['cur_idx'] = str(cur_idx) |
| 591 | input_dict['origin_idx'] = str(origin_idx) |
| 592 | |
| 593 | if hint is not None: |
| 594 | input_dict["hint"] = hint |
| 595 | input_dict["query"] = query |
| 596 | input_dict["init_table_infos"] = init_table_infos |