wehub-resource-sync
917eedffcf
Main / Python 3.11 - Docs (push) Has been cancelled
Main / Python 3.11 - Build (push) Has been cancelled
Main / Python 3.11 - Lint (push) Has been cancelled
Main / Python 3.11 - Style (push) Has been cancelled
Main / Python 3.11 - Test (push) Has been cancelled
Main / GPU CI (push) Has been cancelled
Main / Release (push) Has been cancelled
Main / Build and Push Docker Images (push) Has been cancelled
396 lines
15 KiB
Python
396 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Repackage locally processed OLMoCR-mix style data back into parquet metadata and PDF tarballs.
|
|
|
|
Given a directory that mirrors the layout produced by prepare_olmocrmix.py (folders of markdown/PDF
|
|
pairs), this script rebuilds a HuggingFace-style payload by:
|
|
* walking the processed directory to recover document ids, metadata, and natural text
|
|
* emitting a parquet file with dedicated columns for PageResponse fields plus document helpers
|
|
* chunking PDFs into .tar.gz archives that stay under a user-configurable size (default 1 GiB)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import tarfile
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Dict, Iterator, List, Optional
|
|
|
|
import pandas as pd
|
|
from tqdm import tqdm
|
|
|
|
from olmocr.prompts import PageResponse
|
|
from olmocr.train.front_matter import FrontMatterParser
|
|
|
|
DEFAULT_MAX_TAR_BYTES = 1_073_741_824 # 1 GiB
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class DocumentRecord:
|
|
doc_id: str
|
|
markdown_path: Path
|
|
pdf_path: Path
|
|
pdf_size: int
|
|
primary_language: Optional[str]
|
|
is_rotation_valid: Optional[bool]
|
|
rotation_correction: Optional[int]
|
|
is_table: Optional[bool]
|
|
is_diagram: Optional[bool]
|
|
natural_text: Optional[str]
|
|
page_number: Optional[int]
|
|
url: Optional[str]
|
|
extras_json: Optional[str]
|
|
chunk_name: Optional[str] = None
|
|
pdf_relpath: Optional[str] = None
|
|
|
|
|
|
def infer_doc_id(md_path: Path, processed_root: Path) -> str:
|
|
"""Reconstruct the doc_id used in parquet/index space."""
|
|
rel = md_path.relative_to(processed_root)
|
|
|
|
# Simply preserve the directory structure as the doc_id
|
|
# Convert path to doc_id by removing extension
|
|
return str(rel.with_suffix(""))
|
|
|
|
|
|
def infer_pdf_path(md_path: Path, doc_id: str, pdf_root: Optional[Path]) -> Path:
|
|
"""Locate the PDF file corresponding to the markdown doc."""
|
|
pdf_candidate = md_path.with_suffix(".pdf")
|
|
if pdf_candidate.exists():
|
|
return pdf_candidate.resolve()
|
|
|
|
if pdf_root is not None:
|
|
alt_path = pdf_root / f"{doc_id}.pdf"
|
|
if alt_path.exists():
|
|
return alt_path.resolve()
|
|
|
|
raise FileNotFoundError(f"No PDF found for {md_path}")
|
|
|
|
|
|
def normalize_response_payload(front_matter: Dict[str, object], body_text: str) -> Dict[str, object]:
|
|
"""Merge parsed fields with the natural text payload."""
|
|
payload = dict(front_matter)
|
|
text = body_text if body_text and body_text.strip() else None
|
|
|
|
# Handle primary_language field - convert booleans to None
|
|
if "primary_language" in payload:
|
|
val = payload["primary_language"]
|
|
if isinstance(val, bool):
|
|
# Convert boolean to None (no language detected)
|
|
print(f"[DEBUG] Converting boolean primary_language value '{val}' to None")
|
|
payload["primary_language"] = None
|
|
elif not isinstance(val, (str, type(None))):
|
|
# Convert other types to string or None
|
|
print(f"[DEBUG] Converting unexpected primary_language type {type(val)} value '{val}' to string/None")
|
|
payload["primary_language"] = str(val) if val else None
|
|
else:
|
|
payload["primary_language"] = None
|
|
|
|
payload.setdefault("is_rotation_valid", True)
|
|
payload.setdefault("rotation_correction", 0)
|
|
payload.setdefault("is_table", False)
|
|
payload.setdefault("is_diagram", False)
|
|
payload["natural_text"] = text
|
|
return payload
|
|
|
|
|
|
def load_url_mappings(processed_dir: Path) -> Dict[str, str]:
|
|
"""Load URL mappings from urls.jsonl if it exists."""
|
|
urls_file = processed_dir / "urls.jsonl"
|
|
url_map = {}
|
|
|
|
if urls_file.exists():
|
|
print(f"Loading URL mappings from {urls_file}")
|
|
with open(urls_file, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
entry = json.loads(line.strip())
|
|
url_map[entry["id"]] = entry["url"]
|
|
print(f"Loaded {len(url_map)} URL mappings")
|
|
|
|
return url_map
|
|
|
|
|
|
def guess_url(front_matter: Dict[str, object], doc_id: str, source_url_template: Optional[str]) -> Optional[str]:
|
|
# TODO, we will have to add some better support for this
|
|
return None
|
|
|
|
|
|
def parse_page_number(doc_id: str, front_matter: Dict[str, object]) -> Optional[int]:
|
|
"""Extract page number from front matter or doc_id suffix."""
|
|
if "page_number" in front_matter:
|
|
value = front_matter["page_number"]
|
|
try:
|
|
return int(value)
|
|
except (TypeError, ValueError):
|
|
pass
|
|
|
|
if "-" in doc_id:
|
|
suffix = doc_id.rsplit("-", 1)[-1]
|
|
try:
|
|
return int(suffix)
|
|
except ValueError:
|
|
return None
|
|
return None
|
|
|
|
|
|
def collect_documents(
|
|
processed_dir: Path,
|
|
pdf_root: Optional[Path],
|
|
url_template: Optional[str],
|
|
strict: bool,
|
|
) -> List[DocumentRecord]:
|
|
"""Scan processed markdown/pdf pairs into DocumentRecord objects."""
|
|
records: List[DocumentRecord] = []
|
|
md_files = sorted(processed_dir.rglob("*.md"))
|
|
canonical_keys = {
|
|
"primary_language",
|
|
"is_rotation_valid",
|
|
"rotation_correction",
|
|
"is_table",
|
|
"is_diagram",
|
|
"natural_text",
|
|
}
|
|
|
|
# Load URL mappings from urls.jsonl if it exists
|
|
url_map = load_url_mappings(processed_dir)
|
|
|
|
parser = FrontMatterParser(front_matter_class=PageResponse)
|
|
|
|
for md_path in tqdm(md_files, desc="Scanning markdown files"):
|
|
try:
|
|
doc_id = infer_doc_id(md_path, processed_dir)
|
|
pdf_path = infer_pdf_path(md_path, doc_id, pdf_root)
|
|
markdown_text = md_path.read_text(encoding="utf-8")
|
|
front_matter, body_text = parser._extract_front_matter_and_text(markdown_text)
|
|
response_payload = normalize_response_payload(front_matter, body_text)
|
|
pdf_size = pdf_path.stat().st_size
|
|
page_number = parse_page_number(doc_id, front_matter)
|
|
|
|
# Try to get URL from the loaded url_map
|
|
# Handle both formats: "0001/234567" and "0001234567"
|
|
url = url_map.get(doc_id)
|
|
if not url and "/" in doc_id:
|
|
# Try combining the parts (e.g., "0001/234567" -> "0001234567")
|
|
combined_id = doc_id.replace("/", "")
|
|
url = url_map.get(combined_id)
|
|
if not url:
|
|
# Fall back to guess_url if URL not found in map
|
|
url = guess_url(front_matter, doc_id, url_template)
|
|
|
|
extras = {k: v for k, v in response_payload.items() if k not in canonical_keys}
|
|
extras_json = json.dumps(extras, ensure_ascii=False) if extras else None
|
|
|
|
records.append(
|
|
DocumentRecord(
|
|
doc_id=doc_id,
|
|
markdown_path=md_path,
|
|
pdf_path=pdf_path,
|
|
pdf_size=pdf_size,
|
|
primary_language=response_payload.get("primary_language"),
|
|
is_rotation_valid=response_payload.get("is_rotation_valid"),
|
|
rotation_correction=response_payload.get("rotation_correction"),
|
|
is_table=response_payload.get("is_table"),
|
|
is_diagram=response_payload.get("is_diagram"),
|
|
natural_text=response_payload.get("natural_text"),
|
|
page_number=page_number,
|
|
url=url,
|
|
extras_json=extras_json,
|
|
)
|
|
)
|
|
except Exception as exc:
|
|
if strict:
|
|
raise
|
|
tqdm.write(f"[WARN] Skipping {md_path}: {exc}")
|
|
|
|
return records
|
|
|
|
|
|
def write_parquet(records: List[DocumentRecord], parquet_path: Path, compression: str) -> None:
|
|
"""Emit the textual payload into a parquet file."""
|
|
if not records:
|
|
raise RuntimeError("No records to write into parquet")
|
|
|
|
pdf_relpaths: List[str] = []
|
|
for rec in records:
|
|
path_value = rec.pdf_relpath or f"{rec.doc_id}.pdf"
|
|
pdf_relpaths.append(path_value)
|
|
|
|
data = {
|
|
"url": [rec.url for rec in records],
|
|
"page_number": [rec.page_number for rec in records],
|
|
"pdf_relpath": pdf_relpaths,
|
|
"primary_language": [rec.primary_language for rec in records],
|
|
"is_rotation_valid": [rec.is_rotation_valid for rec in records],
|
|
"rotation_correction": [rec.rotation_correction for rec in records],
|
|
"is_table": [rec.is_table for rec in records],
|
|
"is_diagram": [rec.is_diagram for rec in records],
|
|
"natural_text": [rec.natural_text for rec in records],
|
|
"extras": [rec.extras_json for rec in records],
|
|
}
|
|
index = [rec.doc_id for rec in records]
|
|
df = pd.DataFrame(data, index=index)
|
|
df.index.name = "id"
|
|
|
|
parquet_path.parent.mkdir(parents=True, exist_ok=True)
|
|
df.to_parquet(parquet_path, compression=compression)
|
|
|
|
|
|
def chunk_records_by_size(records: List[DocumentRecord], max_bytes: int) -> Iterator[List[DocumentRecord]]:
|
|
"""Yield batches of records whose summed PDF sizes stay under max_bytes."""
|
|
batch: List[DocumentRecord] = []
|
|
batch_size = 0
|
|
overhead = 1024 # rough tar header allowance per entry
|
|
|
|
for record in records:
|
|
entry_size = record.pdf_size + overhead
|
|
if entry_size > max_bytes:
|
|
raise RuntimeError(f"Single PDF {record.pdf_path} exceeds max tar size {max_bytes} bytes")
|
|
|
|
if batch and batch_size + entry_size > max_bytes:
|
|
yield batch
|
|
batch = []
|
|
batch_size = 0
|
|
|
|
batch.append(record)
|
|
batch_size += entry_size
|
|
|
|
if batch:
|
|
yield batch
|
|
|
|
|
|
def write_pdf_tarballs(
|
|
records: List[DocumentRecord],
|
|
pdf_dir: Path,
|
|
chunk_prefix: str,
|
|
max_bytes: int,
|
|
manifest_path: Path,
|
|
chunk_dir_name: str,
|
|
) -> None:
|
|
"""Bundle PDFs into .tar.gz archives under the size cap."""
|
|
pdf_dir.mkdir(parents=True, exist_ok=True)
|
|
manifest_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
manifest_rows: List[Dict[str, str]] = []
|
|
batches = chunk_records_by_size(records, max_bytes)
|
|
|
|
normalized_dir = chunk_dir_name.strip().strip("/") if chunk_dir_name else ""
|
|
|
|
for chunk_idx, batch in tqdm(enumerate(batches), desc="Writing PDF tarballs"):
|
|
tar_name = f"{chunk_prefix}_{chunk_idx:05d}.tar.gz"
|
|
tar_path = pdf_dir / tar_name
|
|
with tarfile.open(tar_path, "w:gz", dereference=True) as tar:
|
|
for rec in batch:
|
|
tar.add(rec.pdf_path, arcname=f"{rec.doc_id}.pdf", recursive=False)
|
|
rec.chunk_name = tar_name
|
|
inner_ref = f"{tar_name}:{rec.doc_id}.pdf"
|
|
rec.pdf_relpath = f"{normalized_dir}/{inner_ref}" if normalized_dir else inner_ref
|
|
manifest_rows.append({"doc_id": rec.doc_id, "chunk": tar_name, "arcname": f"{rec.doc_id}.pdf", "pdf_relpath": rec.pdf_relpath})
|
|
|
|
actual_size = tar_path.stat().st_size
|
|
if actual_size > max_bytes:
|
|
raise RuntimeError(f"{tar_path} exceeded size cap ({actual_size} bytes > {max_bytes} bytes)")
|
|
|
|
with manifest_path.open("w", encoding="utf-8") as manifest_file:
|
|
for row in manifest_rows:
|
|
manifest_file.write(json.dumps(row) + "\n")
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Repackage processed olmocr-mix data into parquet + PDF tarballs.")
|
|
parser.add_argument("--processed-dir", required=True, type=Path, help="Directory with markdown/PDF pairs (output of prepare_olmocrmix.py).")
|
|
parser.add_argument("--subset", required=True, help="Dataset subset identifier (e.g. 00_documents).")
|
|
parser.add_argument("--split", required=True, help="Dataset split identifier (e.g. train_s2pdf).")
|
|
parser.add_argument(
|
|
"--output-dir",
|
|
required=True,
|
|
type=Path,
|
|
help="Destination directory for the parquet file and pdf tarballs.",
|
|
)
|
|
parser.add_argument(
|
|
"--parquet-name",
|
|
default=None,
|
|
help="Filename for the generated parquet file (defaults to {subset}_{split}.parquet).",
|
|
)
|
|
parser.add_argument(
|
|
"--pdf-chunk-dir",
|
|
default="pdf_tarballs",
|
|
help="Name of the subdirectory (under output-dir) to place PDF tarballs in.",
|
|
)
|
|
parser.add_argument(
|
|
"--pdf-chunk-prefix",
|
|
default=None,
|
|
help="Prefix for generated tarball filenames (defaults to {subset}_{split}).",
|
|
)
|
|
parser.add_argument(
|
|
"--max-tar-size-bytes",
|
|
type=int,
|
|
default=DEFAULT_MAX_TAR_BYTES,
|
|
help="Maximum uncompressed size (in bytes) to pack into a single tarball (default 1 GiB).",
|
|
)
|
|
parser.add_argument(
|
|
"--pdf-root",
|
|
type=Path,
|
|
default=None,
|
|
help="Optional directory containing {doc_id}.pdf files if they are not alongside the markdown.",
|
|
)
|
|
parser.add_argument(
|
|
"--url-template",
|
|
type=str,
|
|
default=None,
|
|
help="Optional template to synthesize URLs, e.g. 's3://bucket/{prefix}/{base_pdf}.pdf'.",
|
|
)
|
|
parser.add_argument(
|
|
"--parquet-compression",
|
|
default="snappy",
|
|
help="Compression codec passed to pandas.to_parquet (default: snappy).",
|
|
)
|
|
parser.add_argument(
|
|
"--manifest-name",
|
|
default="pdf_chunk_manifest.jsonl",
|
|
help="Filename for the emitted chunk manifest (stored under output-dir).",
|
|
)
|
|
parser.add_argument("--strict", action="store_true", help="Fail immediately when a markdown/PDF pair cannot be processed.")
|
|
return parser.parse_args()
|
|
|
|
|
|
def build_dataset_tag(subset: str, split: str) -> str:
|
|
"""Normalize subset/split into a filesystem-friendly tag."""
|
|
return f"{subset.strip().replace('/', '_')}_{split.strip().replace('/', '_')}"
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
|
|
processed_dir = args.processed_dir.expanduser().resolve()
|
|
if not processed_dir.exists():
|
|
raise FileNotFoundError(f"Processed directory not found: {processed_dir}")
|
|
|
|
pdf_root = args.pdf_root.expanduser().resolve() if args.pdf_root else None
|
|
output_dir = args.output_dir.expanduser().resolve()
|
|
dataset_tag = build_dataset_tag(args.subset, args.split)
|
|
parquet_name = args.parquet_name or f"{dataset_tag}.parquet"
|
|
chunk_prefix = args.pdf_chunk_prefix or dataset_tag
|
|
parquet_path = output_dir / parquet_name
|
|
pdf_dir = output_dir / args.pdf_chunk_dir
|
|
manifest_path = output_dir / args.manifest_name
|
|
|
|
records = collect_documents(processed_dir, pdf_root, args.url_template, args.strict)
|
|
if not records:
|
|
raise RuntimeError("No markdown/PDF pairs discovered - nothing to package.")
|
|
|
|
records.sort(key=lambda rec: rec.doc_id)
|
|
|
|
write_pdf_tarballs(records, pdf_dir, chunk_prefix, args.max_tar_size_bytes, manifest_path, args.pdf_chunk_dir)
|
|
write_parquet(records, parquet_path, args.parquet_compression)
|
|
|
|
print(f"Wrote parquet: {parquet_path}")
|
|
print(f"Wrote PDF tarballs to: {pdf_dir}")
|
|
print(f"Wrote manifest: {manifest_path}")
|
|
print(f"Total documents packaged: {len(records)}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|