| 67 | |
| 68 | |
| 69 | class Text2SQL: |
| 70 | def __init__(self, logger_name, filter_mode="complex", prompt_mode="v2", n_shots=2, |
| 71 | few_shot_mode="masked_ques_sim"): |
| 72 | self.model = 'gpt-3.5-turbo' |
| 73 | self.ner = NER() |
| 74 | self.filter_characters = FilterCharacters() |
| 75 | self.sql_type = SqlType.SINGLE_TABLE.value |
| 76 | self.table_character_dict = {} |
| 77 | self.filter_characters_mode = args.filter_mode |
| 78 | self.logger = get_logger(logger_name) |
| 79 | self.question_type = QuestionType.EASY.value |
| 80 | self.aggregation_type = AggregationType.NON.value |
| 81 | self.dataset = "" |
| 82 | self.n_shots = args.n_shots |
| 83 | self.few_shot_mode = args.few_shot_mode |
| 84 | self.sql_prompt_fewshot = SQLPromptFewshot() |
| 85 | self.sql_prompt_error = SQLPromptError() |
| 86 | |
| 87 | def get_features_by_llm(self, query, table_info, limitation, metric, main_metric, mode="en"): |
| 88 | numbers = len(limitation.split(",")) |
| 89 | entities = limitation.split(",") |
| 90 | if self.filter_characters_mode == FilterType.COMPLEX.value: |
| 91 | features_prompt = (get_features_prompt() if mode == 'en' else get_features_prompt_cn()) |
| 92 | output_format = json.dumps(generate_dictionary_for_complex(entities)) |
| 93 | prompt_dict = { |
| 94 | "query": query, |
| 95 | "table_info": table_info, |
| 96 | "limitation": limitation, |
| 97 | # "metric": metric, |
| 98 | "main_metric": main_metric, |
| 99 | "numbers": numbers, |
| 100 | "output_format": output_format |
| 101 | } |
| 102 | else: |
| 103 | features_prompt = get_features_prompt_simple() |
| 104 | output_format = json.dumps(generate_dictionary(entities)) |
| 105 | prompt_dict = { |
| 106 | "table_info": table_info, |
| 107 | # "metric": metric, |
| 108 | "limitation": limitation, |
| 109 | "numbers": numbers, |
| 110 | "output_format": output_format |
| 111 | } |
| 112 | features_prompt = get_prompt_content(features_prompt, prompt_dict) |
| 113 | features = ask_llm(features_prompt, args.sc_filter_temp) |
| 114 | try: |
| 115 | features = get_dict_from_str(features) |
| 116 | except Exception as e: |
| 117 | print(f"==============={features}") |
| 118 | print(f"==============={e.args}") |
| 119 | features = "" |
| 120 | return features |
| 121 | |
| 122 | @self_consistency(args.sc_nums_question_label, ["question_type"], "all") |
| 123 | def get_question_type(self, query, table_info): |
| 124 | sql_prompt = get_questions_label_prompt(self.dataset) |
| 125 | prompt_dict = { |
| 126 | "query": query, |