From b1cf796730901222ad0882ff44efa0ef43a743ee Mon Sep 17 00:00:00 2001 From: Maxim Lysak <101627549+maxmnemonic@users.noreply.github.com> Date: Fri, 31 Jan 2025 17:00:14 +0100 Subject: [PATCH] fix: KeyError in tableformer prediction (#854) * fix for KeyError in tableformer prediction Signed-off-by: Maksym Lysak * chore: rewrite cumbersome dictionary checking Signed-off-by: Christoph Auer --------- Signed-off-by: Maksym Lysak Signed-off-by: Christoph Auer Co-authored-by: Maksym Lysak Co-authored-by: Christoph Auer --- docling/models/table_structure_model.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/docling/models/table_structure_model.py b/docling/models/table_structure_model.py index 1f4fbc7..f17cbed 100644 --- a/docling/models/table_structure_model.py +++ b/docling/models/table_structure_model.py @@ -209,12 +209,16 @@ class TableStructureModel(BasePageModel): tc.bbox = tc.bbox.scaled(1 / self.scale) table_cells.append(tc) + assert "predict_details" in table_out + # Retrieving cols/rows, after post processing: - num_rows = table_out["predict_details"]["num_rows"] - num_cols = table_out["predict_details"]["num_cols"] - otsl_seq = table_out["predict_details"]["prediction"][ - "rs_seq" - ] + num_rows = table_out["predict_details"].get("num_rows", 0) + num_cols = table_out["predict_details"].get("num_cols", 0) + otsl_seq = ( + table_out["predict_details"] + .get("prediction", {}) + .get("rs_seq", []) + ) tbl = Table( otsl_seq=otsl_seq,