in nl2sql/tasks/table_selection/core.py [0:0]
def CURATED_FEW_SHOT_COT_PROMPT(self) -> _CoreTableSelectorPrompt:
def greedy_post_processor(
intermediate_steps: List[Dict[str, Any]], available_tables: set[str]
) -> set[str]:
"""Implements greedy post-processing for yes/no table selection."""
selected_tables = []
for step in intermediate_steps:
if "yes." in step["raw_response"].lower():
selected_tables.append(step["table"])
filtered_tables = set(selected_tables).intersection(available_tables)
return filtered_tables
return _CoreTableSelectorPrompt(
prompt_id="TASK_TABLE_SELECTION_CORE_V1_SPIDER_V1",
prompt_template=FewShotPrompts.TASK_TABLE_SELECTION_CORE_V1_SPIDER_V1,
call_for_each_table=True,
post_processor=lambda x: yes_no_classifier(x) == "True",
greedy_post_processor=greedy_post_processor,
)