allenai--olmocr/olmocr/data/prepare_workspace.py

#!/usr/bin/env python3
"""
Prepare workspace generated by olmocr/pipeline.py for fine-tuning.

This script reads JSONL files from workspace/results, extracts individual pages
from PDFs based on page boundaries, and creates corresponding markdown files.

Usage:
    python prepare_workspace.py workspace_path output_dir [--max-examples N]
"""

import argparse
import json
import logging
import os
import sys
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from urllib.parse import urlparse

import boto3
from pypdf import PdfReader, PdfWriter
from tqdm import tqdm

from olmocr.s3_utils import parse_s3_path

# Set up logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

# Define the metadata columns to write in front matter (in order)
PAGE_RESPONSE_COLUMNS = [
    "primary_language",
    "is_rotation_valid",
    "rotation_correction",
    "is_table",
    "is_diagram",
]


def fetch_s3_file(s3_url: str, local_path: str) -> str:
    """Download a file from an S3 URI (s3://bucket/key) to local_path."""
    parsed = urlparse(s3_url)
    bucket_name = parsed.netloc
    key = parsed.path.lstrip("/")

    # Create directory if it doesn't exist
    os.makedirs(os.path.dirname(local_path), exist_ok=True)

    s3 = boto3.client("s3")
    s3.download_file(bucket_name, key, local_path)
    return local_path


def list_s3_result_files(s3_client, workspace_path: str) -> List[str]:
    """List all JSONL files in the S3 workspace results directory."""
    bucket, prefix = parse_s3_path(workspace_path)
    results_prefix = os.path.join(prefix, "results").rstrip("/") + "/"

    all_files = []
    paginator = s3_client.get_paginator("list_objects_v2")
    for page in paginator.paginate(Bucket=bucket, Prefix=results_prefix):
        if "Contents" in page:
            all_files.extend([f"s3://{bucket}/{obj['Key']}" for obj in page["Contents"] if obj["Key"].endswith(".jsonl")])

    logger.info(f"Found {len(all_files)} JSONL files in S3 workspace")
    return all_files


def download_s3_file(s3_client, s3_path: str) -> str:
    """Download an S3 file and return its contents as a string."""
    bucket, key = parse_s3_path(s3_path)
    response = s3_client.get_object(Bucket=bucket, Key=key)
    return response["Body"].read().decode("utf-8")


def load_jsonl_files(results_dir: Path) -> List[Path]:
    """Load all JSONL files from the workspace results directory."""
    jsonl_files = list(results_dir.glob("*.jsonl"))
    if not jsonl_files:
        logger.error(f"No JSONL files found in {results_dir}")
        return []

    logger.info(f"Found {len(jsonl_files)} JSONL files in {results_dir}")
    return jsonl_files


def parse_jsonl_entry(entry: Dict) -> Optional[Dict]:
    """Parse a single JSONL entry and extract relevant information."""
    try:
        text = entry.get("text", "")
        metadata = entry.get("metadata", {})
        attributes = entry.get("attributes", {})

        source_file = metadata.get("Source-File", "")
        if not source_file:
            logger.warning("Entry missing Source-File in metadata")
            return None

        pdf_page_numbers = attributes.get("pdf_page_numbers", [])
        if not pdf_page_numbers:
            logger.warning(f"Entry for {source_file} missing pdf_page_numbers")
            return None

        # Extract PAGE_RESPONSE_COLUMNS data from attributes
        page_response_data = {}
        for column in PAGE_RESPONSE_COLUMNS:
            page_response_data[column] = attributes.get(column, [])

        return {
            "id": entry.get("id", ""),
            "text": text,
            "source_file": source_file,
            "metadata": metadata,
            "pdf_page_numbers": pdf_page_numbers,
            "page_response_data": page_response_data,
        }
    except Exception as e:
        logger.error(f"Error parsing JSONL entry: {e}")
        return None


def extract_page_text(text: str, page_boundaries: List[List[int]]) -> Dict[int, str]:
    """
    Extract text for each page based on character boundaries.

    Args:
        text: Full document text
        page_boundaries: List of [start_char, end_char, page_num] for each page

    Returns:
        Dictionary mapping page number to extracted text
    """
    page_texts = {}

    for start_char, end_char, page_num in page_boundaries:
        page_text = text[start_char:end_char]
        page_texts[page_num] = page_text

    return page_texts


def extract_pdf_page(pdf_path: str, page_num: int, output_path: str) -> bool:
    """
    Extract a single page from a PDF and save it to output_path.

    Args:
        pdf_path: Path to the source PDF
        page_num: 1-based page number to extract
        output_path: Path where the single-page PDF will be saved

    Returns:
        True if successful, False otherwise
    """
    try:
        reader = PdfReader(pdf_path)

        # Check if page number is valid
        if page_num < 1 or page_num > len(reader.pages):
            logger.error(f"Page {page_num} out of range for {pdf_path} (has {len(reader.pages)} pages)")
            return False

        writer = PdfWriter()
        # PyPDF uses 0-based indexing
        writer.add_page(reader.pages[page_num - 1])

        # Create output directory if it doesn't exist
        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        with open(output_path, "wb") as f:
            writer.write(f)

        return True
    except Exception as e:
        logger.error(f"Error extracting page {page_num} from {pdf_path}: {e}")
        return False


def process_document(entry_data: Dict, output_dir: Path, cache_dir: Path) -> Tuple[int, int]:
    """
    Process a single document: extract pages and create markdown files.

    Returns:
        Tuple of (successful_pages, failed_pages)
    """
    successful = 0
    failed = 0

    source_file = entry_data["source_file"]
    doc_id = entry_data["id"]
    full_text = entry_data["text"]
    pdf_page_numbers = entry_data["pdf_page_numbers"]
    page_response_data = entry_data.get("page_response_data", {})

    # Extract page texts
    page_texts = extract_page_text(full_text, pdf_page_numbers)

    # Download PDF if it's from S3
    if source_file.startswith("s3://"):
        # Create a cache path based on the S3 key
        parsed = urlparse(source_file)
        cache_path = cache_dir / parsed.netloc / parsed.path.lstrip("/")
        local_pdf_path = str(cache_path)

        if not cache_path.exists():
            try:
                logger.info(f"Downloading {source_file} to cache")
                fetch_s3_file(source_file, local_pdf_path)
            except Exception as e:
                logger.error(f"Failed to download {source_file}: {e}")
                return 0, len(page_texts)
        else:
            logger.debug(f"Using cached PDF: {cache_path}")
    else:
        local_pdf_path = source_file

    # Create output subdirectory based on document ID (first 4 characters)
    if len(doc_id) >= 4:
        subdir = doc_id[:4]
        doc_dir = output_dir / subdir
    else:
        doc_dir = output_dir / "misc"

    doc_dir.mkdir(parents=True, exist_ok=True)

    # Create a mapping from page numbers to their indices
    page_num_to_index = {}
    for idx, (_, _, page_num) in enumerate(pdf_page_numbers):
        page_num_to_index[page_num] = idx

    # Process each page
    for page_num, page_text in page_texts.items():
        try:
            # Get the index for this page number
            page_index = page_num_to_index.get(page_num)

            # Create filenames
            base_name = f"{doc_id}_page{page_num}"
            md_path = doc_dir / f"{base_name}.md"
            pdf_path = doc_dir / f"{base_name}.pdf"

            # Write markdown file
            with open(md_path, "w", encoding="utf-8") as f:
                # Write YAML front matter
                f.write("---\n")

                # Write PAGE_RESPONSE_COLUMNS fields in order
                for column in PAGE_RESPONSE_COLUMNS:
                    # Get the value for this page from the attributes lists
                    column_values = page_response_data.get(column, [])

                    if page_index is not None and page_index < len(column_values):
                        value = column_values[page_index]
                        # Convert None to null for YAML
                        if value is None:
                            f.write(f"{column}: null\n")
                        else:
                            f.write(f"{column}: {value}\n")
                    else:
                        # Write default values for missing fields
                        if column == "primary_language":
                            f.write(f"{column}: null\n")
                        elif column == "is_rotation_valid":
                            f.write(f"{column}: true\n")
                        elif column == "rotation_correction":
                            f.write(f"{column}: 0\n")
                        elif column == "is_table":
                            f.write(f"{column}: false\n")
                        elif column == "is_diagram":
                            f.write(f"{column}: false\n")

                # Handle closing delimiter based on whether text exists
                if page_text is not None and len(page_text.strip()) > 0:
                    f.write("---\n")
                    # Write page text
                    f.write(page_text)
                else:
                    # No text or empty text - close without newline
                    f.write("---")

            # Extract PDF page
            if extract_pdf_page(local_pdf_path, page_num, str(pdf_path)):
                successful += 1
                logger.debug(f"Created {md_path} and {pdf_path}")
            else:
                failed += 1
                # Remove the markdown file if PDF extraction failed
                os.remove(md_path)

        except Exception as e:
            logger.error(f"Error processing page {page_num} of document {doc_id}: {e}")
            failed += 1

    return successful, failed


def process_workspace(workspace_path: str, output_dir: Path, max_examples: Optional[int] = None) -> None:
    """
    Process all JSONL files in the workspace and create training data.

    Args:
        workspace_path: Path to the workspace directory (local or S3)
        output_dir: Path to the output directory for training data
        max_examples: Maximum number of documents to process (None for all)
    """
    # Create output and cache directories
    output_dir.mkdir(parents=True, exist_ok=True)
    cache_dir = output_dir / ".pdf_cache"
    cache_dir.mkdir(exist_ok=True)

    # Initialize S3 client if workspace is on S3
    s3_client = None
    if workspace_path.startswith("s3://"):
        s3_client = boto3.client("s3")

    # Parse all entries
    all_entries = []

    if workspace_path.startswith("s3://"):
        # S3 workspace
        jsonl_files = list_s3_result_files(s3_client, workspace_path)
        if not jsonl_files:
            logger.error("No JSONL files found in S3 workspace")
            sys.exit(1)

        for s3_file in jsonl_files:
            logger.info(f"Reading {s3_file}...")
            try:
                content = download_s3_file(s3_client, s3_file)
                for line in content.splitlines():
                    line = line.strip()
                    if not line:
                        continue

                    try:
                        entry = json.loads(line)
                        parsed_entry = parse_jsonl_entry(entry)
                        if parsed_entry:
                            all_entries.append(parsed_entry)
                    except json.JSONDecodeError as e:
                        logger.error(f"JSON decode error in {s3_file}: {e}")
            except Exception as e:
                logger.error(f"Error reading {s3_file}: {e}")
    else:
        # Local workspace
        workspace_path_obj = Path(workspace_path)
        results_dir = workspace_path_obj / "results"
        if not results_dir.exists():
            logger.error(f"Results directory not found: {results_dir}")
            sys.exit(1)

        jsonl_files = load_jsonl_files(results_dir)
        if not jsonl_files:
            sys.exit(1)

        for jsonl_file in jsonl_files:
            logger.info(f"Reading {jsonl_file.name}...")
            with open(jsonl_file, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue

                    try:
                        entry = json.loads(line)
                        parsed_entry = parse_jsonl_entry(entry)
                        if parsed_entry:
                            all_entries.append(parsed_entry)
                    except json.JSONDecodeError as e:
                        logger.error(f"JSON decode error: {e}")

    logger.info(f"Found {len(all_entries)} valid documents to process")

    # Limit entries if max_examples is set
    if max_examples and len(all_entries) > max_examples:
        all_entries = all_entries[:max_examples]
        logger.info(f"Limited to {max_examples} documents")

    # Process documents with progress bar
    total_successful = 0
    total_failed = 0

    with tqdm(total=len(all_entries), desc="Processing documents") as pbar:
        for entry_data in all_entries:
            successful, failed = process_document(entry_data, output_dir, cache_dir)
            total_successful += successful
            total_failed += failed
            pbar.update(1)
            pbar.set_postfix({"pages_ok": total_successful, "pages_failed": total_failed})

    # Print summary
    logger.info("\nProcessing complete!")
    logger.info(f"Successfully processed: {total_successful} pages")
    logger.info(f"Failed: {total_failed} pages")
    logger.info(f"Output directory: {output_dir.absolute()}")


def main():
    parser = argparse.ArgumentParser(description="Prepare workspace data for fine-tuning by extracting individual pages")
    parser.add_argument("workspace_path", type=str, help="Path to the workspace directory containing results folder")
    parser.add_argument("output_dir", type=str, help="Output directory for processed training data")
    parser.add_argument("--max-examples", type=int, default=None, help="Maximum number of documents to process (default: all)")
    parser.add_argument("--debug", action="store_true", help="Enable debug logging")

    args = parser.parse_args()

    if args.debug:
        logger.setLevel(logging.DEBUG)

    workspace_path = args.workspace_path
    output_dir = Path(args.output_dir)

    # Check if workspace exists
    if workspace_path.startswith("s3://"):
        # For S3, we'll check existence when listing files
        logger.info(f"Using S3 workspace: {workspace_path}")
    else:
        workspace_path_obj = Path(workspace_path)
        if not workspace_path_obj.exists():
            logger.error(f"Workspace path does not exist: {workspace_path}")
            sys.exit(1)

    process_workspace(workspace_path, output_dir, args.max_examples)


if __name__ == "__main__":
    main()