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

Preparing pipeline for new format

Browse Source
This commit is contained in:
Jake Poznanski 2025-07-01 17:01:33 +00:00
parent c6c1fbd0eb
commit 9bf8e9e0fa
1 changed files with 10 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
olmocr/pipeline.py
View File

@ -132,7 +132,7 @@ async def build_page_query(local_pdf_path: str, page: int, target_longest_image_
image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8") image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
return { return {
"model": "Qwen/Qwen2-VL-7B-Instruct", "model": "olmocr",
"messages": [ "messages": [
{ {
"role": "user", "role": "user",
@ -260,8 +260,13 @@ async def process_page(args, worker_id: int, pdf_orig_path: str, pdf_local_path:
server_output_tokens=base_response_data["usage"].get("completion_tokens", 0), server_output_tokens=base_response_data["usage"].get("completion_tokens", 0),
) )
model_response_json = json.loads(base_response_data["choices"][0]["message"]["content"]) model_response_markdown = base_response_data["choices"][0]["message"]["content"]
page_response = PageResponse(**model_response_json)
# Somewhat temporary code, will need to refactor
from olmocr.train.dataloader import FrontMatterParser
parser = FrontMatterParser(front_matter_class=PageResponse)
front_matter, text = parser._extract_front_matter_and_text(model_response_markdown)
page_response = parser._parse_front_matter(front_matter, text)
if not page_response.is_rotation_valid and attempt < MAX_RETRIES - 1: if not page_response.is_rotation_valid and attempt < MAX_RETRIES - 1:
logger.info( logger.info(
@ -581,7 +586,7 @@ async def vllm_server_task(model_name_or_path, args, semaphore):
"--uvicorn-log-level", "--uvicorn-log-level",
"warning", "warning",
"--served-model-name", "--served-model-name",
"Qwen/Qwen2-VL-7B-Instruct", "olmocr",
] ]
cmd.extend(mem_fraction_arg) cmd.extend(mem_fraction_arg)
@ -1010,9 +1015,8 @@ async def main():
default="allenai/olmOCR-7B-0225-preview", default="allenai/olmOCR-7B-0225-preview",
) )
parser.add_argument("--model_max_context", type=int, default="8192", help="Maximum context length that the model was fine tuned under") parser.add_argument("--model_max_context", type=int, default="8192", help="Maximum context length that the model was fine tuned under")
parser.add_argument("--model_chat_template", type=str, default="qwen2-vl", help="Chat template to pass to vllm server")
parser.add_argument("--target_longest_image_dim", type=int, help="Dimension on longest side to use for rendering the pdf pages", default=1024) parser.add_argument("--target_longest_image_dim", type=int, help="Dimension on longest side to use for rendering the pdf pages", default=1024)
parser.add_argument("--target_anchor_text_len", type=int, help="Maximum amount of anchor text to use (characters)", default=6000) parser.add_argument("--target_anchor_text_len", type=int, help="Maximum amount of anchor text to use (characters)", default=3000)
# Beaker/job running stuff # Beaker/job running stuff
parser.add_argument("--beaker", action="store_true", help="Submit this job to beaker instead of running locally") parser.add_argument("--beaker", action="store_true", help="Submit this job to beaker instead of running locally")