mirror of
https://github.com/allenai/olmocr.git
synced 2025-10-12 08:43:32 +00:00
338 lines
13 KiB
Python
338 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test suite for GRPO training dataloader.
|
|
Tests the OlmOCRDataset class and its functionality with Olmocr-bench format.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import json
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
import shutil
|
|
|
|
# Add parent directory to path
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from olmocr.train.grpo_train import OlmOCRDataset, collate_fn
|
|
|
|
|
|
class TestGRPODataloader(unittest.TestCase):
|
|
"""Test cases for the GRPO dataloader."""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
"""Create a temporary bench_data folder with test data."""
|
|
cls.temp_dir = tempfile.mkdtemp()
|
|
cls.bench_data_folder = cls.temp_dir
|
|
cls.pdfs_folder = os.path.join(cls.bench_data_folder, "pdfs")
|
|
|
|
# Create folder structure
|
|
os.makedirs(os.path.join(cls.pdfs_folder, "test_pdfs"), exist_ok=True)
|
|
|
|
# Create dummy PDF files
|
|
cls.pdf_files = []
|
|
for i in range(3):
|
|
pdf_path = os.path.join(cls.pdfs_folder, "test_pdfs", f"test_{i}.pdf")
|
|
# Create a minimal valid PDF
|
|
with open(pdf_path, "wb") as f:
|
|
f.write(b"%PDF-1.4\n%%EOF")
|
|
cls.pdf_files.append(pdf_path)
|
|
|
|
# Create test JSONL files
|
|
cls.jsonl_file1 = os.path.join(cls.bench_data_folder, "test1.jsonl")
|
|
cls.jsonl_file2 = os.path.join(cls.bench_data_folder, "test2.jsonl")
|
|
|
|
# Write test data to JSONL files
|
|
test_data1 = [
|
|
{"pdf": "test_pdfs/test_0.pdf", "page": 0, "id": "test_0_001", "type": "math", "math": "x + y = z"},
|
|
{"pdf": "test_pdfs/test_0.pdf", "page": 0, "id": "test_0_002", "type": "text", "text": "Sample text"},
|
|
{"pdf": "test_pdfs/test_1.pdf", "page": 0, "id": "test_1_001", "type": "math", "math": "a^2 + b^2 = c^2"},
|
|
{"pdf": "test_pdfs/test_1.pdf", "page": 1, "id": "test_1_002", "type": "text", "text": "Another sample"},
|
|
]
|
|
|
|
test_data2 = [
|
|
{"pdf": "test_pdfs/test_2.pdf", "page": 0, "id": "test_2_001", "type": "table", "table": "col1,col2"},
|
|
{"pdf": "test_pdfs/test_0.pdf", "page": 0, "id": "test_0_003", "type": "text", "text": "More text"},
|
|
{"pdf": "test_pdfs/test_2.pdf", "page": 0, "id": "test_2_002", "type": "math", "math": "\\int_0^1 x dx"},
|
|
]
|
|
|
|
with open(cls.jsonl_file1, "w") as f:
|
|
for entry in test_data1:
|
|
f.write(json.dumps(entry) + "\n")
|
|
|
|
with open(cls.jsonl_file2, "w") as f:
|
|
for entry in test_data2:
|
|
f.write(json.dumps(entry) + "\n")
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
"""Clean up temporary files."""
|
|
shutil.rmtree(cls.temp_dir)
|
|
|
|
def test_dataset_initialization(self):
|
|
"""Test that dataset initializes correctly."""
|
|
dataset = OlmOCRDataset(
|
|
bench_data_folder=self.bench_data_folder,
|
|
processor=None,
|
|
max_samples=None,
|
|
target_longest_image_dim=1024,
|
|
)
|
|
|
|
self.assertIsNotNone(dataset)
|
|
self.assertEqual(dataset.bench_data_folder, self.bench_data_folder)
|
|
self.assertEqual(dataset.pdf_folder, self.pdfs_folder)
|
|
self.assertTrue(len(dataset) > 0)
|
|
|
|
def test_unique_pdf_loading(self):
|
|
"""Test that unique PDFs are loaded correctly."""
|
|
dataset = OlmOCRDataset(
|
|
bench_data_folder=self.bench_data_folder,
|
|
processor=None,
|
|
max_samples=None,
|
|
target_longest_image_dim=1024,
|
|
)
|
|
|
|
# Should have 4 unique PDF+page combinations:
|
|
# test_0.pdf page 0, test_1.pdf page 0, test_1.pdf page 1, test_2.pdf page 0
|
|
self.assertEqual(len(dataset), 4)
|
|
|
|
# Check that samples have correct structure
|
|
for sample in dataset.samples:
|
|
self.assertIn("pdf_path", sample)
|
|
self.assertIn("pdf_name", sample)
|
|
self.assertIn("page", sample)
|
|
self.assertIn("jsonl_file", sample)
|
|
self.assertIn("test_ids", sample)
|
|
self.assertIn("entries", sample)
|
|
|
|
def test_test_id_aggregation(self):
|
|
"""Test that test IDs are correctly aggregated per PDF+page."""
|
|
dataset = OlmOCRDataset(
|
|
bench_data_folder=self.bench_data_folder,
|
|
processor=None,
|
|
max_samples=None,
|
|
target_longest_image_dim=1024,
|
|
)
|
|
|
|
# Find the sample for test_0.pdf page 0
|
|
test_0_sample = None
|
|
for sample in dataset.samples:
|
|
if "test_0.pdf" in sample["pdf_name"] and sample["page"] == 0:
|
|
test_0_sample = sample
|
|
break
|
|
|
|
self.assertIsNotNone(test_0_sample)
|
|
# Should have 3 test IDs for test_0.pdf page 0
|
|
self.assertEqual(len(test_0_sample["test_ids"]), 3)
|
|
self.assertIn("test_0_001", test_0_sample["test_ids"])
|
|
self.assertIn("test_0_002", test_0_sample["test_ids"])
|
|
self.assertIn("test_0_003", test_0_sample["test_ids"])
|
|
|
|
def test_max_samples_limit(self):
|
|
"""Test that max_samples correctly limits the dataset size."""
|
|
dataset = OlmOCRDataset(
|
|
bench_data_folder=self.bench_data_folder,
|
|
processor=None,
|
|
max_samples=2,
|
|
target_longest_image_dim=1024,
|
|
)
|
|
|
|
self.assertEqual(len(dataset), 2)
|
|
|
|
@patch('olmocr.train.grpo_train.render_pdf_to_base64png')
|
|
@patch('olmocr.train.grpo_train.build_no_anchoring_v4_yaml_prompt')
|
|
def test_getitem_format(self, mock_prompt, mock_render):
|
|
"""Test that __getitem__ returns the correct format."""
|
|
# Mock the rendering and prompt functions
|
|
mock_render.return_value = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" # 1x1 white pixel PNG
|
|
mock_prompt.return_value = "Test prompt"
|
|
|
|
dataset = OlmOCRDataset(
|
|
bench_data_folder=self.bench_data_folder,
|
|
processor=None,
|
|
max_samples=1,
|
|
target_longest_image_dim=1024,
|
|
)
|
|
|
|
item = dataset[0]
|
|
|
|
self.assertIsNotNone(item)
|
|
self.assertIn("prompt", item)
|
|
self.assertIn("pdf_path", item)
|
|
self.assertIn("jsonl_file", item)
|
|
self.assertIn("test_ids", item)
|
|
self.assertIn("image", item)
|
|
|
|
# Check prompt structure
|
|
self.assertIsInstance(item["prompt"], list)
|
|
self.assertEqual(len(item["prompt"]), 1)
|
|
self.assertEqual(item["prompt"][0]["role"], "user")
|
|
self.assertIsInstance(item["prompt"][0]["content"], list)
|
|
self.assertEqual(len(item["prompt"][0]["content"]), 2)
|
|
|
|
# Check other fields
|
|
self.assertIsInstance(item["pdf_path"], str)
|
|
self.assertIsInstance(item["jsonl_file"], str)
|
|
self.assertIsInstance(item["test_ids"], list)
|
|
self.assertTrue(len(item["test_ids"]) > 0)
|
|
|
|
@patch('olmocr.train.grpo_train.render_pdf_to_base64png')
|
|
@patch('olmocr.train.grpo_train.build_no_anchoring_v4_yaml_prompt')
|
|
def test_collate_function(self, mock_prompt, mock_render):
|
|
"""Test that the collate function works correctly."""
|
|
# Mock the rendering and prompt functions
|
|
mock_render.return_value = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
|
|
mock_prompt.return_value = "Test prompt"
|
|
|
|
dataset = OlmOCRDataset(
|
|
bench_data_folder=self.bench_data_folder,
|
|
processor=None,
|
|
max_samples=2,
|
|
target_longest_image_dim=1024,
|
|
)
|
|
|
|
# Create a batch
|
|
batch = [dataset[0], dataset[1]]
|
|
collated = collate_fn(batch)
|
|
|
|
self.assertIsNotNone(collated)
|
|
self.assertIn("prompts", collated)
|
|
self.assertIn("images", collated)
|
|
self.assertIn("pdf_paths", collated)
|
|
self.assertIn("jsonl_files", collated)
|
|
self.assertIn("test_ids", collated)
|
|
|
|
# Check batch size consistency
|
|
self.assertEqual(len(collated["prompts"]), 2)
|
|
self.assertEqual(len(collated["images"]), 2)
|
|
self.assertEqual(len(collated["pdf_paths"]), 2)
|
|
self.assertEqual(len(collated["jsonl_files"]), 2)
|
|
self.assertEqual(len(collated["test_ids"]), 2)
|
|
|
|
def test_collate_with_none_values(self):
|
|
"""Test that collate function handles None values correctly."""
|
|
batch = [None, {"prompt": [], "image": None, "pdf_path": "test.pdf",
|
|
"jsonl_file": "test.jsonl", "test_ids": ["id1"]}, None]
|
|
|
|
collated = collate_fn(batch)
|
|
|
|
self.assertIsNotNone(collated)
|
|
self.assertEqual(len(collated["prompts"]), 1)
|
|
|
|
def test_empty_jsonl_handling(self):
|
|
"""Test handling of empty JSONL files."""
|
|
# Create an empty JSONL file
|
|
empty_jsonl = os.path.join(self.bench_data_folder, "empty.jsonl")
|
|
open(empty_jsonl, "w").close()
|
|
|
|
# Should still work with other non-empty files
|
|
dataset = OlmOCRDataset(
|
|
bench_data_folder=self.bench_data_folder,
|
|
processor=None,
|
|
max_samples=None,
|
|
target_longest_image_dim=1024,
|
|
)
|
|
|
|
self.assertTrue(len(dataset) > 0)
|
|
|
|
# Clean up
|
|
os.remove(empty_jsonl)
|
|
|
|
def test_malformed_jsonl_handling(self):
|
|
"""Test handling of malformed JSONL entries."""
|
|
# Create a JSONL with some malformed entries
|
|
malformed_jsonl = os.path.join(self.bench_data_folder, "malformed.jsonl")
|
|
with open(malformed_jsonl, "w") as f:
|
|
f.write('{"pdf": "test.pdf", "id": "valid_1"}\n')
|
|
f.write('not valid json\n')
|
|
f.write('{"pdf": "test2.pdf", "id": "valid_2"}\n')
|
|
|
|
# Should skip malformed entries but process valid ones
|
|
dataset = OlmOCRDataset(
|
|
bench_data_folder=self.bench_data_folder,
|
|
processor=None,
|
|
max_samples=None,
|
|
target_longest_image_dim=1024,
|
|
)
|
|
|
|
# Should still have entries from valid files
|
|
self.assertTrue(len(dataset) > 0)
|
|
|
|
# Clean up
|
|
os.remove(malformed_jsonl)
|
|
|
|
def test_missing_pdf_folder(self):
|
|
"""Test error handling when pdfs folder is missing."""
|
|
temp_bad_folder = tempfile.mkdtemp()
|
|
|
|
with self.assertRaises(ValueError) as context:
|
|
dataset = OlmOCRDataset(
|
|
bench_data_folder=temp_bad_folder,
|
|
processor=None,
|
|
max_samples=None,
|
|
target_longest_image_dim=1024,
|
|
)
|
|
|
|
self.assertIn("PDFs folder not found", str(context.exception))
|
|
|
|
# Clean up
|
|
shutil.rmtree(temp_bad_folder)
|
|
|
|
def test_no_jsonl_files(self):
|
|
"""Test error handling when no JSONL files are present."""
|
|
temp_folder = tempfile.mkdtemp()
|
|
os.makedirs(os.path.join(temp_folder, "pdfs"))
|
|
|
|
with self.assertRaises(ValueError) as context:
|
|
dataset = OlmOCRDataset(
|
|
bench_data_folder=temp_folder,
|
|
processor=None,
|
|
max_samples=None,
|
|
target_longest_image_dim=1024,
|
|
)
|
|
|
|
self.assertIn("No JSONL files found", str(context.exception))
|
|
|
|
# Clean up
|
|
shutil.rmtree(temp_folder)
|
|
|
|
|
|
class TestIntegrationWithRealData(unittest.TestCase):
|
|
"""Integration tests with real bench data if available."""
|
|
|
|
@unittest.skipUnless(
|
|
os.path.exists("/home/ubuntu/olmocr/olmOCR-bench/bench_data"),
|
|
"Real bench data not available"
|
|
)
|
|
def test_with_real_bench_data(self):
|
|
"""Test with real bench data if available."""
|
|
bench_data_folder = "/home/ubuntu/olmocr/olmOCR-bench/bench_data"
|
|
|
|
dataset = OlmOCRDataset(
|
|
bench_data_folder=bench_data_folder,
|
|
processor=None,
|
|
max_samples=5,
|
|
target_longest_image_dim=1024,
|
|
)
|
|
|
|
self.assertEqual(len(dataset), 5)
|
|
|
|
# Test that we can iterate through the dataset
|
|
for i in range(len(dataset)):
|
|
item = dataset[i]
|
|
if item is not None: # Some PDFs might fail to render
|
|
self.assertIn("prompt", item)
|
|
self.assertIn("pdf_path", item)
|
|
self.assertIn("jsonl_file", item)
|
|
self.assertIn("test_ids", item)
|
|
|
|
# Verify paths exist
|
|
self.assertTrue(os.path.exists(item["pdf_path"]))
|
|
self.assertTrue(os.path.exists(item["jsonl_file"]))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main(verbosity=2) |