vlmparse 0.1.0__py3-none-any.whl → 0.1.3__py3-none-any.whl

vlmparse/benchpdf2md/olmocrbench/repeatdetect.py

@@ -0,0 +1,175 @@
+import random
+import re
+import string
+import time
+import unittest
+
+
+class RepeatDetector:
+    def __init__(self, max_ngram_size: int = 10):
+        self.max_ngram_size = max_ngram_size
+        self.data = ""
+
+    def add_letters(self, new_str: str):
+        self.data += new_str
+
+    def ngram_repeats(self) -> list[int]:
+        result = [0] * self.max_ngram_size
+
+        if not self.data:
+            return result
+
+        # Normalize all whitespace to single spaces
+        text = re.sub(r"\s+", " ", self.data)
+
+        # For each n-gram size
+        for size in range(1, self.max_ngram_size + 1):
+            if len(text) < size:
+                continue
+
+            # Get the last n-gram
+            target = text[-size:]
+
+            # Count backwards from the end to find repeats
+            count = 0
+            pos = len(text) - size  # Start position for previous n-gram
+
+            while pos >= 0:
+                if text[pos : pos + size] == target:
+                    count += 1
+                    pos -= size  # Move back by the size of the n-gram
+                else:
+                    break
+
+            result[size - 1] = count
+
+        return result
+
+
+class RepeatDetectorTest(unittest.TestCase):
+    def test_basicTest1(self):
+        d = RepeatDetector(max_ngram_size=3)
+        d.add_letters("a")
+        self.assertEqual(d.ngram_repeats(), [1, 0, 0])
+
+    def test_basicTest2(self):
+        d = RepeatDetector(max_ngram_size=3)
+        d.add_letters("abab")
+        self.assertEqual(d.ngram_repeats(), [1, 2, 1])
+
+    def test_longer_sequence(self):
+        d = RepeatDetector(max_ngram_size=3)
+        d.add_letters("aabaabaa")
+        self.assertEqual(d.ngram_repeats(), [2, 1, 2])
+
+    def test_no_repeats(self):
+        d = RepeatDetector(max_ngram_size=3)
+        d.add_letters("abc")
+        self.assertEqual(d.ngram_repeats(), [1, 1, 1])
+
+    def test_empty_data(self):
+        d = RepeatDetector(max_ngram_size=3)
+        self.assertEqual(d.ngram_repeats(), [0, 0, 0])
+
+    def test_max_ngram_greater_than_data_length(self):
+        d = RepeatDetector(max_ngram_size=5)
+        d.add_letters("abc")
+        self.assertEqual(d.ngram_repeats(), [1, 1, 1, 0, 0])
+
+    def test_large_single_char(self):
+        d = RepeatDetector(max_ngram_size=5)
+        d.add_letters("a" * 10000)
+        self.assertEqual(d.ngram_repeats(), [10000, 5000, 3333, 2500, 2000])
+
+    def test_repeating_pattern(self):
+        d = RepeatDetector(max_ngram_size=5)
+        d.add_letters("abcabcabcabc")
+        self.assertEqual(d.ngram_repeats(), [1, 1, 4, 1, 1])
+
+    def test_mixed_characters(self):
+        d = RepeatDetector(max_ngram_size=4)
+        d.add_letters("abcdabcabcdabc")
+        self.assertEqual(d.ngram_repeats(), [1, 1, 1, 1])
+
+    def test_palindrome(self):
+        d = RepeatDetector(max_ngram_size=5)
+        d.add_letters("racecar")
+        self.assertEqual(d.ngram_repeats(), [1, 1, 1, 1, 1])
+
+    def test_repeats_not_at_end(self):
+        d = RepeatDetector(max_ngram_size=3)
+        d.add_letters("abcabcxyz")
+        self.assertEqual(d.ngram_repeats(), [1, 1, 1])
+
+    def test_long_repeat_at_end(self):
+        d = RepeatDetector(max_ngram_size=5)
+        d.add_letters("abcabcabcabcabcabcabcabcabcabc")
+        self.assertEqual(d.ngram_repeats(), [1, 1, 10, 1, 1])
+
+    def test_large_repeating_pattern(self):
+        d = RepeatDetector(max_ngram_size=4)
+        pattern = "abcd"
+        repeat_count = 1000
+        d.add_letters(pattern * repeat_count)
+        self.assertEqual(d.ngram_repeats(), [1, 1, 1, repeat_count])
+
+    def test_unicode_characters(self):
+        d = RepeatDetector(max_ngram_size=3)
+        d.add_letters("αβγαβγ")
+        self.assertEqual(d.ngram_repeats(), [1, 1, 2])
+
+    def test_random_data(self):
+        random.seed(42)
+        d = RepeatDetector(max_ngram_size=5)
+        data = "".join(random.choices(string.ascii_letters, k=10000))
+        d.add_letters(data)
+        counts = d.ngram_repeats()
+        for count in counts:
+            self.assertTrue(0 <= count <= len(data))
+
+    def test_special_characters(self):
+        d = RepeatDetector(max_ngram_size=4)
+        d.add_letters("@@##@@##")
+        self.assertEqual(d.ngram_repeats(), [2, 1, 1, 2])
+
+    def test_incremental_addition(self):
+        d = RepeatDetector(max_ngram_size=3)
+        d.add_letters("abc")
+        self.assertEqual(d.ngram_repeats(), [1, 1, 1])
+        d.add_letters("abc")
+        self.assertEqual(d.ngram_repeats(), [1, 1, 2])
+        d.add_letters("abc")
+        self.assertEqual(d.ngram_repeats(), [1, 1, 3])
+
+    def test_long_non_repeating_sequence(self):
+        d = RepeatDetector(max_ngram_size=5)
+        d.add_letters("abcdefghijklmnopqrstuvwxyz")
+        self.assertEqual(d.ngram_repeats(), [1, 1, 1, 1, 1])
+
+    def test_alternating_characters(self):
+        d = RepeatDetector(max_ngram_size=4)
+        d.add_letters("ababababab")
+        self.assertEqual(d.ngram_repeats(), [1, 5, 1, 2])
+
+
+class BenchmarkRepeatDetect(unittest.TestCase):
+    def testLargeRandom(self):
+        all_data = []
+
+        for _ in range(1000):
+            all_data.append("".join(random.choices("a", k=10000)))
+
+        start = time.perf_counter()
+
+        for data in all_data:
+            d = RepeatDetector(max_ngram_size=20)
+            d.add_letters(data)
+            print(d.ngram_repeats())
+
+        end = time.perf_counter()
+
+        print(f"testLargeRandom took {end-start:0.0001f} seconds")
+
+
+if __name__ == "__main__":
+    unittest.main()

vlmparse/benchpdf2md/olmocrbench/run_olmocr_bench.py

@@ -0,0 +1,256 @@
+import datetime
+import json
+import os
+import time
+from dataclasses import asdict
+from pathlib import Path
+
+import fire
+import pandas as pd
+from huggingface_hub import snapshot_download
+from joblib import Parallel, delayed
+from loguru import logger
+from tqdm import tqdm
+
+from vlmparse.benchpdf2md.utils import bootstrap_and_format_results
+from vlmparse.data_model.document import Document
+from vlmparse.registries import converter_config_registry, docker_config_registry
+
+IN_FOLDER = Path(
+    "/mnt/projects/rag-pretraitement/data/docparser/benchmarks/select_difficult_pdf/validated_tests/tiny_test_tests_first_batch/tests/tiny_text_long_text/"
+)
+
+OUT_FOLDER = Path(
+    os.getenv(
+        "OUT_FOLDER_FR_BENCHMARK",
+        "/mnt/projects/rag-pretraitement/data/docparser/benchmarks/fr-bench-pdf2md-preds",
+    )
+)
+IN_FOLDER = Path(
+    "/data/data/docparser/benchmarks/select_difficult_pdf/validated_tests/tiny_test_tests_first_batch/tests/tiny_text_long_text/"
+)
+
+OUT_FOLDER = Path(
+    os.getenv(
+        "OUT_FOLDER_FR_BENCHMARK",
+        "/data/data/docparser/benchmarks/fr-bench-pdf2md-preds",
+    )
+)
+
+
+def process_and_run_benchmark(
+    model="gemini-2.5-flash-lite",
+    uri: str | None = None,
+    retry: str | None = None,
+    concurrency: int = 1,
+    debug: bool = False,
+    gpu: int = 1,
+    regenerate: bool = False,
+    in_folder: Path | str = "allenai/olmOCR-bench",
+    save_folder: Path | str = OUT_FOLDER,
+    retrylast: bool = False,
+    dry_run: bool = True,
+    filter_type: str | list[str] | None = None,
+):
+    save_folder = Path(save_folder)
+
+    # if not in_folder.exists():
+    #     raise ValueError(f"Input folder does not exist: {in_folder}")
+    # if not in_folder.is_dir():
+    #     raise ValueError(f"Input path is not a directory: {in_folder}")
+
+    # ds = create_dataset(in_folder)
+
+    if in_folder == "allenai/olmOCR-bench":
+        local_folder_path = snapshot_download(
+            repo_id=in_folder,
+            repo_type="dataset",  # Use "model" or "space" for other types
+        )
+        in_folder = local_folder_path
+    logger.info(f"In folder: {in_folder}")
+
+    pdfs = list(Path(in_folder).rglob("*.pdf"))
+
+    try:
+        if retrylast:
+            retry = save_folder / model
+            previous_runs = sorted(os.listdir(retry))
+            if len(previous_runs) > 0:
+                retry = retry / previous_runs[-1]
+            else:
+                raise ValueError(
+                    "No previous runs found, do not use the retrylast flag"
+                )
+        files = list(sorted(set(pdfs)))
+        if retry is None or regenerate:
+            files = list(sorted(set(pdfs)))
+            logger.info(f"Number of files to convert: {len(files)}")
+            if retry is not None:
+                already_processed = [
+                    f.removesuffix(".zip") for f in os.listdir(retry / "results")
+                ]
+                files = [
+                    f
+                    for f in files
+                    if Path(f).name.removesuffix(".pdf") not in already_processed
+                ]
+
+                logger.info(f"Number of files after filtering: {len(files)}")
+
+            if len(files) == 0:
+                raise ValueError(
+                    f"No PDF files found in the input folder: {in_folder}\nDataset paths: {pdfs[:5]}"
+                )
+
+            save_folder = (
+                (
+                    save_folder
+                    / model
+                    / (datetime.datetime.now().strftime("%Y-%m-%dT%Hh%Mm%Ss"))
+                )
+                if not retry
+                else retry
+            )
+
+            if uri is None:
+                docker_config = docker_config_registry.get(model)
+                if docker_config is not None:
+                    docker_config.gpu_device_ids = [str(gpu)]
+                    server = docker_config.get_server(auto_stop=True)
+                    server.start()
+                    client = docker_config.get_client()
+                else:
+                    client = converter_config_registry.get(model).get_client()
+            else:
+                client = converter_config_registry.get(model, uri=uri).get_client()
+            client.num_concurrent_pages = concurrency if not debug else 1
+            client.num_concurrent_files = concurrency if not debug else 1
+            client.debug = debug
+
+            if dry_run:
+                client.save_folder = None
+                logger.info("Dry run, converting first 3 files")
+                client.batch(files[:3])
+
+            client.save_folder = str(save_folder)
+            tic = time.perf_counter()
+            client.batch(files)
+            total_time = time.perf_counter() - tic
+            logger.info(
+                f"Time taken to convert {len(files)} files: {total_time:.2f} seconds"
+            )
+
+        else:
+            save_folder = Path(retry)
+            total_time = None
+
+        tests_files = list(Path(in_folder).rglob("**/*.jsonl"))
+        if filter_type is not None:
+            tests_files = [tf for tf in tests_files if filter_type in tf.name]
+
+        df = run_olmocr_benchmark(tests_files, out_folder=save_folder / "results")
+
+        logger.info(
+            f"Number of pages: {df['pdf_path'].unique().shape[0]}, Number of tests: {len(df)}"
+        )
+        if "type" in df.columns:
+            by_type_df = bootstrap_and_format_results(df, "type", "result")
+            logger.info(f"By type:\n{by_type_df}")
+
+        import pdb
+
+        pdb.set_trace()
+
+        if "tests_name" in df.columns:
+            by_tests_name_df = bootstrap_and_format_results(df, "tests_name", "result")
+            logger.info(f"By tests_name:\n{by_tests_name_df}")
+
+        logger.info("average result:")
+        avg = df.loc[df.type != "baseline"]["result"].mean()
+        logger.info(avg)
+
+        if not debug:
+            save_folder_test_results = (
+                save_folder
+                / "test_results"
+                / datetime.datetime.now().strftime("%Y-%m-%dT%Hh%Mm%Ss")
+            )
+            save_folder_test_results.mkdir(parents=True, exist_ok=True)
+            df.to_parquet(save_folder_test_results / "test_results.parquet")
+            by_type_df.to_excel(save_folder_test_results / "by_type.xlsx")
+
+            with open(save_folder_test_results / "metrics.json", "w") as f:
+                json.dump(
+                    {
+                        "total_time": total_time,
+                        "num_pages": len(files),
+                        "num_tests": len(df),
+                        "avg_result": avg,
+                        "avg_doc_latency": df["doc_latency"].mean(),
+                        "avg_page_latency": df["page_latency"].mean(),
+                        "avg_time_per_page": total_time / len(files)
+                        if total_time is not None
+                        else None,
+                    },
+                    f,
+                )
+
+    except Exception:
+        raise
+
+
+def run_olmocr_benchmark(
+    tests_files: list[Path],
+    out_folder: Path,
+    num_workers: int = 64,
+):
+    from vlmparse.benchpdf2md.olmocrbench.tests import load_tests
+
+    files = list(out_folder.rglob("*.zip"))
+    map_files = {path.stem: path for path in files}
+    tests = [test for tf in tests_files for test in load_tests(tf)]
+
+    def worker(test):
+        key = Path(test.pdf).stem
+
+        _dict = {
+            "test_id": test.id,
+        } | asdict(test)
+        if key not in map_files:
+            logger.warning(f"No zip document found for {test.pdf}")
+            _dict["result"] = False
+            _dict["explanation"] = f"No zip document found for {test.pdf}"
+
+        else:
+            file_path = map_files[key]
+
+            doc = Document.from_zip(file_path)
+            md_text = doc.text
+            tests_name = Path(doc.file_path).parent.name
+
+            passed, explanation = test.run(md_text)
+            _dict["result"] = passed
+            _dict["explanation"] = explanation
+            _dict["tests_name"] = tests_name
+            _dict["pdf_path"] = str(doc.file_path)
+            _dict["doc_path"] = str(file_path)
+            _dict["doc_latency"] = doc.latency
+            _dict["page_latency"] = doc.pages[0].latency
+
+        return _dict
+
+    results = Parallel(n_jobs=num_workers)(
+        delayed(worker)(test) for test in tqdm(tests)
+    )
+
+    df = pd.DataFrame(results)
+
+    return df
+
+
+def main():
+    fire.Fire(process_and_run_benchmark)
+
+
+if __name__ == "__main__":
+    main()