olmocr/filter/filter.py

import logging
import re
import subprocess
from collections import Counter
from typing import List

from lingua import Language, LanguageDetectorBuilder
from pypdf import PdfReader

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


class PdfFilter:
    def __init__(
        self,
        languages_to_keep=None,
        apply_form_check=True,
        apply_download_spam_check=True,
        download_spam_threshold=0.004,
    ):
        super().__init__()
        self.language_detector = LanguageDetectorBuilder.from_all_languages().with_preloaded_language_models().build()
        self.languages_to_keep = languages_to_keep if languages_to_keep is not None else [Language.ENGLISH]
        self.apply_form_check = apply_form_check
        self.apply_download_spam_check = apply_download_spam_check
        self.download_spam_threshold = download_spam_threshold

    def _is_form(self, pdf_reader) -> bool:
        # Check if the PDF is a form
        if pdf_reader.get_form_text_fields():
            return True
        return False  # Not a form

    def _is_download_spam(self, base_text: str) -> bool:
        seo_words = {
            "download",
            "pdf",
            "epub",
            "mobi",
            "free",
            "ebook",
            "file",
            "save",
            "casino",
            "viagra",
            "cialis",
            "ciprofloxacin",
        }

        base_text = base_text.strip().lower()
        clean_text = re.sub(r"\W+", " ", base_text)

        word_counts = Counter(clean_text.split())
        total_words = len(clean_text.split())

        if total_words == 0:
            return False

        seo_score = sum(word_counts[word] for word in seo_words if word in word_counts)

        return (seo_score / total_words) > self.download_spam_threshold

    # Returns True if there is something wrong with this PDF
    def filter_out_pdf(self, local_pdf_path: str) -> bool:
        try:
            # Attempt to read the PDF at the beginning
            pdf_reader = PdfReader(local_pdf_path)

            # Form check
            if self.apply_form_check and self._is_form(pdf_reader):
                logger.info(f"Filtering out {local_pdf_path} because it's a form")
                return True  # Filter out
        except Exception as e:
            logger.warning(f"Error reading PDF {local_pdf_path}: {e}")
            return True  # Filter out the PDF if an exception occurs

        # Read the first five pages of text for language calculation
        pdftotext_result = subprocess.run(
            ["pdftotext", "-f", "1", "-l", "5", local_pdf_path, "-"],
            timeout=60,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        if pdftotext_result.returncode != 0:
            logger.warning(f"pdftotext returned {pdftotext_result.returncode} on {local_pdf_path}")
            return True  # Filter out

        base_text = pdftotext_result.stdout.decode("utf-8")

        alpha_count = sum(c.isalpha() for c in base_text)

        if len(base_text) < 200:
            logger.info(f"Keeping {local_pdf_path} on the safe side because not enough text exists in it to analyze")
            return False  # keep the pdf

        if alpha_count / len(base_text) < 0.50:
            logger.info(f"Keeping {local_pdf_path} on the safe side because it's text does not contain many letters so it might be OCRed badly")
            return False  # keep the pdf

        # Language check
        language = self.language_detector.detect_language_of(base_text)
        if language not in self.languages_to_keep:
            logger.info(f"Filtering out {local_pdf_path} because language was {language}")
            return True  # Filter out

        # Download spam check
        if self.apply_download_spam_check and self._is_download_spam(base_text):
            logger.info(f"Filtering out {local_pdf_path} because of SEO/download spam")
            return True  # Filter out

        return False  # Keep the PDF


if __name__ == "__main__":
    import tempfile
    from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait

    import boto3
    from tqdm import tqdm

    from olmocr.s3_utils import parse_s3_path

    # Quiet logs from pypdf
    logging.getLogger("pypdf").setLevel(logging.ERROR)

    def process_pdf(s3_path):
        """
        Process a single PDF file to determine if it should be kept or removed.
        """
        s3_bucket, s3_key = parse_s3_path(s3_path)
        pdf_s3 = boto3.client("s3")

        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=True) as tmp_file:
            pdf_s3.download_fileobj(s3_bucket, s3_key, tmp_file)
            tmp_file.flush()

            # Perform filtering logic
            if filter.filter_out_pdf(tmp_file.name):
                return s3_path, "remove"
            else:
                return s3_path, "keep"

    # Load the list of S3 paths with a progress bar
    with open("/home/ubuntu/s2pdf_paths_1M.txt", "r") as f:
        s3_work_paths: List[str] = list(filter(None, (line.strip() for line in tqdm(f, desc="Loading paths"))))

    # Initialize the PDF filter
    filter = PdfFilter(
        languages_to_keep={Language.ENGLISH, None},  # none means could not detect language, that's okay keep it, might be an OCR
        apply_download_spam_check=True,
        apply_form_check=True,
    )

    # Output files
    keep_path = "/home/ubuntu/s2pdf_paths_filter_keep.txt"
    remove_path = "/home/ubuntu/s2pdf_paths_filter_remove.txt"

    max_pending = 20  # Limit on the number of concurrent futures
    total_pdfs = len(s3_work_paths)
    pdf_iter = iter(s3_work_paths)  # Iterator for PDFs

    # Process the PDFs with limited concurrent futures
    with open(keep_path, "w") as fkeep, open(remove_path, "w") as fremove:
        with ProcessPoolExecutor(max_workers=max_pending) as executor:
            pending_futures = {}

            with tqdm(total=total_pdfs, desc="Processing PDFs") as pbar:
                # Submit initial batch of futures
                for _ in range(min(max_pending, total_pdfs)):
                    s3_path = next(pdf_iter)
                    future = executor.submit(process_pdf, s3_path)
                    pending_futures[future] = s3_path

                while pending_futures:
                    # Wait for the next future to complete
                    done, _ = wait(  # type: ignore
                        pending_futures.keys(),
                        timeout=0.1,
                        return_when=FIRST_COMPLETED,
                    )

                    for future in done:
                        s3_path = pending_futures.pop(future)
                        try:
                            s3_path, result = future.result()
                            if result == "keep":
                                fkeep.write(s3_path + "\n")
                            elif result == "remove":
                                fremove.write(s3_path + "\n")
                        except Exception as e:
                            print(f"Error processing {s3_path}: {e}")

                        # Update the progress bar
                        pbar.update(1)

                        # Submit a new future if there are more PDFs
                        try:
                            s3_path = next(pdf_iter)
                            future = executor.submit(process_pdf, s3_path)
                            pending_futures[future] = s3_path
                        except StopIteration:
                            pass  # No more PDFs to process