parsearena 0.1.3__py3-none-any.whl

parsearena/pipeline.py

@@ -0,0 +1,276 @@
+"""Run parsers over PDFs, one isolated subprocess per parser."""
+
+from __future__ import annotations
+
+import json
+import os
+import platform
+import re
+import signal
+import subprocess
+import time
+from dataclasses import asdict
+from datetime import UTC, datetime
+from pathlib import Path
+
+import pypdfium2 as pdfium
+
+from . import __version__
+from .errors import UserError
+from .evaluate import DocOutput
+from .parsers import ParserSpec, ensure_env, runner_dir
+
+
+def doc_id_for(path: Path, taken: set[str]) -> str:
+    base = re.sub(r"[^A-Za-z0-9._-]+", "-", path.stem).strip("-") or "doc"
+    doc_id, n = base, 2
+    while doc_id in taken:
+        doc_id, n = f"{base}-{n}", n + 1
+    taken.add(doc_id)
+    return doc_id
+
+
+def collect_docs(paths: list[Path]) -> list[dict]:
+    docs, taken = [], set()
+    for p in paths:
+        if not p.exists():
+            raise UserError(f"File not found: {p}")
+        if p.suffix.lower() != ".pdf":
+            raise UserError(f"Not a PDF: {p}", "parsearena v0.1 reads PDF files only.")
+        try:
+            pdf = pdfium.PdfDocument(str(p))
+            n = len(pdf)
+            pdf.close()
+        except pdfium.PdfiumError as exc:
+            raise UserError(f"Could not open {p}: {exc}") from exc
+        docs.append({"id": doc_id_for(p, taken), "file": p.name, "path": str(p.resolve()), "num_pages": n})
+    if not docs:
+        raise UserError("No PDFs given.", "Example: parsearena run docs/*.pdf")
+    return docs
+
+
+def hardware() -> str:
+    machine = platform.machine()
+    cpu = platform.processor() or machine
+    if platform.system() == "Darwin":
+        try:
+            cpu = (
+                subprocess.run(
+                    ["sysctl", "-n", "machdep.cpu.brand_string"], capture_output=True, text=True, timeout=5
+                ).stdout.strip()
+                or cpu
+            )
+            mem = int(
+                subprocess.run(
+                    ["sysctl", "-n", "hw.memsize"], capture_output=True, text=True, timeout=5
+                ).stdout
+            )
+            return f"{cpu}, {round(mem / 2**30)} GB RAM, macOS {platform.mac_ver()[0]}"
+        except (OSError, ValueError, subprocess.SubprocessError):
+            pass
+    mem = ""
+    try:
+        mem = f", {round(os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES') / 2**30)} GB RAM"
+    except (ValueError, OSError, AttributeError):
+        pass
+    return f"{cpu}{mem}, {platform.system()} {platform.release()}"
+
+
+def make_warmup_pdf(src: Path, dst: Path) -> dict:
+    """Two pages: the first page of ``src`` as is, and the same page as an image only."""
+    source = pdfium.PdfDocument(str(src))
+    out = pdfium.PdfDocument.new()
+    try:
+        out.import_pages(source, [0])
+        page = source[0]
+        w, h = page.get_size()
+        bitmap = page.render(scale=150 / 72)
+        image = pdfium.PdfImage.new(out)
+        image.set_bitmap(bitmap)
+        image.set_matrix(pdfium.PdfMatrix().scale(w, h))
+        new = out.new_page(w, h)
+        new.insert_obj(image)
+        new.gen_content()
+        out.save(str(dst))
+    finally:
+        source.close()
+        out.close()
+    return {"path": str(dst), "num_pages": 2}
+
+
+def run_parser(
+    spec: ParserSpec,
+    docs: list[dict],
+    work: Path,
+    device: str = "auto",
+    timeout: int = 3600,
+    warmup: bool = True,
+    log=print,
+) -> dict[str, DocOutput]:
+    python = ensure_env(spec, log=log)
+    # The runner's working directory is pdir, so every path handed to it is absolute.
+    pdir = work.resolve() / spec.name
+    pdir.mkdir(parents=True, exist_ok=True)
+    warmup_doc = make_warmup_pdf(Path(docs[0]["path"]), pdir / "warmup.pdf") if (warmup and docs) else None
+    manifest = {
+        "device": device,
+        "warmup_doc": warmup_doc,
+        "docs": [
+            {
+                "id": d["id"],
+                "path": d["path"],
+                "num_pages": d["num_pages"],
+                "out": str(pdir / f"{d['id']}.json"),
+            }
+            for d in docs
+        ],
+    }
+    for d in manifest["docs"]:
+        Path(d["out"]).unlink(missing_ok=True)
+    mpath = pdir / "manifest.json"
+    mpath.write_text(json.dumps(manifest, indent=2))
+    env = {**dict(spec.env), **os.environ, "PYTHONUNBUFFERED": "1", "TOKENIZERS_PARALLELISM": "false"}
+    env.pop("PYTHONPATH", None)
+    if device == "cpu":
+        env["CUDA_VISIBLE_DEVICES"] = ""
+    if device != "auto":
+        env["TORCH_DEVICE"] = device
+    log_path = pdir / "stderr.log"
+    started = time.perf_counter()
+    error = None
+    with open(log_path, "w") as lf:
+        proc = subprocess.Popen(
+            [str(python), str(runner_dir() / spec.runner), str(mpath)],
+            stdout=lf,
+            stderr=subprocess.STDOUT,
+            env=env,
+            cwd=pdir,
+            start_new_session=True,
+        )
+        try:
+            code = proc.wait(timeout=timeout)
+            if code != 0:
+                error = f"runner exited with code {code}"
+        except subprocess.TimeoutExpired:
+            error = f"timed out after {timeout}s"
+        finally:
+            stop_process_tree(proc)
+    wall = time.perf_counter() - started
+    results: dict[str, DocOutput] = {}
+    for d in manifest["docs"]:
+        out = Path(d["out"])
+        if out.exists():
+            raw = json.loads(out.read_text())
+            results[d["id"]] = DocOutput(
+                pages=raw.get("pages", []),
+                seconds=raw.get("seconds", 0.0),
+                error=raw.get("error"),
+                device=raw.get("device", ""),
+                warmup_seconds=raw.get("warmup_seconds"),
+                extra={"version": raw.get("version"), "wall_seconds": wall},
+            )
+        else:
+            tail = _tail(log_path)
+            results[d["id"]] = DocOutput(
+                pages=[], error=f"{error or 'no output'}: {tail}"[:1500], extra={"wall_seconds": wall}
+            )
+    return results
+
+
+def _descendants(pid: int) -> list[int]:
+    try:
+        out = subprocess.run(
+            ["ps", "-A", "-o", "pid=,ppid="], capture_output=True, text=True, timeout=10
+        ).stdout
+    except (OSError, subprocess.SubprocessError):
+        return []
+    children: dict[int, list[int]] = {}
+    for line in out.splitlines():
+        parts = line.split()
+        if len(parts) == 2 and parts[0].isdigit() and parts[1].isdigit():
+            children.setdefault(int(parts[1]), []).append(int(parts[0]))
+    found, stack = [], [pid]
+    while stack:
+        for child in children.get(stack.pop(), []):
+            found.append(child)
+            stack.append(child)
+    return found
+
+
+def stop_process_tree(proc: subprocess.Popen, grace: float = 5.0) -> None:
+    """Make sure neither the runner nor anything it started outlives it.
+
+    The runner kills its own helpers on a normal exit; this covers timeouts and
+    crashes: every descendant (even one in its own session) and the runner's
+    process group get SIGTERM, then SIGKILL.
+    """
+    running = proc.poll() is None
+    # Only signal PIDs we know are ours: after a normal exit the runner is reaped and
+    # its PID may be reused, so only its process group is signalled.
+    targets = [*_descendants(proc.pid), proc.pid] if running else []
+    for sig in (signal.SIGTERM, signal.SIGKILL):
+        for p in targets:
+            try:
+                os.kill(p, sig)
+            except (ProcessLookupError, PermissionError):
+                pass
+        try:
+            os.killpg(proc.pid, sig)
+        except (ProcessLookupError, PermissionError):
+            pass
+        try:
+            proc.wait(timeout=grace)
+        except subprocess.TimeoutExpired:
+            continue
+        if sig == signal.SIGTERM:
+            time.sleep(0.2)
+            targets = [p for p in targets if _pid_alive(p)]
+            if not targets:
+                break
+
+
+def _pid_alive(pid: int) -> bool:
+    try:
+        os.kill(pid, 0)
+        return True
+    except (ProcessLookupError, PermissionError):
+        return False
+
+
+def _tail(path: Path, lines: int = 6) -> str:
+    try:
+        text = path.read_text(errors="replace").strip().splitlines()
+    except OSError:
+        return ""
+    return " | ".join(text[-lines:])
+
+
+def save_outputs(path: Path, docs: list[dict], specs: list[ParserSpec], outputs, device: str) -> dict:
+    data = {
+        "meta": {
+            "tool": "parsearena",
+            "tool_version": __version__,
+            "created": datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"),
+            "hardware": hardware(),
+            "device_requested": device,
+        },
+        "parsers": [
+            {
+                "name": s.name,
+                "label": s.label,
+                "version": s.version,
+                "packages": list(s.packages),
+                "license": s.license,
+                "url": s.url,
+            }
+            for s in specs
+        ],
+        "docs": [{k: v for k, v in d.items() if k != "path"} for d in docs],
+        "outputs": {p: {d: asdict(o) for d, o in per.items()} for p, per in outputs.items()},
+    }
+    path.write_text(json.dumps(data, ensure_ascii=False, indent=1))
+    return data
+
+
+def load_outputs(data: dict) -> dict[str, dict[str, DocOutput]]:
+    return {p: {d: DocOutput(**o) for d, o in per.items()} for p, per in data["outputs"].items()}

parsearena/questions.py

@@ -0,0 +1,125 @@
+"""Question sets: hand-written YAML or drafted by an LLM from the PDF text layer."""
+
+from __future__ import annotations
+
+import json
+import re
+from pathlib import Path
+
+import yaml
+
+from .errors import BudgetExceeded, UserError
+from .evaluate import Question
+from .text import contains, normalize
+
+EXAMPLE = """questions:
+  - doc: report.pdf          # file name (or its stem) of one of the PDFs
+    q: What was the median household income in 2023?
+    answer: "80,610"          # a short string that must appear in the document
+  - doc: report.pdf
+    q: Which region had the highest median income?
+    answers: ["West", "the West"]   # any of these counts
+"""
+
+
+def load_questions(path: Path, docs: list[dict]) -> list[Question]:
+    if not path.exists():
+        raise UserError(f"Questions file not found: {path}", "Format:\n" + EXAMPLE)
+    try:
+        raw = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
+    except yaml.YAMLError as exc:
+        raise UserError(f"Could not parse {path}: {exc}", "Format:\n" + EXAMPLE) from exc
+    items = raw.get("questions") if isinstance(raw, dict) else raw
+    if not isinstance(items, list):
+        raise UserError(f"{path} has no `questions:` list.", "Format:\n" + EXAMPLE)
+    by_name = {}
+    for d in docs:
+        by_name[d["id"]] = d["id"]
+        by_name[d["file"]] = d["id"]
+        by_name[Path(d["file"]).stem] = d["id"]
+    out: list[Question] = []
+    for i, item in enumerate(items, start=1):
+        if not isinstance(item, dict):
+            raise UserError(f"Question {i} in {path} is not a mapping.", "Format:\n" + EXAMPLE)
+        doc = str(item.get("doc", ""))
+        if doc not in by_name:
+            raise UserError(
+                f"Question {i} refers to unknown doc {doc!r}.",
+                "Use the file name of one of the PDFs passed to run: " + ", ".join(d["file"] for d in docs),
+            )
+        answers = item.get("answers") or ([item["answer"]] if "answer" in item else [])
+        answers = [str(a) for a in answers if str(a).strip()]
+        if not item.get("q") or not answers:
+            raise UserError(f"Question {i} in {path} needs `q` and `answer` (or `answers`).")
+        out.append(
+            Question(id=str(item.get("id") or f"q{i}"), doc=by_name[doc], q=str(item["q"]), answers=answers)
+        )
+    return out
+
+
+def dump_questions(questions: list[Question], docs: list[dict], path: Path, header: str = "") -> None:
+    file_of = {d["id"]: d["file"] for d in docs}
+    data = {
+        "questions": [
+            {"id": q.id, "doc": file_of.get(q.doc, q.doc), "q": q.q, "answers": q.answers} for q in questions
+        ]
+    }
+    text = yaml.safe_dump(data, sort_keys=False, allow_unicode=True, width=100)
+    path.write_text((header + "\n" if header else "") + text, encoding="utf-8")
+
+
+SYSTEM = (
+    "You write retrieval test questions for a document. Each question must be answerable from a "
+    "single passage, and its answer must be a SHORT string (a number, a name, a few words) copied "
+    "verbatim from the text. Reply with JSON only."
+)
+
+PROMPT = """Text of the document "{name}" (may be truncated):
+<<<
+{text}
+>>>
+Write {n} questions about specific facts in this text (numbers in tables, definitions, names).
+Return a JSON list like [{{"q": "...", "answer": "..."}}].
+The answer must appear verbatim in the text above."""
+
+
+def _parse_json_list(reply: str) -> list[dict]:
+    reply = re.sub(r"<think>.*?</think>", "", reply, flags=re.S)
+    m = re.search(r"\[.*\]", reply, flags=re.S)
+    if not m:
+        return []
+    try:
+        data = json.loads(m.group(0))
+    except json.JSONDecodeError:
+        return []
+    return [d for d in data if isinstance(d, dict)]
+
+
+def generate_questions(llm, docs: list[dict], texts: dict[str, str], per_doc: int = 5, log=print):
+    """Draft questions per document; keep only those whose answer is in the source text.
+
+    Returns (questions, skipped_docs, budget_error). Stops early, keeping what it
+    has, if the next call would exceed the budget.
+    """
+    questions: list[Question] = []
+    skipped: list[str] = []
+    for d in docs:
+        text = texts.get(d["id"], "")
+        if len(text.strip()) < 200:
+            skipped.append(d["file"])
+            log(f"Skipping {d['file']}: no usable text layer to draft questions from (scanned?).")
+            continue
+        prompt = PROMPT.format(name=d["file"], text=text[:12000], n=per_doc)
+        try:
+            reply = llm.complete(SYSTEM, prompt, max_tokens=4000)
+        except BudgetExceeded as exc:
+            return questions, skipped, exc
+        norm = normalize(text)
+        kept = 0
+        for item in _parse_json_list(reply):
+            q, a = str(item.get("q", "")).strip(), str(item.get("answer", "")).strip()
+            if q and a and len(a) <= 80 and contains(norm, a):
+                kept += 1
+                questions.append(Question(id=f"{d['id']}-g{kept}", doc=d["id"], q=q, answers=[a]))
+        log(f"{d['file']}: kept {kept} questions")
+    return questions, skipped, None

parsearena/render.py

@@ -0,0 +1,40 @@
+"""Render PDF pages to images for the viewer (pypdfium2, no external tools)."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pypdfium2 as pdfium
+
+
+def render_pages(pdf_path: Path, out_dir: Path, width: int = 900) -> list[str]:
+    """Write one WebP per page and return their file names (relative to out_dir)."""
+    out_dir.mkdir(parents=True, exist_ok=True)
+    names = []
+    pdf = pdfium.PdfDocument(str(pdf_path))
+    try:
+        for i in range(len(pdf)):
+            page = pdf[i]
+            scale = width / max(page.get_width(), 1)
+            image = page.render(scale=scale).to_pil().convert("RGB")
+            name = f"{i + 1:03d}.webp"
+            image.save(out_dir / name, "WEBP", quality=80, method=4)
+            names.append(name)
+            page.close()
+    finally:
+        pdf.close()
+    return names
+
+
+def page_text(pdf_path: Path) -> list[str]:
+    """Text layer of every page (used only to draft questions)."""
+    pdf = pdfium.PdfDocument(str(pdf_path))
+    try:
+        out = []
+        for i in range(len(pdf)):
+            tp = pdf[i].get_textpage()
+            out.append(tp.get_text_range())
+            tp.close()
+        return out
+    finally:
+        pdf.close()

parsearena/report.py

@@ -0,0 +1,57 @@
+"""Static HTML viewer: one index.html plus rendered page images."""
+
+from __future__ import annotations
+
+import json
+import shutil
+from datetime import UTC, datetime
+from importlib import resources
+from pathlib import Path
+
+from .render import render_pages
+
+
+def build_report(
+    out_dir: Path,
+    recorded: dict,
+    pdf_paths: dict[str, Path],
+    questions: list,
+    evaluation: dict,
+    timings: dict,
+    is_recorded: bool,
+    doc_info: dict[str, dict] | None = None,
+) -> Path:
+    """Write ``out_dir/index.html`` and ``out_dir/pages/<doc>/NNN.webp``."""
+    out_dir.mkdir(parents=True, exist_ok=True)
+    pages_root = out_dir / "pages"
+    if pages_root.exists():
+        shutil.rmtree(pages_root)
+    docs = []
+    for d in recorded["docs"]:
+        names = render_pages(pdf_paths[d["id"]], pages_root / d["id"])
+        info = (doc_info or {}).get(d["id"], {})
+        docs.append({**d, **info, "images": [f"pages/{d['id']}/{n}" for n in names]})
+    evaluation = {**evaluation, "computed": datetime.now(UTC).strftime("%Y-%m-%d %H:%M UTC")}
+    data = {
+        "meta": {**recorded["meta"], "recorded": is_recorded},
+        "parsers": recorded["parsers"],
+        "docs": docs,
+        "outputs": {
+            p: {d: {"pages": o["pages"], "error": o["error"]} for d, o in per.items()}
+            for p, per in recorded["outputs"].items()
+        },
+        "questions": [{"id": q.id, "doc": q.doc, "q": q.q, "answers": q.answers} for q in questions],
+        "evaluation": evaluation,
+        "timings": timings,
+    }
+    payload = json.dumps(data, ensure_ascii=False, separators=(",", ":")).replace("</", "<\\/")
+    template = (resources.files("parsearena") / "report_template.html").read_text(encoding="utf-8")
+    index = out_dir / "index.html"
+    index.write_text(template.replace("__DATA__", payload), encoding="utf-8")
+    (out_dir / "results.json").write_text(
+        json.dumps(
+            {k: data[k] for k in ("meta", "parsers", "evaluation", "timings")}, indent=1, ensure_ascii=False
+        ),
+        encoding="utf-8",
+    )
+    return index