| 307 | return input_dict |
| 308 | |
| 309 | def step_get_fewshots(self, input_dict, pre_sql, mode): |
| 310 | query = input_dict["query"] |
| 311 | limitation = ",".join(input_dict["ner_results"]["limitation"]) |
| 312 | if args.has_error_case: |
| 313 | input_dict["few_shots"] = self.sql_prompt_error.error_case |
| 314 | return input_dict |
| 315 | |
| 316 | embedding_base_model = args.embedding_base_model |
| 317 | if self.few_shot_mode in ["random", "ques_sim", "masked_ques_sim", "query_sim"]: |
| 318 | if not args.reduce_ql: |
| 319 | few_shots_dict = get_fewshots( |
| 320 | question=query, |
| 321 | entities=limitation, |
| 322 | sql=pre_sql, |
| 323 | question_type=list(set([self.question_type])), |
| 324 | mode=self.few_shot_mode, |
| 325 | n_shots=self.n_shots, |
| 326 | model=embedding_base_model, |
| 327 | index_version=args.few_shot_data, |
| 328 | ques_type_mode='all' |
| 329 | ) |
| 330 | else: |
| 331 | few_shots_dict = get_fewshots( |
| 332 | question=query, |
| 333 | entities=limitation, |
| 334 | sql=pre_sql, |
| 335 | question_type=[QuestionType.EASY.value, QuestionType.NESTED.value, |
| 336 | QuestionType.JOIN.value, QuestionType.JOIN_NESTED.value], |
| 337 | mode=self.few_shot_mode, |
| 338 | n_shots=self.n_shots, |
| 339 | model=embedding_base_model, |
| 340 | index_version=args.few_shot_data, |
| 341 | ques_type_mode='all' |
| 342 | ) |
| 343 | few_shots = "Some example questions and corresponding SQL queries are provided based on similar problems:\n" |
| 344 | # few_shots_dict.reverse() |
| 345 | for item_ in few_shots_dict: |
| 346 | few_shots += f'{item_["query"]}\n' |
| 347 | temp_sql = {"sql": item_["gold_sql"]} |
| 348 | few_shots += f'{json.dumps(temp_sql, ensure_ascii=False)}\n\n' |
| 349 | |
| 350 | else: |
| 351 | few_shots = "" |
| 352 | input_dict["few_shots"] = few_shots |
| 353 | |
| 354 | if mode == "debug": |
| 355 | print("---------------get fewshots---------------------") |
| 356 | print(few_shots) |
| 357 | |
| 358 | return input_dict |
| 359 | |
| 360 | def step_get_pre_sql(self, input_dict, mode): |
| 361 | """Generate SQL for fewshot=query_sim to use""" |