PDF to Markdown Conversion Guide

import os
import json
import base64
import time
import argparse
from pathlib import Path
from mistralai import Mistral


def encode_pdf_to_base64(pdf_path):
    """Encode a PDF file to base64."""
    try:
        with open(pdf_path, "rb") as pdf_file:
            return base64.b64encode(pdf_file.read()).decode('utf-8')
    except Exception as e:
        print(f"Error encoding {pdf_path}: {e}")
        return None


def find_pdf_files(input_folder, max_size_mb=36):
    """Recursively find all PDF files in the input folder that are within size limit."""
    pdf_files = []
    skipped_files = []
    max_size_bytes = max_size_mb * 1024 * 1024  # Convert MB to bytes

    for root, dirs, files in os.walk(input_folder):
        for file in files:
            if file.lower().endswith('.pdf'):
                full_path = os.path.join(root, file)
                relative_path = os.path.relpath(full_path, input_folder)

                # Check file size
                try:
                    file_size = os.path.getsize(full_path)
                    if file_size <= max_size_bytes:
                        pdf_files.append((full_path, relative_path))
                    else:
                        size_mb = file_size / (1024 * 1024)
                        skipped_files.append((relative_path, size_mb))
                        print(f"Skipping {relative_path} ({size_mb:.1f}MB) - exceeds {max_size_mb}MB limit")
                except OSError as e:
                    print(f"Error checking size of {relative_path}: {e}")

    if skipped_files:
        print(f"\nSkipped {len(skipped_files)} file(s) due to size limit:")
        for path, size in skipped_files[:5]:  # Show first 5
            print(f"  - {path} ({size:.1f}MB)")
        if len(skipped_files) > 5:
            print(f"  ... and {len(skipped_files) - 5} more")

    return pdf_files


def create_batch_file(pdf_files, batch_file_path, input_folder):
    """Create a JSONL batch file for OCR processing."""
    entries = []
    with open(batch_file_path, 'w') as file:
        for index, (pdf_path, relative_path) in enumerate(pdf_files):
            print(f"Encoding PDF {index + 1}/{len(pdf_files)}: {relative_path}")
            base64_pdf = encode_pdf_to_base64(pdf_path)

            if base64_pdf:
                entry = {
                    "custom_id": f"{index}|{relative_path}",  # Store relative path in custom_id
                    "body": {
                        "document": {
                            "type": "document_url",
                            "document_url": f"data:application/pdf;base64,{base64_pdf}"
                        },
                        "include_image_base64": False  # We don't need image base64 for text extraction
                    }
                }
                file.write(json.dumps(entry) + '\n')
                entries.append(entry)
            else:
                print(f"Failed to encode {pdf_path}, skipping...")

    return len(entries)


def create_output_structure(input_folder, output_folder, pdf_files):
    """Create the output folder structure matching the input."""
    for _, relative_path in pdf_files:
        # Get the directory part of the relative path
        relative_dir = os.path.dirname(relative_path)
        if relative_dir:
            output_dir = os.path.join(output_folder, relative_dir)
            os.makedirs(output_dir, exist_ok=True)


def process_batch_results(results_file, output_folder):
    """Process the batch results and save markdown files."""
    processed_count = 0

    with open(results_file, 'r') as f:
        for line in f:
            if line.strip():
                try:
                    result = json.loads(line)
                    custom_id = result.get('custom_id', '')

                    # Extract index and relative path from custom_id
                    if '|' in custom_id:
                        index, relative_path = custom_id.split('|', 1)

                        # Change extension from .pdf to .md
                        md_relative_path = os.path.splitext(relative_path)[0] + '.md'
                        output_path = os.path.join(output_folder, md_relative_path)

                        # Extract the OCR content
                        response = result.get('response', {})
                        if response and response.get('status_code') == 200:
                            body = response.get('body', {})
                            pages = body.get('pages', [])

                            # Create output directory if needed
                            os.makedirs(os.path.dirname(output_path), exist_ok=True)

                            # Write markdown content
                            with open(output_path, 'w', encoding='utf-8') as md_file:
                                for i, page in enumerate(pages):
                                    md_file.write(f"# Page {page.get('index', i) + 1}\n\n")
                                    md_file.write(page.get('text', ''))
                                    md_file.write("\n\n---\n\n")

                            processed_count += 1
                        else:
                            status = response.get('status_code') if response else 'Unknown'
                            print(f"Failed to process {relative_path}: status {status}")
                except json.JSONDecodeError as e:
                    print(f"Error parsing result line: {e}")

    return processed_count


def main():
    parser = argparse.ArgumentParser(description="Batch OCR PDF conversion using Mistral OCR API.")
    parser.add_argument("input_folder", type=str, help="Path to folder containing PDF files.")
    parser.add_argument("output_folder", type=str, help="Path to folder where Markdown files will be saved.")
    parser.add_argument("--api-key", type=str, default=None, help="Mistral API key (optional if env var is set).")
    parser.add_argument("--max-size", type=int, default=36, help="Maximum PDF size in MB (default: 36).")

    args = parser.parse_args()

    # Resolve absolute paths
    input_folder = Path(args.input_folder).resolve()
    output_folder = Path(args.output_folder).resolve()

    if not input_folder.exists() or not input_folder.is_dir():
        print(f"Input folder does not exist or is not a directory: {input_folder}")
        return

    print(f"Input folder: {input_folder}")
    print(f"Output folder: {output_folder}")

    # Ensure output folder exists
    output_folder.mkdir(parents=True, exist_ok=True)

    # Resolve API key
    api_key = args.api_key or os.getenv("MISTRAL_API_KEY")
    if not api_key:
        print("Error: No API key provided. Use --api-key or set MISTRAL_API_KEY environment variable.")
        return

    # Initialize Mistral client
    try:
        client = Mistral(api_key=api_key)
        print("Mistral client initialized successfully.")
    except Exception as e:
        print(f"Failed to initialize Mistral client: {e}")
        return

    # Find all PDF files
    print(f"\nSearching for PDF files in: {input_folder}")
    pdf_files = find_pdf_files(str(input_folder), max_size_mb=args.max_size)

    if not pdf_files:
        print("No PDF files found in the input folder.")
        return

    print(f"Found {len(pdf_files)} PDF file(s)")

    # Create output folder structure
    create_output_structure(str(input_folder), str(output_folder), pdf_files)

    # Create batch file
    batch_file_path = "ocr_batch_requests.jsonl"
    print(f"\nCreating batch file: {batch_file_path}")
    num_entries = create_batch_file(pdf_files, batch_file_path, str(input_folder))

    if num_entries == 0:
        print("No valid entries created for batch processing.")
        return

    print(f"Created batch file with {num_entries} entries")

    # Upload batch file
    print("\nUploading batch file...")
    try:
        with open(batch_file_path, "rb") as f:
            batch_data = client.files.upload(
                file={
                    "file_name": batch_file_path,
                    "content": f
                },
                purpose="batch"
            )
        print(f"Batch file uploaded successfully. ID: {batch_data.id}")
    except Exception as e:
        print(f"Error uploading batch file: {e}")
        return

    # Create batch job
    print("\nCreating batch job...")
    try:
        created_job = client.batch.jobs.create(
            input_files=[batch_data.id],
            model="mistral-ocr-latest",
            endpoint="/v1/ocr",
            metadata={"job_type": "batch_ocr_processing"}
        )
        print(f"Batch job created. ID: {created_job.id}")
    except Exception as e:
        print(f"Error creating batch job: {e}")
        return

    # Monitor job progress
    print("\nMonitoring job progress...")
    while True:
        try:
            retrieved_job = client.batch.jobs.get(job_id=created_job.id)

            status = retrieved_job.status
            total = retrieved_job.total_requests
            failed = retrieved_job.failed_requests
            succeeded = retrieved_job.succeeded_requests

            print(f"\rStatus: {status} | Total: {total} | Succeeded: {succeeded} | Failed: {failed}", end='', flush=True)

            if status not in ["QUEUED", "RUNNING"]:
                print()  # New line
                break

            time.sleep(5)  # Check every 5 seconds

        except Exception as e:
            print(f"\nError checking job status: {e}")
            return

    # Download results
    if retrieved_job.status in ["SUCCEEDED", "SUCCESS"] and retrieved_job.output_file:
        print(f"\nJob completed successfully! Downloading results...")
        try:
            # Download the results file
            results_response = client.files.download(file_id=retrieved_job.output_file)
            results_file_path = "ocr_batch_results.jsonl"

            # Handle httpx.Response object
            if hasattr(results_response, 'iter_bytes'):
                results_content = b''.join(results_response.iter_bytes())
            else:
                # Fallback for other response types
                results_content = results_response.content if hasattr(results_response, 'content') else bytes(results_response)

            with open(results_file_path, "wb") as f:
                f.write(results_content)

            print(f"Results downloaded to: {results_file_path}")

            # Process results and create markdown files
            print("\nProcessing results and creating markdown files...")
            processed = process_batch_results(results_file_path, str(output_folder))
            print(f"\nSuccessfully processed {processed} files")

            # Clean up temporary files
            if os.path.exists(batch_file_path):
                os.remove(batch_file_path)
            if os.path.exists(results_file_path):
                os.remove(results_file_path)

        except Exception as e:
            print(f"Error downloading or processing results: {e}")
    else:
        print(f"\nJob failed with status: {retrieved_job.status}")
        if getattr(retrieved_job, "errors", None):
            print(f"Errors: {retrieved_job.errors}")


if __name__ == "__main__":
    main()