Move MaxPagesForClassification from template parameters to config

Author: Taniya Mathur

config_library/pattern-2/bank-statement-sample/config.yaml

@@ -68,6 +68,7 @@ classes:
         description: List of all transactions in the statement period
         attributeType: list
 classification:
+  maxPagesForClassification: "ALL"
   image:
     target_height: ''
     target_width: ''

config_library/pattern-2/lending-package-sample/config.yaml

@@ -914,6 +914,7 @@ classes:
         attributeType: group
 classification:
   classificationMethod: multimodalPageLevelClassification
+  maxPagesForClassification: "ALL"
   image:
     target_height: ''
     target_width: ''

config_library/pattern-2/rvl-cdip-package-sample-with-few-shot-examples/config.yaml

@@ -647,6 +647,7 @@ classes:
         imagePath: config_library/pattern-2/few_shot_example_with_multimodal_page_classification/example-images/bank-statement-pages/
 
 classification:
+  maxPagesForClassification: "ALL"
   image:
     target_height: ''
     target_width: ''

config_library/pattern-2/rvl-cdip-package-sample/config.yaml

@@ -307,6 +307,7 @@ classes:
       - name: comments
         description: Additional notes or remarks about the document. Look for sections labeled 'notes', 'remarks', or 'comments'.
 classification:
+  maxPagesForClassification: "ALL"
   image:
     target_height: ''
     target_width: ''

docs/classification.md

@@ -181,6 +181,22 @@ When deciding between Text-Based Holistic Classification and MultiModal Page-Lev
 
 ## Customizing Classification in Pattern 2
 
+### Configuration Settings
+
+#### Page Limit Configuration
+
+Control how many pages are used for classification:
+
+```yaml
+classification:
+  maxPagesForClassification: "ALL"  # Default: use all pages
+  # Or: "1", "2", "3", etc. - use only first N pages
+```
+
+**Important**: When set to a number (e.g., `"3"`), only the first N pages are classified, but the result is applied to ALL pages in the document. This forces the entire document to be assigned a single class with one section.
+
+### Prompt Components
+
 In Pattern 2, you can customize classification behavior through various prompt components:
 
 ### System Prompts

lib/idp_common_pkg/tests/test_max_pages_classification.py

@@ -8,32 +8,33 @@
 from idp_common.models import Document, Page
 
 
+@pytest.mark.unit
 class TestMaxPagesForClassification:
     @pytest.fixture
     def mock_config(self):
         return {
             "classification": {
                 "maxPagesForClassification": "ALL",
                 "classificationMethod": "multimodalPageLevelClassification",
+                "model": "us.amazon.nova-pro-v1:0",
+                "system_prompt": "Test system prompt",
+                "task_prompt": "Test task prompt",
             }
         }
 
     @pytest.fixture
     def classification_service(self, mock_config):
-        with patch(
-            "idp_common.classification.service.get_config", return_value=mock_config
-        ):
-            return ClassificationService(backend="bedrock")
+        return ClassificationService(backend="bedrock", config=mock_config)
 
     @pytest.fixture
     def sample_document(self):
         doc = Document(id="test-doc")
         doc.pages = {
-            "1": Page(page_id="1", page_number=1),
-            "2": Page(page_id="2", page_number=2),
-            "3": Page(page_id="3", page_number=3),
-            "4": Page(page_id="4", page_number=4),
-            "5": Page(page_id="5", page_number=5),
+            "1": Page(page_id="1"),
+            "2": Page(page_id="2"),
+            "3": Page(page_id="3"),
+            "4": Page(page_id="4"),
+            "5": Page(page_id="5"),
         }
         return doc
 
@@ -84,16 +85,16 @@ def test_apply_limited_classification_single_type(self, classification_service):
         # Original document with 3 pages
         original_doc = Document(id="original")
         original_doc.pages = {
-            "1": Page(page_id="1", page_number=1),
-            "2": Page(page_id="2", page_number=2),
-            "3": Page(page_id="3", page_number=3),
+            "1": Page(page_id="1"),
+            "2": Page(page_id="2"),
+            "3": Page(page_id="3"),
         }
 
         # Classified document with 2 pages, both classified as "invoice"
         classified_doc = Document(id="classified")
         classified_doc.pages = {
-            "1": Page(page_id="1", page_number=1),
-            "2": Page(page_id="2", page_number=2),
+            "1": Page(page_id="1"),
+            "2": Page(page_id="2"),
         }
         classified_doc.pages["1"].classification = "invoice"
         classified_doc.pages["2"].classification = "invoice"
@@ -122,17 +123,17 @@ def test_apply_limited_classification_tie_breaker(self, classification_service):
         # Original document with 4 pages
         original_doc = Document(id="original")
         original_doc.pages = {
-            "1": Page(page_id="1", page_number=1),
-            "2": Page(page_id="2", page_number=2),
-            "3": Page(page_id="3", page_number=3),
-            "4": Page(page_id="4", page_number=4),
+            "1": Page(page_id="1"),
+            "2": Page(page_id="2"),
+            "3": Page(page_id="3"),
+            "4": Page(page_id="4"),
         }
 
         # Classified document with 2 pages, different classifications
         classified_doc = Document(id="classified")
         classified_doc.pages = {
-            "1": Page(page_id="1", page_number=1),
-            "2": Page(page_id="2", page_number=2),
+            "1": Page(page_id="1"),
+            "2": Page(page_id="2"),
         }
         classified_doc.pages["1"].classification = "payslip"
         classified_doc.pages["2"].classification = "drivers_license"
@@ -164,7 +165,7 @@ def test_apply_limited_classification_tie_breaker(self, classification_service):
     def test_apply_limited_classification_empty_sections(self, classification_service):
         """Test handling of empty sections"""
         original_doc = Document(id="original")
-        original_doc.pages = {"1": Page(page_id="1", page_number=1)}
+        original_doc.pages = {"1": Page(page_id="1")}
 
         classified_doc = Document(id="classified")
         classified_doc.sections = []
@@ -176,27 +177,31 @@ def test_apply_limited_classification_empty_sections(self, classification_servic
         # Should return original document unchanged
         assert result == original_doc
 
-    @patch("idp_common.classification.service.get_config")
-    def test_config_integration(self, mock_get_config):
+    def test_config_integration(self):
         """Test that maxPagesForClassification is read from config"""
-        mock_get_config.return_value = {
+        mock_config = {
             "classification": {
                 "maxPagesForClassification": "2",
                 "classificationMethod": "multimodalPageLevelClassification",
+                "model": "us.amazon.nova-pro-v1:0",
+                "system_prompt": "Test system prompt",
+                "task_prompt": "Test task prompt",
             }
         }
 
-        service = ClassificationService(backend="bedrock")
+        service = ClassificationService(backend="bedrock", config=mock_config)
         assert service.max_pages_for_classification == "2"
 
-    @patch("idp_common.classification.service.get_config")
-    def test_config_default_value(self, mock_get_config):
+    def test_config_default_value(self):
         """Test default value when maxPagesForClassification not in config"""
-        mock_get_config.return_value = {
+        mock_config = {
             "classification": {
-                "classificationMethod": "multimodalPageLevelClassification"
+                "classificationMethod": "multimodalPageLevelClassification",
+                "model": "us.amazon.nova-pro-v1:0",
+                "system_prompt": "Test system prompt",
+                "task_prompt": "Test task prompt",
             }
         }
 
-        service = ClassificationService(backend="bedrock")
+        service = ClassificationService(backend="bedrock", config=mock_config)
         assert service.max_pages_for_classification == "ALL"

patterns/pattern-2/template.yaml

@@ -91,10 +91,7 @@ Parameters:
       - "false"
     Description: "Enable Human In The Loop (A2I) for document review"
 
-  MaxPagesForClassification:
-    Type: String
-    Default: "ALL"
-    Description: "Number of pages to use for document classification"
+
 
   SageMakerA2IReviewPortalURL:
     Type: String
@@ -976,7 +973,6 @@ Resources:
       ServiceToken: !Ref UpdateConfigurationFunctionArn
       Default: !Ref ConfigurationDefaultS3Uri
       ConfigLibraryHash: !Ref ConfigLibraryHash
-      MaxPagesForClassification: !Ref MaxPagesForClassification
       CustomClassificationModelARN: !Ref CustomClassificationModelARN
       CustomExtractionModelARN: !Ref CustomExtractionModelARN
 

patterns/pattern-3/template.yaml

@@ -105,10 +105,7 @@ Parameters:
       - "false"
     Description: "Enable Human In The Loop (A2I) for document review"
 
-  MaxPagesForClassification:
-    Type: String
-    Default: "ALL"
-    Description: "Number of pages to use for document classification"
+
 
   SageMakerA2IReviewPortalURL:
     Type: String
@@ -878,7 +875,6 @@ Resources:
       ServiceToken: !Ref UpdateConfigurationFunctionArn
       Default: !Ref ConfigurationDefaultS3Uri
       ConfigLibraryHash: !Ref ConfigLibraryHash
-      MaxPagesForClassification: !Ref MaxPagesForClassification
 
 
   OCRFunction:

src/lambda/update_configuration/index.py

@@ -157,13 +157,6 @@ def handler(event: Dict[str, Any], context: Any) -> None:
                         if 'extraction' in resolved_default:
                             resolved_default['extraction']['model'] = properties['CustomExtractionModelARN']
                             logger.info(f"Updated extraction model to: {properties['CustomExtractionModelARN']}")
-                    
-                    # Add MaxPagesForClassification if provided
-                    if 'MaxPagesForClassification' in properties:
-                        if 'classification' not in resolved_default:
-                            resolved_default['classification'] = {}
-                        resolved_default['classification']['maxPagesForClassification'] = properties['MaxPagesForClassification']
-                        logger.info(f"Updated maxPagesForClassification to: {properties['MaxPagesForClassification']}")
 
                 update_configuration('Default', resolved_default)
 

template.yaml

@@ -114,17 +114,7 @@ Parameters:
     Description: >-
      Select the configuration preset for Pattern 1. Each configuration contains pre-tuned settings for specific document processing scenarios - see https://github.com/aws-samples/sample-genai-idp/blob/main/config_library/README.md. Note: This selected configuration will be replaced by the Custom Configuration Path if specified.
 
-  MaxPagesForClassification:
-    Type: String
-    Default: "ALL"
-    AllowedValues:
-      - "ALL"
-      - "1"
-      - "2"
-      - "3"
-      - "5"
-      - "10"
-    Description: "Number of pages to use for document classification. 'ALL' uses all pages, or specify a number (1-10) to limit classification to first N pages and apply result to all pages. Only applies to Pattern 2 and Pattern 3."
+
     
   # Pattern 2 Parameters
 
@@ -918,7 +908,6 @@ Resources:
             - "s3://${ConfigurationBucket}/config_library/pattern-2/${ConfigPath}/config.yaml"
             - ConfigPath: !FindInMap [Pattern2ConfigurationMap, !Ref Pattern2Configuration, ConfigPath]
         ConfigLibraryHash: "<CONFIG_LIBRARY_HASH_TOKEN>"
-        MaxPagesForClassification: !Ref MaxPagesForClassification
         EnableHITL: !Ref EnableHITL
         SageMakerA2IReviewPortalURL: !If
           - IsHITLEnabled
@@ -960,7 +949,6 @@ Resources:
             - "s3://${ConfigurationBucket}/config_library/pattern-3/${ConfigPath}/config.yaml"
             - ConfigPath: !FindInMap [Pattern3ConfigurationMap, !Ref Pattern3Configuration, ConfigPath]
         ConfigLibraryHash: "<CONFIG_LIBRARY_HASH_TOKEN>"
-        MaxPagesForClassification: !Ref MaxPagesForClassification
         EnableHITL: !Ref EnableHITL
         SageMakerA2IReviewPortalURL: !If
           - IsHITLEnabled