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
399 lines
16 KiB
Python
399 lines
16 KiB
Python
import argparse
|
|
import json
|
|
import shutil
|
|
import tarfile
|
|
from concurrent.futures import ProcessPoolExecutor, as_completed
|
|
from os import PathLike
|
|
from pathlib import Path
|
|
from typing import Any, Optional
|
|
|
|
import pandas as pd
|
|
from huggingface_hub import snapshot_download
|
|
from tqdm import tqdm
|
|
|
|
|
|
def extract_tarball(tarball_path: Path, extract_dir: Path) -> int:
|
|
"""Extract a single tarball and return the number of files extracted."""
|
|
try:
|
|
with tarfile.open(tarball_path, "r") as tar:
|
|
# Extract with overwrite for existing files
|
|
members = tar.getmembers()
|
|
for member in members:
|
|
try:
|
|
tar.extract(member, extract_dir)
|
|
except (OSError, IOError) as e:
|
|
# If extraction fails due to existing file, try to remove and re-extract
|
|
target_path = extract_dir / member.name
|
|
if target_path.exists():
|
|
if target_path.is_dir():
|
|
# Skip existing directories
|
|
continue
|
|
else:
|
|
# Remove existing file and re-extract
|
|
target_path.unlink()
|
|
tar.extract(member, extract_dir)
|
|
else:
|
|
# Re-raise if it's not a file exists issue
|
|
raise e
|
|
return len(members)
|
|
except Exception as e:
|
|
print(f"Error extracting {tarball_path}: {e}")
|
|
return 0
|
|
|
|
|
|
PAGE_RESPONSE_COLUMNS = [
|
|
"primary_language",
|
|
"is_rotation_valid",
|
|
"rotation_correction",
|
|
"is_table",
|
|
"is_diagram",
|
|
"natural_text",
|
|
]
|
|
|
|
|
|
def _coerce_optional(value: Any) -> Optional[Any]:
|
|
"""Convert pandas nulls to None."""
|
|
if pd.isna(value):
|
|
return None
|
|
return value
|
|
|
|
|
|
def _coerce_bool(value: Any, default: bool) -> bool:
|
|
if value is None or pd.isna(value):
|
|
return default
|
|
if isinstance(value, bool):
|
|
return value
|
|
if isinstance(value, (int, float)):
|
|
return bool(int(value))
|
|
if isinstance(value, str):
|
|
lowered = value.strip().lower()
|
|
if lowered in {"true", "1", "yes", "y"}:
|
|
return True
|
|
if lowered in {"false", "0", "no", "n"}:
|
|
return False
|
|
return default
|
|
|
|
|
|
def _coerce_rotation(value: Any, default: int = 0) -> int:
|
|
if value is None or pd.isna(value):
|
|
return default
|
|
try:
|
|
rotation = int(value)
|
|
if rotation in {0, 90, 180, 270}:
|
|
return rotation
|
|
except (TypeError, ValueError):
|
|
pass
|
|
return default
|
|
|
|
|
|
def _coerce_text(value: Any) -> Optional[str]:
|
|
if value is None or pd.isna(value):
|
|
return None
|
|
text = str(value)
|
|
return text if text.strip() else None
|
|
|
|
|
|
def extract_response_from_row(row: pd.Series) -> dict[str, Any]:
|
|
"""Return a PageResponse-like dict regardless of parquet schema."""
|
|
response_data: dict[str, Any] = {}
|
|
raw_response = row.get("response")
|
|
|
|
if isinstance(raw_response, str):
|
|
stripped = raw_response.strip()
|
|
if stripped:
|
|
try:
|
|
response_data = json.loads(stripped)
|
|
except json.JSONDecodeError:
|
|
response_data = {}
|
|
elif isinstance(raw_response, dict):
|
|
response_data = dict(raw_response)
|
|
|
|
if not response_data:
|
|
for column in PAGE_RESPONSE_COLUMNS:
|
|
if column in row:
|
|
response_data[column] = _coerce_optional(row[column])
|
|
|
|
extras = row.get("extras")
|
|
if isinstance(extras, str):
|
|
extras = extras.strip()
|
|
if extras:
|
|
try:
|
|
response_data.update(json.loads(extras))
|
|
except json.JSONDecodeError:
|
|
pass
|
|
elif isinstance(extras, dict):
|
|
response_data.update(extras)
|
|
|
|
response_data["primary_language"] = _coerce_optional(response_data.get("primary_language"))
|
|
response_data["is_rotation_valid"] = _coerce_bool(response_data.get("is_rotation_valid"), True)
|
|
response_data["rotation_correction"] = _coerce_rotation(response_data.get("rotation_correction"), 0)
|
|
response_data["is_table"] = _coerce_bool(response_data.get("is_table"), False)
|
|
response_data["is_diagram"] = _coerce_bool(response_data.get("is_diagram"), False)
|
|
response_data["natural_text"] = _coerce_text(response_data.get("natural_text"))
|
|
|
|
return response_data
|
|
|
|
|
|
def prepare_olmocr_mix(dataset_path: str, subset: str, split: str, destination: str | PathLike, max_examples: Optional[int] = None) -> str:
|
|
"""
|
|
Prepare OLMoCR mix dataset by downloading from HuggingFace and organizing into a folder structure.
|
|
|
|
Args:
|
|
dataset_path: HuggingFace dataset path
|
|
subset: Dataset subset name
|
|
split: Dataset split (train/validation/test)
|
|
destination: Destination directory path
|
|
max_examples: Maximum number of examples to process (None for all)
|
|
"""
|
|
# Step 1: Download dataset using hugging face hub snapshot_download to destination/hugging_face folder
|
|
dest_path = Path(destination)
|
|
hugging_face_dir = dest_path / "hugging_face"
|
|
hugging_face_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
if Path(dataset_path).exists():
|
|
print("Dataset path is a local folder, using that")
|
|
local_dir = dataset_path
|
|
shutil.copytree(local_dir, hugging_face_dir, dirs_exist_ok=True)
|
|
else:
|
|
print(f"Downloading dataset {dataset_path} to {hugging_face_dir}...")
|
|
|
|
# For allenai/olmOCR-mix-0225, download everything as before
|
|
# For other datasets, filter to only download needed files
|
|
if dataset_path == "allenai/olmOCR-mix-0225":
|
|
# Download the entire repository including PDFs and parquet files
|
|
local_dir = snapshot_download(
|
|
repo_id=dataset_path,
|
|
repo_type="dataset",
|
|
local_dir=hugging_face_dir,
|
|
)
|
|
else:
|
|
# For other datasets, only download the specific parquet file and related PDF tarballs
|
|
# Construct the dataset tag for filtering
|
|
dataset_tag = f"{subset}_{split}"
|
|
|
|
# Define patterns to allow:
|
|
# 1. The specific parquet file
|
|
# 2. Related PDF tarballs in pdf_tarballs directory
|
|
# 3. README and metadata files (for dataset info)
|
|
# 4. urls.jsonl for URL mappings if it exists
|
|
allow_patterns = [
|
|
f"{dataset_tag}.parquet",
|
|
f"pdf_tarballs/{dataset_tag}_*.tar.gz",
|
|
"README.md",
|
|
"*.json", # Include any metadata JSON files
|
|
]
|
|
|
|
print(f"Filtering download to patterns: {allow_patterns}")
|
|
|
|
local_dir = snapshot_download(
|
|
repo_id=dataset_path,
|
|
repo_type="dataset",
|
|
local_dir=hugging_face_dir,
|
|
allow_patterns=allow_patterns,
|
|
)
|
|
|
|
print(f"Downloaded to: {local_dir}")
|
|
|
|
# Step 2: Create destination folder structure for processed markdown files
|
|
processed_dir = dest_path / f"processed_{subset}_{split}"
|
|
processed_dir.mkdir(exist_ok=True)
|
|
|
|
# Manual map to parquet files for now
|
|
if dataset_path == "allenai/olmOCR-mix-0225":
|
|
if subset == "00_documents" and split == "train_s2pdf":
|
|
parquet_files = [dest_path / "hugging_face" / "train-s2pdf.parquet"]
|
|
elif subset == "00_documents" and split == "eval_s2pdf":
|
|
parquet_files = [dest_path / "hugging_face" / "eval-s2pdf.parquet"]
|
|
elif subset == "01_books" and split == "train_iabooks":
|
|
parquet_files = [dest_path / "hugging_face" / "train-iabooks.parquet"]
|
|
elif subset == "01_books" and split == "eval_iabooks":
|
|
parquet_files = [dest_path / "hugging_face" / "eval-iabooks.parquet"]
|
|
else:
|
|
raise NotImplementedError()
|
|
else:
|
|
parquet_files = [dest_path / "hugging_face" / f"{subset}_{split}.parquet"]
|
|
|
|
# Step 3: Extract PDF tarballs
|
|
pdf_tarballs_dir = dest_path / "hugging_face" / "pdf_tarballs"
|
|
if pdf_tarballs_dir.exists():
|
|
extracted_dir = pdf_tarballs_dir / "extracted"
|
|
extracted_dir.mkdir(exist_ok=True)
|
|
|
|
# Check if PDFs are already extracted
|
|
existing_pdfs = list(extracted_dir.glob("*.pdf"))
|
|
if existing_pdfs:
|
|
print(f"Found {len(existing_pdfs)} already extracted PDFs in {extracted_dir}, skipping extraction step")
|
|
else:
|
|
# Find tarball files based on dataset type
|
|
if dataset_path == "allenai/olmOCR-mix-0225":
|
|
# Extract all tarballs for the full dataset
|
|
tarball_files = list(pdf_tarballs_dir.glob("*.tar*")) + list(pdf_tarballs_dir.glob("*.tgz"))
|
|
else:
|
|
# Only extract tarballs matching the dataset_tag pattern
|
|
dataset_tag = f"{subset}_{split}"
|
|
tarball_files = list(pdf_tarballs_dir.glob(f"{dataset_tag}_*.tar*")) + list(pdf_tarballs_dir.glob(f"{dataset_tag}_*.tgz"))
|
|
print(f"Filtering tarballs to pattern: {dataset_tag}_*")
|
|
|
|
if tarball_files:
|
|
print(f"\nFound {len(tarball_files)} PDF tarballs to extract...")
|
|
|
|
# Use ProcessPoolExecutor for parallel extraction
|
|
with ProcessPoolExecutor() as executor:
|
|
# Submit all tasks
|
|
future_to_tarball = {}
|
|
for tarball in tarball_files:
|
|
future = executor.submit(extract_tarball, tarball, extracted_dir)
|
|
future_to_tarball[future] = tarball
|
|
|
|
# Process results as they complete with progress bar
|
|
total_files_extracted = 0
|
|
with tqdm(total=len(tarball_files), desc="Extracting tarballs") as pbar:
|
|
for future in as_completed(future_to_tarball):
|
|
tarball = future_to_tarball[future]
|
|
try:
|
|
files_extracted = future.result()
|
|
total_files_extracted += files_extracted
|
|
pbar.set_postfix({"files": total_files_extracted})
|
|
except Exception as e:
|
|
print(f"\nError with {tarball.name}: {e}")
|
|
pbar.update(1)
|
|
|
|
print(f"Extracted {total_files_extracted} files from tarballs to {extracted_dir}")
|
|
else:
|
|
print(f"No PDF tarballs directory found at {pdf_tarballs_dir}")
|
|
|
|
# Step 4: Process parquet files
|
|
total_processed = 0
|
|
total_errors = 0
|
|
|
|
# Create urls.jsonl file for id-to-url mappings
|
|
urls_file_path = processed_dir / "urls.jsonl"
|
|
urls_file = open(urls_file_path, "w", encoding="utf-8")
|
|
|
|
for parquet_file in parquet_files:
|
|
print(f"Processing {parquet_file.name}...")
|
|
df = pd.read_parquet(parquet_file)
|
|
|
|
# Process each row
|
|
for idx, row in df.iterrows():
|
|
if max_examples and total_processed >= max_examples:
|
|
break
|
|
|
|
try:
|
|
|
|
response = extract_response_from_row(row)
|
|
doc_id = str(idx)
|
|
|
|
assert len(doc_id) > 4
|
|
|
|
# Extract URL from row and write to urls.jsonl
|
|
url = row.get("url", None)
|
|
if url:
|
|
url_entry = {"id": doc_id, "url": url}
|
|
urls_file.write(json.dumps(url_entry) + "\n")
|
|
|
|
# Create folder structure
|
|
# For allenai/olmOCR-mix-0225: use first 4 characters as folder
|
|
# For other datasets: preserve the existing structure
|
|
|
|
if dataset_path == "allenai/olmOCR-mix-0225":
|
|
# Standard format: use first 4 characters as folder
|
|
folder_name = doc_id[:4]
|
|
file_name = f"{doc_id[4:]}.md"
|
|
|
|
# Create directory
|
|
output_dir = processed_dir / folder_name
|
|
output_dir.mkdir(exist_ok=True)
|
|
else:
|
|
# Custom format: preserve directory structure from doc_id
|
|
# The doc_id already contains the full path structure
|
|
if "/" in doc_id:
|
|
# doc_id contains path separators
|
|
path_parts = doc_id.rsplit("/", 1)
|
|
folder_path = Path(path_parts[0])
|
|
file_name = f"{path_parts[1]}.md"
|
|
output_dir = processed_dir / folder_path
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
else:
|
|
# No path separator, put at root
|
|
file_name = f"{doc_id}.md"
|
|
output_dir = processed_dir
|
|
|
|
# Write markdown file with front matter and natural text
|
|
output_file = output_dir / file_name
|
|
with open(output_file, "w", encoding="utf-8") as f:
|
|
# Extract natural_text and other fields for front matter
|
|
natural_text = response.get("natural_text", "")
|
|
# Create front matter from other fields
|
|
front_matter = {k: v for k, v in response.items() if k != "natural_text"}
|
|
|
|
# Write front matter
|
|
f.write("---\n")
|
|
for k, v in front_matter.items():
|
|
f.write(f"{k}: {v}\n")
|
|
|
|
if natural_text is not None and len(natural_text.strip()) > 0:
|
|
f.write("---\n")
|
|
|
|
# Write natural text
|
|
f.write(natural_text)
|
|
else:
|
|
f.write("---")
|
|
|
|
# Look for matching PDF in extracted directory and create symlinks
|
|
extracted_pdfs_dir = dest_path / "hugging_face" / "pdf_tarballs" / "extracted"
|
|
|
|
# Find PDFs that match the ID pattern
|
|
matched_pdf_path = extracted_pdfs_dir / f"{doc_id}.pdf"
|
|
assert matched_pdf_path.exists(), "Matching PDF not found"
|
|
|
|
# Create symlink path based on dataset type
|
|
if dataset_path == "allenai/olmOCR-mix-0225":
|
|
symlink_path = output_dir / f"{doc_id[4:]}.pdf"
|
|
else:
|
|
# For custom datasets, use the same filename as the markdown
|
|
symlink_path = output_file.with_suffix(".pdf")
|
|
|
|
# Create relative symlink to the PDF
|
|
if not symlink_path.exists():
|
|
symlink_path.symlink_to(matched_pdf_path)
|
|
|
|
total_processed += 1
|
|
if total_processed % 1000 == 0:
|
|
print(f"Processed {total_processed} examples...")
|
|
except Exception as ex:
|
|
print(f"Error processing line: {ex}")
|
|
total_errors += 1
|
|
|
|
if max_examples and total_processed >= max_examples:
|
|
break
|
|
|
|
# Close the urls.jsonl file
|
|
urls_file.close()
|
|
print(f"Created urls.jsonl with {total_processed} id-to-url mappings")
|
|
|
|
print(f"Completed! Processed {total_processed} examples to {processed_dir}")
|
|
print(f"Total errors: {total_errors}")
|
|
|
|
return str(processed_dir)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Prepare OLMoCR mix dataset")
|
|
parser.add_argument("--dataset-path", type=str, default="allenai/olmOCR-mix-0225", help="HuggingFace dataset path (e.g., 'allenai/olmocr-mix')")
|
|
|
|
# Add subset and split to the parser (not the group) but they'll be validated later
|
|
parser.add_argument("--subset", type=str, default=None, help="Dataset subset name")
|
|
parser.add_argument("--split", type=str, default=None, help="Dataset split ex eval_s2pdf")
|
|
|
|
parser.add_argument("--destination", type=str, required=True, help="Destination directory path")
|
|
parser.add_argument("--max-examples", type=int, default=None, help="Maximum number of examples to process (default: all)")
|
|
|
|
args = parser.parse_args()
|
|
|
|
prepare_olmocr_mix(dataset_path=args.dataset_path, subset=args.subset, split=args.split, destination=args.destination, max_examples=args.max_examples)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|