Fix encoder runtime failures and CoS table selection flow

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 = {

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:

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,

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)