fix: KeyError in tableformer prediction (#854)

* fix for KeyError in tableformer prediction Signed-off-by: Maksym Lysak <[email protected]> * chore: rewrite cumbersome dictionary checking Signed-off-by: Christoph Auer <[email protected]> --------- Signed-off-by: Maksym Lysak <[email protected]> Signed-off-by: Christoph Auer <[email protected]> Co-authored-by: Maksym Lysak <[email protected]> Co-authored-by: Christoph Auer <[email protected]>
DS4SD · Jan 31, 2025 · b1cf796 · b1cf796
1 parent 70d68b6
commit b1cf796
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/docling/models/table_structure_model.py b/docling/models/table_structure_model.py
@@ -209,12 +209,16 @@ def __call__(
                                     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,