Trying to get new CI to work

Author: jakep-allenai

.github/workflows/main.yml

@@ -112,7 +112,7 @@ jobs:
     needs: [checks]
     env:
       BEAKER_TOKEN: ${{ secrets.BEAKER_TOKEN }}
-      BEAKER_IMAGE: chrisw/olmocr-gpu-ci
+      BEAKER_IMAGE: jakep/olmocr-gpu-ci
       BEAKER_BUDGET: ai2/oe-data
       BEAKER_WORKSPACE: ai2/olmocr
     steps:

scripts/beaker/gpu-ci-script.sh

@@ -9,7 +9,7 @@ git clone https://github.com/allenai/olmocr.git olmocr \
     .[gpu] \
     pytest \
     --find-links https://flashinfer.ai/whl/cu124/torch2.4/flashinfer/ \
-  && bash tests/gnarly_pdfs/test_gnarly_pdfs.sh
+  && bash scripts/run_integration_test.sh
 
 
 

scripts/run_integration_test.sh

@@ -0,0 +1,6 @@
+#/usr/bin/bash
+
+set -ex
+
+python -m olmocr.pipeline ./localworkspace --pdfs tests/gnarly_pdfs/ambiguous.pdf tests/gnarly_pdfs/edgar.pdf tests/gnarly_pdfs/dolma-page-1.pdf \
+  && pytest tests/test_integration.py

tests/gnarly_pdfs/test_gnarly_pdfs.py

@@ -0,0 +1,27 @@
+import glob
+import json
+import os
+import unittest
+
+import pytest
+
+
+@pytest.mark.nonci
+class TestPipelineIntegration(unittest.TestCase):
+    def setUp(self):
+        self.data = []
+
+        for file in glob.glob(os.path.join("localworkspace", "results", "*.jsonl")):
+            with open(file, "r") as jf:
+                for line in jf:
+                    if len(line.strip()) > 0:
+                        self.data.append(json.loads(line))
+
+    def test_edgar(self) -> None:
+        self.assertTrue(any("King of England" in line["text"] for line in self.data))
+
+    def test_ambig(self) -> None:
+        self.assertTrue(any("Apples and Bananas" in line["text"] for line in self.data))
+
+    def test_dolma(self) -> None:
+        self.assertTrue(any("We extensively document Dolma" in line["text"] for line in self.data))