Cumulative rotation support

Author: jakep-allenai

olmocr/pipeline.py

@@ -217,7 +217,7 @@ async def process_page(args, worker_id: int, pdf_orig_path: str, pdf_local_path:
     MODEL_MAX_CONTEXT = 16384
     TEMPERATURE_BY_ATTEMPT = [0.1, 0.1, 0.2, 0.3, 0.5, 0.8, 0.9, 1.0]
     exponential_backoffs = 0
-    local_image_rotation = 0
+    cumulative_rotation = 0  # Track cumulative rotation instead of local
     attempt = 0
     await tracker.track_work(worker_id, f"{pdf_orig_path}-{page_num}", "started")
 
@@ -227,7 +227,7 @@ async def process_page(args, worker_id: int, pdf_orig_path: str, pdf_local_path:
             pdf_local_path,
             page_num,
             args.target_longest_image_dim,
-            image_rotation=local_image_rotation,
+            image_rotation=cumulative_rotation,
         )
         # Change temperature as number of attempts increases to overcome repetition issues at expense of quality
         query["temperature"] = TEMPERATURE_BY_ATTEMPT[lookup_attempt]
@@ -273,7 +273,9 @@ async def process_page(args, worker_id: int, pdf_orig_path: str, pdf_local_path:
                 logger.info(
                     f"Got invalid_page rotation for {pdf_orig_path}-{page_num} attempt {attempt}, retrying with {page_response.rotation_correction} rotation"
                 )
-                local_image_rotation = page_response.rotation_correction
+                # Add the rotation correction to the cumulative rotation
+                cumulative_rotation = (cumulative_rotation + page_response.rotation_correction) % 360
+                logger.info(f"Cumulative rotation is now {cumulative_rotation} degrees")
                 raise ValueError(f"invalid_page rotation for {pdf_orig_path}-{page_num}")
 
             metrics.add_metrics(**{"completed_pages": 1, f"finished_on_attempt_{attempt}": 1})

tests/test_pipeline.py

@@ -296,3 +296,211 @@ async def mock_build_page_query(local_pdf_path, page, target_longest_image_dim,
         # Verify tracker was called correctly
         mock_tracker.track_work.assert_any_call(0, "test-edgar-rotated90.pdf-1", "started")
         mock_tracker.track_work.assert_any_call(0, "test-edgar-rotated90.pdf-1", "finished")
+
+    @pytest.mark.asyncio
+    async def test_process_page_with_cumulative_rotation(self):
+        """Test that process_page correctly accumulates rotations across multiple attempts."""
+
+        # Path to the test PDF (can use any test PDF)
+        test_pdf_path = "tests/gnarly_pdfs/edgar-rotated90.pdf"
+
+        # Mock arguments
+        args = MockArgs()
+
+        # Counter to track number of API calls
+        call_count = 0
+
+        async def mock_apost(url, json_data):
+            nonlocal call_count
+            call_count += 1
+
+            # First call - model detects rotation is needed (90 degrees)
+            if call_count == 1:
+                response_content = """---
+primary_language: en
+is_rotation_valid: false
+rotation_correction: 90
+is_table: false
+is_diagram: false
+---
+
+This document appears to be rotated and needs correction."""
+
+            # Second call - model still detects rotation is needed (another 90 degrees)
+            elif call_count == 2:
+                response_content = """---
+primary_language: en
+is_rotation_valid: false
+rotation_correction: 90
+is_table: false
+is_diagram: false
+---
+
+Document still needs rotation."""
+
+            # Third call - after 180 total degrees of rotation, model says it's correct
+            elif call_count == 3:
+                response_content = """---
+primary_language: en
+is_rotation_valid: true
+rotation_correction: 0
+is_table: false
+is_diagram: false
+---
+
+UNITED STATES
+SECURITIES AND EXCHANGE COMMISSION
+Washington, D.C. 20549
+
+Document is now correctly oriented after 180 degree rotation."""
+
+            else:
+                raise ValueError(f"Unexpected call count: {call_count}")
+
+            # Mock response structure
+            response_body = {
+                "choices": [{"message": {"content": response_content}, "finish_reason": "stop"}],
+                "usage": {"prompt_tokens": 1000, "completion_tokens": 100, "total_tokens": 1100},
+            }
+
+            return 200, json.dumps(response_body).encode()
+
+        # Mock the worker tracker
+        mock_tracker = AsyncMock()
+
+        # Ensure the test PDF exists
+        assert os.path.exists(test_pdf_path), f"Test PDF not found at {test_pdf_path}"
+
+        # Track calls to build_page_query
+        build_page_query_calls = []
+        original_build_page_query = build_page_query
+
+        async def mock_build_page_query(local_pdf_path, page, target_longest_image_dim, image_rotation=0):
+            build_page_query_calls.append(image_rotation)
+            return await original_build_page_query(local_pdf_path, page, target_longest_image_dim, image_rotation)
+
+        with patch("olmocr.pipeline.apost", side_effect=mock_apost):
+            with patch("olmocr.pipeline.tracker", mock_tracker):
+                with patch("olmocr.pipeline.build_page_query", side_effect=mock_build_page_query):
+                    result = await process_page(args=args, worker_id=0, pdf_orig_path="test-cumulative-rotation.pdf", pdf_local_path=test_pdf_path, page_num=1)
+
+        # Verify the result
+        assert isinstance(result, PageResult)
+        assert result.page_num == 1
+        assert result.is_fallback == False
+        assert result.response.is_rotation_valid == True
+        assert result.response.rotation_correction == 0
+        assert result.response.natural_text is not None
+        assert "180 degree rotation" in result.response.natural_text
+
+        # Verify that exactly 3 API calls were made
+        assert call_count == 3
+
+        # Verify build_page_query was called with correct cumulative rotations
+        assert len(build_page_query_calls) == 3
+        assert build_page_query_calls[0] == 0    # First call with no rotation
+        assert build_page_query_calls[1] == 90   # Second call with 90 degree rotation
+        assert build_page_query_calls[2] == 180  # Third call with cumulative 180 degree rotation
+
+        # Verify tracker was called correctly
+        mock_tracker.track_work.assert_any_call(0, "test-cumulative-rotation.pdf-1", "started")
+        mock_tracker.track_work.assert_any_call(0, "test-cumulative-rotation.pdf-1", "finished")
+
+    @pytest.mark.asyncio
+    async def test_process_page_rotation_wraps_around(self):
+        """Test that cumulative rotation correctly wraps around at 360 degrees."""
+
+        # Path to the test PDF
+        test_pdf_path = "tests/gnarly_pdfs/edgar-rotated90.pdf"
+
+        # Mock arguments
+        args = MockArgs()
+
+        # Counter to track number of API calls
+        call_count = 0
+
+        async def mock_apost(url, json_data):
+            nonlocal call_count
+            call_count += 1
+
+            # First call - model detects rotation is needed (270 degrees)
+            if call_count == 1:
+                response_content = """---
+primary_language: en
+is_rotation_valid: false
+rotation_correction: 270
+is_table: false
+is_diagram: false
+---
+
+Document needs 270 degree rotation."""
+
+            # Second call - model detects more rotation is needed (180 degrees)
+            # Total would be 450, but should wrap to 90
+            elif call_count == 2:
+                response_content = """---
+primary_language: en
+is_rotation_valid: false
+rotation_correction: 180
+is_table: false
+is_diagram: false
+---
+
+Document needs additional rotation."""
+
+            # Third call - after wrapped rotation (90 degrees), model says it's correct
+            elif call_count == 3:
+                response_content = """---
+primary_language: en
+is_rotation_valid: true
+rotation_correction: 0
+is_table: false
+is_diagram: false
+---
+
+Document correctly oriented at 90 degrees total rotation."""
+
+            else:
+                raise ValueError(f"Unexpected call count: {call_count}")
+
+            # Mock response structure
+            response_body = {
+                "choices": [{"message": {"content": response_content}, "finish_reason": "stop"}],
+                "usage": {"prompt_tokens": 1000, "completion_tokens": 100, "total_tokens": 1100},
+            }
+
+            return 200, json.dumps(response_body).encode()
+
+        # Mock the worker tracker
+        mock_tracker = AsyncMock()
+
+        # Ensure the test PDF exists
+        assert os.path.exists(test_pdf_path), f"Test PDF not found at {test_pdf_path}"
+
+        # Track calls to build_page_query
+        build_page_query_calls = []
+        original_build_page_query = build_page_query
+
+        async def mock_build_page_query(local_pdf_path, page, target_longest_image_dim, image_rotation=0):
+            build_page_query_calls.append(image_rotation)
+            return await original_build_page_query(local_pdf_path, page, target_longest_image_dim, image_rotation)
+
+        with patch("olmocr.pipeline.apost", side_effect=mock_apost):
+            with patch("olmocr.pipeline.tracker", mock_tracker):
+                with patch("olmocr.pipeline.build_page_query", side_effect=mock_build_page_query):
+                    result = await process_page(args=args, worker_id=0, pdf_orig_path="test-rotation-wrap.pdf", pdf_local_path=test_pdf_path, page_num=1)
+
+        # Verify the result
+        assert isinstance(result, PageResult)
+        assert result.page_num == 1
+        assert result.is_fallback == False
+        assert result.response.is_rotation_valid == True
+
+        # Verify that exactly 3 API calls were made
+        assert call_count == 3
+
+        # Verify build_page_query was called with correct cumulative rotations
+        assert len(build_page_query_calls) == 3
+        assert build_page_query_calls[0] == 0    # First call with no rotation
+        assert build_page_query_calls[1] == 270  # Second call with 270 degree rotation
+        assert build_page_query_calls[2] == 90   # Third call with wrapped rotation (270 + 180 = 450 % 360 = 90)