Skip to content

Commit c47cd02

Browse files
davanstrienclaude
davanstrien
and
claude
authored
ocr: guard output-column collisions across the OCR recipes (+ --overwrite) (#66)
* ocr: guard output-column collisions across the OCR recipes (+ --overwrite) Generalises the pp-ocrv6 fix from #65 to the remaining output-writing recipes. Each script that adds an output column could previously build a *duplicate* column when that name already existed in the input dataset -- crashing on push after inference, or silently clobbering a ground-truth column (e.g. 'text' / 'markdown'). - Add a shared ensure_output_columns_free() startup guard (copied per-script; no shared lib in this repo) that fails fast if an output column already exists in the input, with a new --overwrite flag to opt into replacing it. Sink scripts (pp-doclayout) carry the equivalent guard inline; surya (in #65) guards both of its columns. - The 5 recipes that hardcoded the "markdown" output column (nanonets-ocr, abot-ocr, deepseek-ocr, deepseek-ocr-vllm) also gained a configurable --output-column. - Bonus, found via a context-length audit: nanonets-ocr.py had --max-tokens 15000 > --max-model-len 8192 (the output budget alone couldn't fit the window); raised the --max-model-len default to 32768 (model max is 128k). - deepseek-ocr.py: tidied a pre-existing bare `except:` -> `except Exception:` so ruff passes. Static-verified: ruff + AST + call-site wiring on all 25 files. The guard pattern itself is Jobs-proven via pp-ocrv6 in #65. A light Jobs smoke on one hardcoded convert + one uniform script before merge is cheap insurance for the arg wiring. Note: nanonets-ocr2.py is intentionally NOT included -- it carries separate in-progress vLLM-image edits; its collision guard will land with that work. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * ocr: guard nuextract3 reasoning column too (review nit #1) --enable-thinking also writes {output_column}_reasoning; include it in the collision guard (conditional on enable_thinking) so a pre-existing _reasoning column fails fast instead of crashing on push after inference. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent a7f5234 commit c47cd02

25 files changed

Lines changed: 881 additions & 21 deletions

‎ocr/abot-ocr.py‎

Lines changed: 44 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,28 @@ def check_cuda_availability():
9696
logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
9797

9898

99+
def ensure_output_columns_free(dataset, columns, overwrite=False):
100+
"""Fail fast if an output column would collide with an existing input column.
101+
102+
Adding a column that already exists silently overwrites it (e.g. a ground-truth
103+
`text`/`markdown` column) or crashes on push with a duplicate-column error only
104+
*after* inference has run. Catch it up front. With overwrite=True, drop the clashing
105+
column(s) here instead (logged) so the later add_column is clean.
106+
"""
107+
clash = [c for c in columns if c in dataset.column_names]
108+
if not clash:
109+
return dataset
110+
if overwrite:
111+
logger.warning(f"--overwrite: replacing existing column(s) {clash}")
112+
return dataset.remove_columns(clash)
113+
logger.error(
114+
f"Output column(s) {clash} already exist in the input dataset "
115+
f"(columns: {dataset.column_names})."
116+
)
117+
logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
118+
sys.exit(1)
119+
120+
99121
def post_process_text(text: str, threshold: int = 8000) -> str:
100122
"""Trim runaway repetition loops that VLM OCR can fall into on dense pages.
101123
@@ -258,6 +280,8 @@ def main(
258280
private: bool = False,
259281
shuffle: bool = False,
260282
seed: int = 42,
283+
output_column: str = "markdown",
284+
overwrite: bool = False,
261285
verbose: bool = False,
262286
):
263287
"""Process images from a HF dataset through the ABot-OCR model."""
@@ -286,6 +310,9 @@ def main(
286310
f"Column '{image_column}' not found. Available: {dataset.column_names}"
287311
)
288312

313+
# Fail fast if the output column would collide with an existing input column
314+
dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
315+
289316
# Shuffle if requested
290317
if shuffle:
291318
logger.info(f"Shuffling dataset with seed {seed}")
@@ -335,17 +362,17 @@ def main(
335362
logger.error(f"Error processing batch: {e}")
336363
all_markdown.extend(["[OCR FAILED]"] * len(batch_images))
337364

338-
# Add markdown column to dataset
339-
logger.info("Adding markdown column to dataset")
340-
dataset = dataset.add_column("markdown", all_markdown)
365+
# Add output column to dataset
366+
logger.info(f"Adding '{output_column}' column to dataset")
367+
dataset = dataset.add_column(output_column, all_markdown)
341368

342369
# Handle inference_info tracking
343370
logger.info("Updating inference_info...")
344371

345372
inference_entry = {
346373
"model_id": model,
347374
"model_name": "ABot-OCR",
348-
"column_name": "markdown",
375+
"column_name": output_column,
349376
"timestamp": datetime.now().isoformat(),
350377
"batch_size": batch_size,
351378
"max_tokens": max_tokens,
@@ -533,6 +560,17 @@ def update_inference_info(example):
533560
default=42,
534561
help="Random seed for shuffling (default: 42)",
535562
)
563+
parser.add_argument(
564+
"--output-column",
565+
default="markdown",
566+
help="Column name for the OCR output text (default: markdown)",
567+
)
568+
parser.add_argument(
569+
"--overwrite",
570+
action="store_true",
571+
help="Replace the output column if it already exists in the input dataset "
572+
"(default: error out to avoid clobbering an existing column).",
573+
)
536574
parser.add_argument(
537575
"--verbose",
538576
action="store_true",
@@ -556,5 +594,7 @@ def update_inference_info(example):
556594
private=args.private,
557595
shuffle=args.shuffle,
558596
seed=args.seed,
597+
output_column=args.output_column,
598+
overwrite=args.overwrite,
559599
verbose=args.verbose,
560600
)

‎ocr/deepseek-ocr-vllm.py‎

Lines changed: 44 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,28 @@ def check_cuda_availability():
7878
logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
7979

8080

81+
def ensure_output_columns_free(dataset, columns, overwrite=False):
82+
"""Fail fast if an output column would collide with an existing input column.
83+
84+
Adding a column that already exists silently overwrites it (e.g. a ground-truth
85+
`text`/`markdown` column) or crashes on push with a duplicate-column error only
86+
*after* inference has run. Catch it up front. With overwrite=True, drop the clashing
87+
column(s) here instead (logged) so the later add_column is clean.
88+
"""
89+
clash = [c for c in columns if c in dataset.column_names]
90+
if not clash:
91+
return dataset
92+
if overwrite:
93+
logger.warning(f"--overwrite: replacing existing column(s) {clash}")
94+
return dataset.remove_columns(clash)
95+
logger.error(
96+
f"Output column(s) {clash} already exist in the input dataset "
97+
f"(columns: {dataset.column_names})."
98+
)
99+
logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
100+
sys.exit(1)
101+
102+
81103
def to_pil(image: Union[Image.Image, Dict[str, Any], str]) -> Image.Image:
82104
"""Convert various image formats to PIL Image."""
83105
if isinstance(image, Image.Image):
@@ -211,6 +233,8 @@ def main(
211233
private: bool = False,
212234
shuffle: bool = False,
213235
seed: int = 42,
236+
output_column: str = "markdown",
237+
overwrite: bool = False,
214238
config: str = None,
215239
create_pr: bool = False,
216240
verbose: bool = False,
@@ -253,6 +277,9 @@ def main(
253277
f"Column '{image_column}' not found. Available: {dataset.column_names}"
254278
)
255279

280+
# Fail fast if the output column would collide with an existing input column
281+
dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
282+
256283
# Shuffle if requested
257284
if shuffle:
258285
logger.info(f"Shuffling dataset with seed {seed}")
@@ -326,17 +353,17 @@ def main(
326353
processing_duration = datetime.now() - start_time
327354
processing_time_str = f"{processing_duration.total_seconds() / 60:.1f} min"
328355

329-
# Add markdown column to dataset
330-
logger.info("Adding markdown column to dataset")
331-
dataset = dataset.add_column("markdown", all_markdown)
356+
# Add output column to dataset
357+
logger.info(f"Adding '{output_column}' column to dataset")
358+
dataset = dataset.add_column(output_column, all_markdown)
332359

333360
# Handle inference_info tracking
334361
logger.info("Updating inference_info...")
335362

336363
inference_entry = {
337364
"model_id": model,
338365
"model_name": "DeepSeek-OCR",
339-
"column_name": "markdown",
366+
"column_name": output_column,
340367
"timestamp": datetime.now().isoformat(),
341368
"prompt_mode": prompt_mode if prompt is None else "custom",
342369
"batch_size": batch_size,
@@ -569,6 +596,17 @@ def update_inference_info(example):
569596
default=42,
570597
help="Random seed for shuffling (default: 42)",
571598
)
599+
parser.add_argument(
600+
"--output-column",
601+
default="markdown",
602+
help="Column name for the OCR output text (default: markdown)",
603+
)
604+
parser.add_argument(
605+
"--overwrite",
606+
action="store_true",
607+
help="Replace the output column if it already exists in the input dataset "
608+
"(default: error out to avoid clobbering an existing column).",
609+
)
572610
parser.add_argument(
573611
"--verbose",
574612
action="store_true",
@@ -594,6 +632,8 @@ def update_inference_info(example):
594632
private=args.private,
595633
shuffle=args.shuffle,
596634
seed=args.seed,
635+
output_column=args.output_column,
636+
overwrite=args.overwrite,
597637
config=args.config,
598638
create_pr=args.create_pr,
599639
verbose=args.verbose,

‎ocr/deepseek-ocr.py‎

Lines changed: 45 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,28 @@ def check_cuda_availability():
7676
logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
7777

7878

79+
def ensure_output_columns_free(dataset, columns, overwrite=False):
80+
"""Fail fast if an output column would collide with an existing input column.
81+
82+
Adding a column that already exists silently overwrites it (e.g. a ground-truth
83+
`text`/`markdown` column) or crashes on push with a duplicate-column error only
84+
*after* inference has run. Catch it up front. With overwrite=True, drop the clashing
85+
column(s) here instead (logged) so the later add_column is clean.
86+
"""
87+
clash = [c for c in columns if c in dataset.column_names]
88+
if not clash:
89+
return dataset
90+
if overwrite:
91+
logger.warning(f"--overwrite: replacing existing column(s) {clash}")
92+
return dataset.remove_columns(clash)
93+
logger.error(
94+
f"Output column(s) {clash} already exist in the input dataset "
95+
f"(columns: {dataset.column_names})."
96+
)
97+
logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
98+
sys.exit(1)
99+
100+
79101
def create_dataset_card(
80102
source_dataset: str,
81103
model: str,
@@ -243,6 +265,8 @@ def main(
243265
private: bool = False,
244266
shuffle: bool = False,
245267
seed: int = 42,
268+
output_column: str = "markdown",
269+
overwrite: bool = False,
246270
):
247271
"""Process images from HF dataset through DeepSeek-OCR model."""
248272

@@ -294,6 +318,9 @@ def main(
294318
f"Column '{image_column}' not found. Available: {dataset.column_names}"
295319
)
296320

321+
# Fail fast if the output column would collide with an existing input column
322+
dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
323+
297324
# Shuffle if requested
298325
if shuffle:
299326
logger.info(f"Shuffling dataset with seed {seed}")
@@ -379,12 +406,12 @@ def main(
379406
try:
380407
shutil.rmtree(temp_dir)
381408
shutil.rmtree(temp_output_dir)
382-
except:
409+
except Exception:
383410
pass
384411

385-
# Add markdown column to dataset
386-
logger.info("Adding markdown column to dataset")
387-
dataset = dataset.add_column("markdown", all_markdown)
412+
# Add output column to dataset
413+
logger.info(f"Adding '{output_column}' column to dataset")
414+
dataset = dataset.add_column(output_column, all_markdown)
388415

389416
# Handle inference_info tracking
390417
logger.info("Updating inference_info...")
@@ -403,7 +430,7 @@ def main(
403430

404431
# Add new inference info
405432
new_info = {
406-
"column_name": "markdown",
433+
"column_name": output_column,
407434
"model_id": model,
408435
"processing_date": datetime.now().isoformat(),
409436
"resolution_mode": resolution_mode,
@@ -581,6 +608,17 @@ def main(
581608
default=42,
582609
help="Random seed for shuffling (default: 42)",
583610
)
611+
parser.add_argument(
612+
"--output-column",
613+
default="markdown",
614+
help="Column name for the OCR output text (default: markdown)",
615+
)
616+
parser.add_argument(
617+
"--overwrite",
618+
action="store_true",
619+
help="Replace the output column if it already exists in the input dataset "
620+
"(default: error out to avoid clobbering an existing column).",
621+
)
584622

585623
args = parser.parse_args()
586624

@@ -600,4 +638,6 @@ def main(
600638
private=args.private,
601639
shuffle=args.shuffle,
602640
seed=args.seed,
641+
output_column=args.output_column,
642+
overwrite=args.overwrite,
603643
)

‎ocr/deepseek-ocr2-vllm.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,28 @@ def check_cuda_availability():
8686
logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
8787

8888

89+
def ensure_output_columns_free(dataset, columns, overwrite=False):
90+
"""Fail fast if an output column would collide with an existing input column.
91+
92+
Adding a column that already exists silently overwrites it (e.g. a ground-truth
93+
`text`/`markdown` column) or crashes on push with a duplicate-column error only
94+
*after* inference has run. Catch it up front. With overwrite=True, drop the clashing
95+
column(s) here instead (logged) so the later add_column is clean.
96+
"""
97+
clash = [c for c in columns if c in dataset.column_names]
98+
if not clash:
99+
return dataset
100+
if overwrite:
101+
logger.warning(f"--overwrite: replacing existing column(s) {clash}")
102+
return dataset.remove_columns(clash)
103+
logger.error(
104+
f"Output column(s) {clash} already exist in the input dataset "
105+
f"(columns: {dataset.column_names})."
106+
)
107+
logger.error("Choose a different --output-column, or pass --overwrite to replace them.")
108+
sys.exit(1)
109+
110+
89111
def to_pil(image: Union[Image.Image, Dict[str, Any], str]) -> Image.Image:
90112
"""Convert various image formats to PIL Image."""
91113
if isinstance(image, Image.Image):
@@ -225,6 +247,7 @@ def main(
225247
shuffle: bool = False,
226248
seed: int = 42,
227249
output_column: str = "markdown",
250+
overwrite: bool = False,
228251
config: str = None,
229252
create_pr: bool = False,
230253
verbose: bool = False,
@@ -267,6 +290,9 @@ def main(
267290
f"Column '{image_column}' not found. Available: {dataset.column_names}"
268291
)
269292

293+
# Fail fast if the output column would collide with an existing input column
294+
dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)
295+
270296
# Shuffle if requested
271297
if shuffle:
272298
logger.info(f"Shuffling dataset with seed {seed}")
@@ -574,6 +600,12 @@ def update_inference_info(example):
574600
default="markdown",
575601
help="Column name for OCR output (default: markdown). Use a different name to add alongside existing OCR.",
576602
)
603+
parser.add_argument(
604+
"--overwrite",
605+
action="store_true",
606+
help="Replace the output column if it already exists in the input dataset "
607+
"(default: error out to avoid clobbering an existing column).",
608+
)
577609
parser.add_argument(
578610
"--shuffle",
579611
action="store_true",
@@ -620,6 +652,7 @@ def update_inference_info(example):
620652
shuffle=args.shuffle,
621653
seed=args.seed,
622654
output_column=args.output_column,
655+
overwrite=args.overwrite,
623656
config=args.config,
624657
create_pr=args.create_pr,
625658
verbose=args.verbose,

0 commit comments

Comments
 (0)