diff --git a/Spreadsheet_LLM_Encoder.py b/Spreadsheet_LLM_Encoder.py index 4f7a6a6..1fa8796 100644 --- a/Spreadsheet_LLM_Encoder.py +++ b/Spreadsheet_LLM_Encoder.py @@ -309,10 +309,6 @@ def find_boundary_candidates(sheet): col_candidates.add(c) col_candidates.add(c + 1) - # Filter out candidates that are part of a detected header region - header_rows = {idx for idx in range(1, sheet.max_row + 1) if is_header_row(sheet, idx)} - row_candidates = {r for r in row_candidates if r not in header_rows} - # Step 2: Compose candidate boundaries candidates = [] if row_candidates and col_candidates: @@ -787,6 +783,93 @@ def aggregate_formats(sheet, format_map): return dict(aggregated_formats) +def aggregate_regions_dfs(sheet, format_map): + """Aggregate connected cells by semantic key using DFS.""" + aggregated_regions = {} + + for key, cells in format_map.items(): + coords = set() + for cell_ref in cells: + try: + col_letter, row = split_cell_ref(cell_ref) + col = get_column_index(col_letter) + except Exception: + continue + coords.add((row, col)) + + if not coords: + aggregated_regions[key] = [] + continue + + visited = set() + regions = [] + + for start in sorted(coords): + if start in visited: + continue + + stack = [start] + component = set() + + while stack: + row, col = stack.pop() + if (row, col) in visited or (row, col) not in coords: + continue + visited.add((row, col)) + component.add((row, col)) + stack.extend([ + (row - 1, col), + (row + 1, col), + (row, col - 1), + (row, col + 1), + ]) + + component_rows = {} + for row, col in sorted(component): + component_rows.setdefault(row, []).append(col) + + pending = [] + for row in sorted(component_rows): + cols = sorted(component_rows[row]) + start_col = cols[0] + prev_col = cols[0] + for col in cols[1:]: + if col == prev_col + 1: + prev_col = col + continue + pending.append([row, row, start_col, prev_col]) + start_col = col + prev_col = col + pending.append([row, row, start_col, prev_col]) + + merged = [] + for row_start, row_end, col_start, col_end in pending: + extended = False + for existing in merged: + if ( + existing[2] == col_start + and existing[3] == col_end + and existing[1] == row_start - 1 + ): + existing[1] = row_end + extended = True + break + if not extended: + merged.append([row_start, row_end, col_start, col_end]) + + for row_start, row_end, col_start, col_end in merged: + start_ref = f"{get_column_letter(col_start)}{row_start}" + end_ref = f"{get_column_letter(col_end)}{row_end}" + if start_ref == end_ref: + regions.append(start_ref) + else: + regions.append(f"{start_ref}:{end_ref}") + + aggregated_regions[key] = regions + + return aggregated_regions + + def cluster_numeric_ranges(sheet, format_map): """Aggregate numeric cells with identical formatting into ranges.""" numeric_map = { diff --git a/chain_of_spreadsheet.py b/chain_of_spreadsheet.py index 57b3dde..c358a83 100644 --- a/chain_of_spreadsheet.py +++ b/chain_of_spreadsheet.py @@ -48,7 +48,7 @@ def identify_table(encoding: Dict, query: str) -> Optional[str]: Identifies the most relevant table for a query using an LLM. (CoS Stage 1) """ - sheet_name = _find_relevant_sheet(encoding, query) + sheet_name = find_relevant_sheet(encoding, query) if not sheet_name: logger.warning("Could not identify a relevant sheet for the query.") return None @@ -71,7 +71,7 @@ def identify_table(encoding: Dict, query: str) -> Optional[str]: logger.warning(f"Could not parse table range from LLM response: {llm_response}") return None -def _find_relevant_sheet(encoding: Dict, query: str) -> Optional[str]: +def find_relevant_sheet(encoding: Dict, query: str) -> Optional[str]: """Helper to find the most relevant sheet using simple keyword matching.""" query_tokens = {t.lower() for t in query.split()} best_score = 0 @@ -86,7 +86,21 @@ def _find_relevant_sheet(encoding: Dict, query: str) -> Optional[str]: if score > best_score: best_score = score best_sheet = sheet_name - return best_sheet + if best_sheet: + return best_sheet + + sheet_names = list(encoding.get("sheets", {})) + if len(sheet_names) == 1: + # Fall back to the only available sheet so the CoS flow can still run + # when token matching finds nothing useful or there is no competing + # sheet to disambiguate against. + return sheet_names[0] + return None + + +def _find_relevant_sheet(encoding: Dict, query: str) -> Optional[str]: + """Deprecated wrapper; use find_relevant_sheet directly.""" + return find_relevant_sheet(encoding, query) def generate_response(sheet_data: Dict, query: str) -> str: diff --git a/evaluation.py b/evaluation.py index eb511cd..bf5eb69 100644 --- a/evaluation.py +++ b/evaluation.py @@ -49,7 +49,11 @@ def load_spreadsheet_dataset(path: str) -> List[Dict[str, object]]: with open(ann_path, 'r') as f: annotations = json.load(f) - bboxes = [range_to_bbox(t['range']) for t in annotations.get("tables", [])] + tables = annotations.get("tables", []) + if not tables: + continue + + bboxes = [range_to_bbox(t['range']) for t in tables] dataset.append({ "spreadsheet_path": spreadsheet_path, diff --git a/example_chain_usage.py b/example_chain_usage.py index 193e105..628bcb8 100644 --- a/example_chain_usage.py +++ b/example_chain_usage.py @@ -1,7 +1,7 @@ """Example for the two-stage chain-of-spreadsheet pipeline.""" import sys from Spreadsheet_LLM_Encoder import spreadsheet_llm_encode -from chain_of_spreadsheet import identify_table, generate_response +from chain_of_spreadsheet import find_relevant_sheet, identify_table, generate_response def main(): @@ -16,14 +16,16 @@ def main(): print("Failed to encode spreadsheet") return - table_name = identify_table(encoding, query) - if not table_name: + sheet_name = find_relevant_sheet(encoding, query) + table_range = identify_table(encoding, query) + if not sheet_name or not table_range: print("Could not identify a relevant table") return - sheet_data = encoding["sheets"][table_name] + sheet_data = encoding["sheets"][sheet_name] answer = generate_response(sheet_data, query) - print(f"Selected sheet: {table_name}") + print(f"Selected sheet: {sheet_name}") + print(f"Selected range: {table_range}") print(answer)