yujunjun/olmocr
1
0
Fork 0
mirror of https://github.com/allenai/olmocr.git synced 2025-10-13 17:22:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixing for conv format

Browse Source
This commit is contained in:
Jake Poznanski 2025-08-21 18:47:53 +00:00
parent 0f8d515d8c
commit 33d889c748
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
olmocr/train/grpo_train.py
View File

@ -187,7 +187,7 @@ def load_tests_cached(jsonl_file: str):
return load_tests(jsonl_file) return load_tests(jsonl_file)
def olmocr_bench_reward(prompts, completions: list[str], completion_ids: list[list[int]], pdf_path: list[str], jsonl_file: list[str], test_ids: list[list[str]], **kwargs): def olmocr_bench_reward(prompts, completions: list[str] | list[list[dict]], completion_ids: list[list[int]], pdf_path: list[str], jsonl_file: list[str], test_ids: list[list[str]], **kwargs):
""" """
Reward function that runs unit tests on completions and returns average pass rate. Reward function that runs unit tests on completions and returns average pass rate.
@ -219,7 +219,7 @@ def olmocr_bench_reward(prompts, completions: list[str], completion_ids: list[li
logger.info(f"Completion {i}: PDF: {comp_pdf_path}, JSONL: {comp_jsonl_file}, Test IDs: {comp_test_ids}") logger.info(f"Completion {i}: PDF: {comp_pdf_path}, JSONL: {comp_jsonl_file}, Test IDs: {comp_test_ids}")
if completion is None or not isinstance(completion, str): if completion is None or not isinstance(completion, str) or not isinstance(completion, list):
logger.warning(f"Invalid completion at index {i}: {type(completion)}") logger.warning(f"Invalid completion at index {i}: {type(completion)}")
logger.warning(f"completion: {completion}") logger.warning(f"completion: {completion}")
rewards.append(None) rewards.append(None)
@ -230,6 +230,9 @@ def olmocr_bench_reward(prompts, completions: list[str], completion_ids: list[li
rewards.append(None) rewards.append(None)
continue continue
if isinstance(completion, list):
completion = completion[0]["content"]
try: try:
# Load all tests from the JSONL file (cached) # Load all tests from the JSONL file (cached)
all_tests = load_tests_cached(comp_jsonl_file) all_tests = load_tests_cached(comp_jsonl_file)