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

vlmparse/benchpdf2md/run_benchmark.py

@@ -0,0 +1,296 @@
+import datetime
+import json
+import os
+import tempfile
+import time
+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.bench_tests.benchmark_tsts import (
+    BaselineTest,
+    load_single_test,
+)
+from vlmparse.benchpdf2md.create_dataset import create_dataset
+from vlmparse.benchpdf2md.utils import bootstrap_and_format_results
+from vlmparse.converter_with_server import ConverterWithServer
+from vlmparse.data_model.document import Document
+from vlmparse.servers.utils import get_model_from_uri
+
+
+def process_and_run_benchmark(
+    model: str | None = None,
+    uri: str | None = None,
+    retry: str | None = None,
+    concurrency: int = 1,
+    debug: bool = False,
+    gpu: int = 2,
+    regenerate: bool = False,
+    in_folder: Path | str | None = None,
+    save_folder: Path | str | None = None,
+    retrylast: bool = False,
+    dry_run: bool = True,
+    filter_type: str | list[str] | None = None,
+    filter_category: str | list[str] | None = None,
+    dpi: int | None = None,
+    port: int | None = None,
+    with_vllm_server: bool = False,
+):
+    if in_folder is None:
+        in_folder = os.getenv("IN_FOLDER_FR_BENCHMARK", "pulseia/fr-bench-pdf2md")
+    if save_folder is None:
+        save_folder = os.getenv("OUT_FOLDER_FR_BENCHMARK", ".")
+
+    in_folder = Path(in_folder)
+    save_folder = Path(save_folder)
+
+    if uri is not None:
+        model = get_model_from_uri(uri)
+
+    if model is None:
+        model = "gemini-2.5-flash-lite"
+
+    save_folder = Path(save_folder)
+
+    if in_folder == "pulseia/fr-bench-pdf2md":
+        local_folder_path = snapshot_download(
+            repo_id=in_folder,
+            repo_type="dataset",
+        )
+        in_folder = local_folder_path
+    logger.info(f"In folder: {in_folder}")
+
+    ds = create_dataset(in_folder)
+
+    if filter_type is not None:
+        if isinstance(filter_type, str):
+            filter_type = [filter_type]
+        ds = ds[ds.type.isin(filter_type)]
+
+    if filter_category is not None:
+        assert (
+            filter_category in ds.category.unique()
+        ), f"Filter category {filter_category} not in dataset categories: {ds.category.unique()}"
+        if isinstance(filter_category, str):
+            filter_category = [filter_category]
+        ds = ds[ds.category.isin(filter_category)]
+
+    try:
+        if retrylast:
+            retry = save_folder / (model + "_" + str(dpi) if dpi is not None else 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(ds["pdf_path"])))
+        if retry is None or regenerate:
+            files = list(sorted(set(ds["pdf_path"])))
+            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: {ds['pdf_path'][:5]}"
+                )
+            model_folder = model
+            if dpi is not None:
+                model_folder = model + "_" + str(dpi)
+            save_folder = save_folder / model_folder
+
+            batch_parser = ConverterWithServer(
+                model=model,
+                uri=uri,
+                gpus=str(gpu),
+                with_vllm_server=with_vllm_server,
+                concurrency=concurrency,
+                port=port,
+            )
+
+            if dry_run:
+                logger.info("Dry run, converting first 3 files")
+                batch_parser.parse(
+                    files[:3],
+                    out_folder=tempfile.mkdtemp(),
+                    mode="document",
+                    dpi=dpi,
+                    debug=debug,
+                    retrylast=retrylast,
+                )
+
+            tic = time.perf_counter()
+            batch_parser.parse(
+                files,
+                out_folder=str(save_folder),
+                mode="document",
+                dpi=dpi,
+                debug=debug,
+                retrylast=retrylast,
+            )
+            save_folder = batch_parser.get_out_folder()
+
+            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
+
+        df = run_pb_benchmark(ds, out_folder=save_folder / "results")
+
+        logger.info(
+            f"Number of pages: {ds['pdf_path'].unique().shape[0]}, Number of tests: {len(ds)}"
+        )
+        for col in ["type", "category"]:
+            if col in df.columns:
+                by_col_df = bootstrap_and_format_results(df, col, "result")
+                logger.info(f"By {col}:\n{by_col_df}")
+            if not debug:
+                by_col_df.to_excel(save_folder / f"by_{col}.xlsx")
+
+        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")
+
+            with open(save_folder_test_results / "metrics.json", "w") as f:
+                metrics = {
+                    "total_time": total_time,
+                    "num_pages": len(files),
+                    "num_tests": len(df),
+                    "avg_result": avg,
+                    "avg_doc_latency": df["doc_latency"].mean()
+                    if "doc_latency" in df.columns
+                    else None,
+                    "avg_page_latency": df["page_latency"].mean()
+                    if "page_latency" in df.columns
+                    else None,
+                    "avg_time_per_page": total_time / len(files)
+                    if total_time is not None
+                    else None,
+                    "dpi": dpi,
+                    "concurrency": concurrency,
+                    "model": model,
+                }
+                for k, v in df.groupby("category")["result"].mean().items():
+                    metrics[f"{k}"] = v
+                json.dump(metrics, f)
+
+    except Exception:
+        raise
+
+
+def run_pb_benchmark(
+    ds: pd.DataFrame,
+    out_folder: Path,
+    num_workers: int = -1,
+):
+    files = list(out_folder.rglob("*.zip"))
+    stem_to_zip_path = {path.stem: path for path in files}
+    pdf_to_zip = {}
+    for pdf_path in ds.pdf_path.unique():
+        if Path(pdf_path).stem not in stem_to_zip_path.keys():
+            logger.warning(f"No zip document found for {pdf_path}")
+            continue
+        pdf_to_zip[Path(pdf_path).stem] = stem_to_zip_path[Path(pdf_path).stem]
+
+    def worker(row):
+        zip_path = pdf_to_zip.get(Path(row["pdf_path"]).stem)
+
+        if zip_path is None:
+            return [
+                row
+                | {
+                    "result": False,
+                    "explanation": f"No zip document found for {row['pdf_path']}",
+                    "best_match_score": 0.0,
+                    "document_processed": False,
+                },
+                dict(
+                    result=False,
+                    explanation=f"No zip document found for {row['pdf_path']}",
+                    pdf=row["pdf_path"],
+                    page=row["page"],
+                    id=f"{Path(row['pdf_path']).stem}-baseline",
+                    type="baseline",
+                    category="baseline",
+                ),
+            ]
+        doc = Document.from_zip(zip_path)
+        md_text = doc.text
+        tests_name = Path(doc.file_path).parent.name
+        tests = [load_single_test(row)]
+
+        tests.append(
+            BaselineTest(
+                pdf=row["pdf_path"],
+                page=row["page"],
+                id=f"{tests_name}-baseline",
+                type="baseline",
+                category="baseline",
+            )
+        )
+
+        results = []
+
+        for test in tests:
+            passed, explanation, best_match_score = test.run(md_text)
+            _dict = {
+                "test_id": test.id,
+                "result": passed,
+                "explanation": explanation,
+                "tests_name": tests_name,
+                "pdf_path": str(doc.file_path),
+                "doc_path": str(zip_path),
+                "doc_latency": doc.latency,
+                "page_latency": doc.pages[0].latency,
+                "best_match_score": best_match_score,
+                "document_processed": True,
+            } | test.model_dump()
+
+            results.append(_dict)
+
+        return results
+
+    results = Parallel(n_jobs=num_workers)(
+        delayed(worker)(row) for row in tqdm(ds.to_dict(orient="records"))
+    )
+
+    df = pd.DataFrame([r for r in results for r in r])
+
+    return df
+
+
+def main():
+    fire.Fire(process_and_run_benchmark)
+
+
+if __name__ == "__main__":
+    main()

vlmparse/benchpdf2md/st_visu_benchmark/app.py

@@ -0,0 +1,271 @@
+import argparse
+import subprocess
+import sys
+from glob import glob
+from pathlib import Path
+
+import pandas as pd
+import streamlit as st
+from huggingface_hub import snapshot_download
+from pypdfium2.internal.bases import uuid
+from streamlit import runtime
+
+from vlmparse.benchpdf2md.bench_tests.benchmark_tsts import load_tests, save_tests
+from vlmparse.benchpdf2md.st_visu_benchmark.highligh_text import highlight_text
+from vlmparse.benchpdf2md.st_visu_benchmark.test_form import edit_test_form
+from vlmparse.benchpdf2md.st_visu_benchmark.ui_elements import download_pdf_page
+from vlmparse.benchpdf2md.st_visu_benchmark.utils import get_doc, save_new_test
+
+
+@st.cache_data
+def load_df(results_file):
+    return pd.read_parquet(results_file).set_index("test_id")
+
+
+@st.cache_data
+def get_pdf_map(folder: Path) -> dict[str, Path]:
+    return {path.name: path for path in Path(folder).rglob("*.pdf")}
+
+
+@st.cache_data
+def get_doc_zip_map(folder: Path) -> dict[str, Path]:
+    return {path.name: path for path in Path(folder).rglob("*.zip")}
+
+
+def run_streamlit(folder: str, dataset_path="pulseia/fr-bench-pdf2md") -> None:
+    st.set_page_config(layout="wide")
+    # tests_folder = Path(folder) / "tests"
+    preds_folder = Path(folder)
+
+    # tests = glob(str(tests_folder / "**/**/tests.jsonl"))
+    files = glob(str(preds_folder / "**/**/test_results/**/test_results.parquet"))
+
+    # map_tests = {Path(t).parent.name: t for t in tests}
+    if dataset_path == "pulseia/fr-bench-pdf2md":
+        local_folder_path = snapshot_download(
+            repo_id="pulseia/fr-bench-pdf2md",
+            repo_type="dataset",
+        )
+        dataset_path = local_folder_path
+
+    tests = glob(str(Path(dataset_path) / "**/*.jsonl"), recursive=True)
+
+    map_tests = {Path(t).parent.name: t for t in tests}
+    with st.sidebar:
+        sel_folders = [
+            (
+                Path(f).parent.parent.parent.parent.name,
+                Path(f).parent.parent.parent.name,
+                Path(f).parent.name,
+            )
+            for f in files
+        ]
+
+        if len(sel_folders) == 0:
+            st.error(f"No results found in folder {preds_folder}")
+            return
+        pipe_folder, date1, date2 = st.selectbox("Dir", sel_folders, index=0)
+        res_folder = preds_folder / pipe_folder / date1 / "test_results" / date2
+        df = load_df(res_folder / "test_results.parquet")
+
+        test_type = st.selectbox("Test type", ["present", "absent", "order", "table"])
+        if "category" not in df.columns:
+            df["category"] = None
+        df["category"] = df["category"].map(str)
+        test_category = st.selectbox("Test category", df.category.map(str).unique())
+
+        only_failed = st.checkbox("Only failed", value=False)
+        only_not_checked = st.checkbox("Only not checked", value=False)
+
+        display_image = st.checkbox("Display image", value=False)
+
+        preds_folder = preds_folder / pipe_folder / date1 / "results"
+
+        df_sel = df.loc[(df.type == test_type) & (df.category == test_category)]
+        if only_failed:
+            df_sel = df_sel[~df_sel.result]
+        if only_not_checked:
+            df_sel = df_sel[df_sel.checked != True]  # noqa: E712
+
+        if df_sel.shape[0] == 0:
+            st.markdown("No failed tests found")
+            st.stop()
+        idx = st.number_input(
+            f"Test index (out of {df_sel.shape[0]})",
+            value=0,
+            min_value=0,
+            max_value=df_sel.shape[0] - 1,
+            step=1,
+        )
+
+        row = df_sel.iloc[idx]
+
+        display_markdown = st.checkbox("Display markdown", value=True)
+        show_layout = st.checkbox("Show layout", value=False)
+        display_original_text = st.checkbox("Display original text", value=False)
+        pdf_map = get_pdf_map(Path(dataset_path))
+
+        pdf_path = pdf_map[row.pdf_path.split("/")[-1]]
+
+        download_pdf_page(pdf_path, page_no=0, file_name=f"{row.tests_name}.pdf")
+
+    doc_path = get_doc_zip_map(preds_folder)[row.doc_path.split("/")[-1]]
+    doc = get_doc(doc_path)
+
+    col1_head, col2_head = st.columns(2)
+    with col1_head:
+        pos_buttons = st.container()
+        st.markdown(f"Test: {row.id}" + ("✅" if row.checked else ""))
+        st.markdown("Success: " + str(row.result))
+        st.markdown("Reason: " + row.explanation)
+
+    tests_path = map_tests[row.tests_name]
+
+    if (
+        "tests" not in st.session_state
+        or st.session_state.get("current_tests_path") != tests_path
+    ):
+        st.session_state["tests"] = load_tests(tests_path)
+        st.session_state["current_tests_path"] = tests_path
+
+    if "current_tests_path" not in st.session_state:
+        st.session_state["current_tests_path"] = tests_path
+
+    if display_original_text:
+        res = doc.pages[0].text
+    else:
+
+        @st.cache_data
+        def get_doc_page_md(doc_path):
+            return doc.pages[0].text
+
+        res = get_doc_page_md(row.doc_path)
+
+    with col2_head:
+        _tests = [test for test in st.session_state["tests"] if test.id == row.id]
+
+        if len(_tests) < 1:
+            st.error("No test found")
+            st.stop()
+        elif len(_tests) > 1:
+            st.error("Multiple tests found")
+        test_obj = _tests[0]
+
+        if st.button("Run test"):
+            success, message, best_match_score = test_obj.run(res)
+            st.markdown(f"Success: {success}, score: {best_match_score:.3f}")
+            st.markdown(message)
+
+        add_presence_test = st.checkbox("Add presence test")
+        if add_presence_test:
+            from vlmparse.benchpdf2md.bench_tests.benchmark_tsts import TextPresenceTest
+
+            test_obj_edited = edit_test_form(
+                TextPresenceTest(
+                    pdf=row.pdf_path,
+                    page=0,
+                    id=f"presence_test_{uuid.uuid4()}",
+                    type="present",
+                    text="",
+                ),
+                "present",
+            )
+            if test_obj_edited is not None:
+                st.session_state["tests"].append(test_obj_edited)
+        else:
+            test_obj_edited = edit_test_form(
+                test_obj,
+                test_type,
+            )
+
+        if test_obj_edited is not None:
+            save_new_test(
+                st.session_state["tests"],
+                test_obj_edited,
+                st.session_state["current_tests_path"],
+            )
+
+    col1_button, col2_button, col3_button = pos_buttons.columns(3)
+
+    with col1_button:
+        if st.button("✅ Validate"):
+            test_obj.checked = True
+            save_new_test(
+                st.session_state["tests"],
+                test_obj,
+                st.session_state["current_tests_path"],
+            )
+
+    with col2_button:
+        if test_type != "baseline":
+            if st.button("❌ Reject"):
+                st.session_state["tests"] = [
+                    test for test in st.session_state["tests"] if test.id != row.id
+                ]
+                save_tests(
+                    st.session_state["tests"], st.session_state["current_tests_path"]
+                )
+    with col3_button:
+        if st.button("Supress page (Warning, this is irreversible)"):
+            import shutil
+
+            shutil.rmtree(Path(row.pdf_path).parent)
+
+    def show_text(res):
+        if test_obj:
+            res = highlight_text(test_obj, res)
+
+        with st.container(height=700):
+            if display_markdown:
+                st.markdown(res, unsafe_allow_html=True)
+            else:
+                st.text(res)
+
+    if display_image:
+        with col1_head:
+            show_text(res)
+
+        with col2_head:
+
+            @st.cache_data
+            def get_image(pipe_folder, date, test_id, show_layout):
+                return doc.pages[0].image
+
+            st.image(get_image(pipe_folder, date1, row.id, show_layout))
+    else:
+        show_text(res)
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(description="Document viewer with Streamlit")
+    parser.add_argument(
+        "folder", type=str, nargs="?", default=".", help="Root folder path"
+    )
+    parser.add_argument(
+        "--ds", type=str, default="pulseia/fr-bench-pdf2md", help="Dataset path"
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    """Main entry point."""
+    args = parse_args()
+    folder = args.folder
+
+    if runtime.exists():
+        run_streamlit(folder, dataset_path=args.ds)
+    else:
+        try:
+            subprocess.run(
+                [sys.executable, "-m", "streamlit", "run", __file__, "--", folder],
+                check=True,
+            )
+        except KeyboardInterrupt:
+            print("\nStreamlit app terminated by user.")
+        except subprocess.CalledProcessError as e:
+            print(f"Error while running Streamlit: {e}")
+
+
+if __name__ == "__main__":
+    main()

vlmparse/benchpdf2md/st_visu_benchmark/highligh_text.py

@@ -0,0 +1,117 @@
+import streamlit as st
+from rapidfuzz import process
+from rapidfuzz.distance import Levenshtein
+
+
+class Match:
+    def __init__(self, start, end, dist):
+        self.start = start
+        self.end = end
+        self.dist = dist
+
+
+def find_near_matches(pattern: str, text: str, max_l_dist: int):
+    if not pattern or not text:
+        return []
+
+    matches = []
+    pattern_len = len(pattern)
+
+    for window_size in [pattern_len, pattern_len - 1, pattern_len + 1]:
+        if window_size <= 0 or window_size > len(text):
+            continue
+
+        chunks = [
+            (text[i : i + window_size], i) for i in range(len(text) - window_size + 1)
+        ]
+        if not chunks:
+            continue
+
+        result = process.extractOne(
+            pattern, [c[0] for c in chunks], scorer=Levenshtein.distance
+        )
+
+        if result:
+            matched_text, score, idx = result
+            dist = int(score)
+            if dist <= max_l_dist:
+                start_pos = chunks[idx][1]
+                matches.append(Match(start_pos, start_pos + window_size, dist))
+
+    return matches
+
+
+@st.cache_data
+def highlight_text(test_obj, res):
+    texts_to_highlight = []
+    if hasattr(test_obj, "text"):
+        texts_to_highlight.append(("text", test_obj.text))
+    if hasattr(test_obj, "before"):
+        texts_to_highlight.append(("before", test_obj.before))
+    if hasattr(test_obj, "after"):
+        texts_to_highlight.append(("after", test_obj.after))
+    if hasattr(test_obj, "cell") and test_obj.cell:
+        texts_to_highlight.append(("cell", test_obj.cell))
+    if hasattr(test_obj, "up") and test_obj.up:
+        texts_to_highlight.append(("up", test_obj.up))
+    if hasattr(test_obj, "down") and test_obj.down:
+        texts_to_highlight.append(("down", test_obj.down))
+    if hasattr(test_obj, "left") and test_obj.left:
+        texts_to_highlight.append(("left", test_obj.left))
+    if hasattr(test_obj, "right") and test_obj.right:
+        texts_to_highlight.append(("right", test_obj.right))
+    if hasattr(test_obj, "top_heading") and test_obj.top_heading:
+        texts_to_highlight.append(("top_heading", test_obj.top_heading))
+    if hasattr(test_obj, "left_heading") and test_obj.left_heading:
+        texts_to_highlight.append(("left_heading", test_obj.left_heading))
+
+    matches_with_pos = []
+    for label, txt in texts_to_highlight:
+        if txt and txt.strip():
+            fuzzy_matches = find_near_matches(
+                txt, res, max_l_dist=min(20, len(txt) // 2)
+            )
+            for match in fuzzy_matches:
+                matches_with_pos.append((match.start, match.end, label, match.dist))
+
+    def remove_overlaps(matches):
+        matches = sorted(matches, key=lambda x: (x[3], x[0]))
+        result = []
+        for match in matches:
+            s1, e1, _, _ = match
+            overlapping = False
+            for s2, e2, _, _ in result:
+                if not (e1 <= s2 or e2 <= s1):
+                    overlapping = True
+                    break
+            if not overlapping:
+                result.append(match)
+        return result
+
+    matches_with_pos = remove_overlaps(matches_with_pos)
+    matches_with_pos.sort(key=lambda x: x[0], reverse=True)
+
+    colors = [
+        "yellow",
+        "lightgreen",
+        "lightblue",
+        "lightcoral",
+        "lightyellow",
+        "lightpink",
+        "lightgray",
+        "lavender",
+        "peachpuff",
+        "palegreen",
+    ]
+    label_to_color = {}
+
+    for start, end, label, dist in matches_with_pos:
+        if label not in label_to_color:
+            label_to_color[label] = colors[len(label_to_color) % len(colors)]
+        color = label_to_color[label]
+        res = (
+            res[:start]
+            + f'<span style="background-color: {color}; font-weight: bold;" title="{label}: edit distance={dist}">{res[start:end]}</span>'
+            + res[end:]
+        )
+    return res