✨ add support for classifications in custom docs

Author: ianardee

mindee/documents/custom_document.py

@@ -6,6 +6,8 @@
 class CustomDocument(Document):
     fields: Dict[str, dict]
     """Dictionary of all fields in the document"""
+    classifications: Dict[str, dict]
+    """Dictionary of all classifications in the document"""
 
     def __init__(
         self,
@@ -39,14 +41,24 @@ def _build_from_api_prediction(
         :param page_n: Page number for multi pages pdf input
         """
         self.fields = {}
+        self.classifications = {}
         for field_name in api_prediction:
             field = api_prediction[field_name]
-            field["page_n"] = page_n
-            self.fields[field_name] = field
+            # Only classifications have the 'value' attribute.
+            if "value" in field:
+                self.classifications[field_name] = field
+            # Only value lists have the 'values' attribute.
+            elif "values" in field:
+                field["page_n"] = page_n
+                self.fields[field_name] = field
             setattr(self, field_name, field)
 
     def __str__(self) -> str:
-        custom_doc_str = f"----- {self.type} -----\n"
+        custom_doc_str = (
+            f"----- {self.type} -----\nFilename: {self.filename or ''}".rstrip() + "\n"
+        )
+        for name, info in self.classifications.items():
+            custom_doc_str += f"{name}: {info['value']}\n"
         for name, info in self.fields.items():
             custom_doc_str += "%s: %s\n" % (
                 name,