From 2a5792e5ed9044f14577706f1e4ca942448a5bbf Mon Sep 17 00:00:00 2001 From: aman-17 Date: Fri, 26 Sep 2025 13:29:48 -0700 Subject: [PATCH] add if else for vllm local usage bug for api argument --- olmocr/pipeline.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/olmocr/pipeline.py b/olmocr/pipeline.py index 2b57a94..90ffce0 100644 --- a/olmocr/pipeline.py +++ b/olmocr/pipeline.py @@ -253,8 +253,13 @@ async def process_page(args, worker_id: int, pdf_orig_path: str, pdf_local_path: while attempt < MAX_RETRIES: lookup_attempt = min(attempt, len(TEMPERATURE_BY_ATTEMPT) - 1) - # Use the model name from args if provided, otherwise default to 'olmocr' - model_name = getattr(args, 'model', 'olmocr') if args.server else 'olmocr' + # For external servers (like DeepInfra), use the model name from args + # For local inference, always use 'olmocr' + if args.server and hasattr(args, 'model'): + model_name = args.model + else: + model_name = 'olmocr' + query = await build_page_query( pdf_local_path, page_num, @@ -274,8 +279,11 @@ async def process_page(args, worker_id: int, pdf_orig_path: str, pdf_local_path: logger.debug(f"Built page query for {pdf_orig_path}-{page_num}") try: - # Pass API key if provided - api_key = getattr(args, 'api_key', None) + # Pass API key only for external servers that need authentication + if args.server and hasattr(args, 'api_key'): + api_key = args.api_key + else: + api_key = None status_code, response_body = await apost(COMPLETION_URL, json_data=query, api_key=api_key) if status_code == 400: