test_parallel.py

#!/usr/bin/env python3

import os
import glob
import logging
import torch.multiprocessing as mp
from parallel_multiproc import process_batch_in_parallel

logging.basicConfig(level=logging.INFO)

def main():
    pdf_dir = "/home/user/app/test_pdf"
    output_dir = "/home/user/app/pdf_output"
    os.makedirs(output_dir, exist_ok=True)
    pdf_files = glob.glob(os.path.join(pdf_dir, "*.pdf"))
    logging.info(f"Found {len(pdf_files)} PDF files to process")

    process_batch_in_parallel(
        pdf_paths=pdf_files,
        output_dir=output_dir,
        num_workers=2,  #for our T4 small specifically, do not change it 
        num_gpus=1   
    )

if __name__ == "__main__":
    mp.set_start_method("spawn", force=True)

    main()