structured-data-extraction 0.1.0__py3-none-any.whl

structured_data_extraction/filter_records.py

@@ -0,0 +1,133 @@
+"""Filter extracted records by consensus confidence tier.
+
+After multi-model extraction and rating, each record carries a confidence score
+(fraction of models that agreed on it). This module lets the user keep only the
+tiers they trust — e.g. only records all models agreed on (1.0), or everything
+above majority (>= 0.67).
+
+Usable both as a CLI tool (structured-filter-records) and programmatically from
+the pipeline's end-of-run step. When no --levels are given, an interactive menu
+lets the user pick tiers; with a single extraction model every record is
+confidence 1.0, so the filter is a no-op.
+"""
+import argparse
+import json
+from collections import Counter
+from pathlib import Path
+from . import config
+
+RECORDS = Path(config.OUTPUT_DIR) / "records.json"
+_META = {"type", "generated_at", "provenance"}   # envelope keys that aren't extracted fields
+
+
+def load(path=RECORDS):
+    return json.loads(Path(path).read_text(encoding="utf-8"))
+
+
+def _records(data):
+    """Yield every extracted record across all documents (skipping envelope meta)."""
+    for env in data.values():
+        for key, recs in env.items():
+            if key not in _META and isinstance(recs, list):
+                yield from recs
+
+
+def counts(data):
+    """{confidence tier: record count}, highest tier first."""
+    tally = Counter(round(r["confidence"], 2) for r in _records(data)
+                    if r.get("confidence") is not None)
+    return dict(sorted(tally.items(), reverse=True))
+
+
+def tier_summary(data):
+    """One line per tier: the cumulative '>= tier' count (and its share of the total), plus the
+    count exactly at the tier. Returned as lines so a caller can print (CLI) or log (pipeline)."""
+    tiers = counts(data)
+    total = sum(tiers.values()) or 1
+    lines = []
+    for tier, at in tiers.items():
+        kept = sum(n for t, n in tiers.items() if t >= tier)
+        lines.append(f"  >= {tier:<4}  {kept:>4} kept ({kept / total:.0%})   {at:>4} at this tier")
+    return lines
+
+
+def keep(data, tiers):
+    """Return a subset with the same envelope shape, only records whose confidence is in `tiers`.
+    A field left with no records is dropped; a document left with no fields is omitted."""
+    chosen = {round(t, 2) for t in tiers}
+    out = {}
+    for stem, env in data.items():
+        kept = {}
+        for key, val in env.items():
+            if key in _META:
+                kept[key] = val
+            elif isinstance(val, list):
+                recs = [r for r in val
+                        if r.get("confidence") is not None and round(r["confidence"], 2) in chosen]
+                if recs:
+                    kept[key] = recs
+        if any(k not in _META for k in kept):
+            out[stem] = kept
+    return out
+
+
+def _pick(tiers):
+    """Interactive menu: choose which confidence tiers to keep."""
+    for i, tier in enumerate(tiers, 1):
+        print(f"  {i}. {tier}")
+    print("  a. all")
+    try:
+        raw = input("Select tiers (e.g. '1,3' or 'a'): ").strip().lower()
+    except EOFError:
+        return tiers   # no one on stdin (piped/non-interactive) → keep all rather than crash
+    if raw in ("", "a", "all"):
+        return tiers
+    return [tiers[int(x) - 1] for x in raw.replace(" ", "").split(",")
+            if x.isdigit() and 1 <= int(x) <= len(tiers)] or tiers
+
+
+def run(path=RECORDS, levels=None, json_out=None):
+    """Filter the records at `path` by confidence and print (or save) the subset.
+
+    Shared by the CLI and the pipeline's end-of-run step.
+    levels: tiers to keep (e.g. [1.0, 0.75]) or ["all"]; None => interactive menu.
+    """
+    data = load(path)
+    tiers = list(counts(data))          # highest first
+    if not tiers:
+        print(f"No scored records in {path}")
+        return
+    print("Confidence tiers (of all records):")
+    print("\n".join(tier_summary(data)))
+
+    if len(config.EXTRACTION_MODELS) < 2:
+        print("Only one extraction model — every record is confidence 1.0, nothing to filter.")
+        chosen = tiers
+    elif levels:
+        chosen = tiers if "all" in levels else [float(x) for x in levels]
+    else:
+        chosen = _pick(tiers)
+
+    out = keep(data, chosen)
+    payload = json.dumps(out, indent=2, ensure_ascii=False)
+    n = sum(1 for _ in _records(out))
+    if json_out:
+        Path(json_out).write_text(payload, encoding="utf-8")
+        print(f"wrote {n} record(s) across {len(out)} document(s) to {json_out}")
+    else:
+        print(payload)
+        print(f"\n{n} record(s) across {len(out)} document(s)")
+
+
+def main():
+    ap = argparse.ArgumentParser(description="Filter records by consensus confidence")
+    ap.add_argument("--levels", nargs="+", metavar="LVL",
+                    help="tiers to keep, e.g. 1.0 0.75 (or 'all'); omit for an interactive menu")
+    ap.add_argument("--file", default=str(RECORDS), help="records.json to read")
+    ap.add_argument("--json", help="write the kept records here instead of printing")
+    args = ap.parse_args()
+    run(args.file, args.levels, args.json)
+
+
+if __name__ == "__main__":
+    main()

structured_data_extraction/main.py

@@ -0,0 +1,329 @@
+"""Pipeline orchestrator : drives the full extraction workflow end to end.
+
+Processes one PDF or a directory of PDFs through every stage in order:
+  convert → vocab → classify → tables → embed → search → extract → rate → records
+
+Each stage writes its output to intermediate/ or output/ and is skipped on rerun
+if the artefact already exists (unless --force is set). Key-dependent caches
+(embeddings, matches, extractions) are automatically invalidated when the user
+changes their key list between runs.
+
+After extraction, finalize_outputs() fuses the per-model results via rater.py,
+writes provenance records, and optionally filters by confidence tier.
+"""
+import argparse
+import json
+import logging
+import sys
+import time
+from contextlib import contextmanager
+from importlib import resources
+from pathlib import Path
+from . import config, records, filter_records, matching, rater
+from .paths import inter, read_json
+from .convert import convert_pdf
+from .vocab import extract_vocab
+from .classify import classify_sections
+from .tables import process_tables, refine_matched_tables, refine_symbol_only_tables
+from .embed import embed_sections_and_keys
+from .search import search
+from .extract import extract_parallel
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
+logger = logging.getLogger("main")
+
+
+def _invalidate_stale_key_caches(stem, user_keys):
+    keys_path = inter(stem, "original_keys.json")
+    if not keys_path.exists():
+        return
+    if sorted(user_keys) == sorted(read_json(keys_path)):
+        return
+    logger.info(f"Keys changed; invalidating key-dependent caches for {stem}")
+    for name in ("key_embeddings.npy", "key_texts.json", "original_keys.json",
+                  "matches.json", "schema.json", "run.json"):
+        p = inter(stem, name)
+        if p.exists():
+            p.unlink()
+    for p in Path(config.OUTPUT_DIR).glob(f"{stem}_extraction_*.json"):
+        p.unlink()
+
+
+def process_single_pdf(pdf_path, user_keys):
+    """Run the full extraction pipeline on one PDF.
+
+    Executes every stage in order: convert → vocab → classify → tables →
+    embed → search → extract, writing intermediate artefacts to
+    INTERMEDIATE_DIR and per-model extractions to OUTPUT_DIR. Each stage is
+    skipped when its output already exists on disk (unless --force is set).
+    If the user's key list has changed since the last run, key-dependent
+    caches are automatically invalidated.
+
+    Args:
+        pdf_path: Path to the input PDF (str or Path).
+        user_keys: List of extraction keys, e.g. ["gain bandwidth", "PSRR"].
+                   Cleaned and validated internally (whitespace, quotes, and
+                   duplicates are handled; raises ValueError on unusable input).
+
+    Returns:
+        A list of dicts on success — one dict per extraction model. Each dict
+        maps a normalised field name to a list of extracted records::
+
+            [
+                {                                       # model 1
+                    "gain_bandwidth": [
+                        {"qualifier": "...", "condition": "...",
+                         "value": 4.0, "unit": "MHz",
+                         "metadata": {"source": "...", "section_type": "table",
+                                      "specified": True, "document_type": "...",
+                                      "model": "gemma4:31b-cloud"}},
+                        ...
+                    ],
+                    ...
+                },
+                { ... },                                # model 2
+            ]
+
+        Returns None if the PDF cannot be found or conversion fails.
+
+        These per-model dicts are the raw, un-fused extractions. To get the
+        final confidence-scored output, call finalize_outputs().
+    """
+    pdf_path = Path(pdf_path)
+    if not pdf_path.exists():
+        logger.error(f"PDF not found: {pdf_path}")
+        return None
+
+    user_keys = matching.validate_keys(user_keys)
+    stem = pdf_path.stem
+    int_dir = Path(config.INTERMEDIATE_DIR)
+    sections_path = int_dir / f"{stem}_sections.json"
+    tables_file = int_dir / f"{stem}_tables.json"
+    logger.info(f"=== Processing {pdf_path.name} ===")
+
+    timings = {}
+
+    @contextmanager
+    def timed(label):
+        t0 = time.perf_counter()
+        try:
+            yield
+        finally:
+            timings[label] = timings.get(label, 0.0) + time.perf_counter() - t0
+
+    with timed("convert"):
+        md_path = convert_pdf(str(pdf_path))
+    if md_path is None:
+        logger.error(f"Conversion failed for {pdf_path.name}")
+        return None
+
+    with timed("vocab"):
+        extract_vocab(md_path)
+    with timed("classify"):
+        classify_sections(stem)
+    with timed("tables"):
+        process_tables(str(sections_path))
+
+    # tables Docling stripped to bare symbols can't be retrieved, so they'd never reach the
+    # lazy-vision step below; de-garble them now so their recovered text gets embedded
+    if tables_file.exists():
+        with timed("pre-refine"):
+            refine_symbol_only_tables(str(tables_file), pdf_path=str(pdf_path))
+
+    _invalidate_stale_key_caches(stem, user_keys)
+    with timed("embed"):
+        embed_sections_and_keys(stem, user_keys=user_keys)
+
+    with timed("search"):
+        matches = search(stem, original_keys=user_keys)
+
+    # refine matched complex tables with vision model
+    matched_idxs = {
+        match.get("section_idx", 0)
+        for key_matches in matches.values()
+        for match in key_matches
+    }
+
+    if matched_idxs and tables_file.exists():
+        logger.info("Refining matched complex tables...")
+        with timed("refine"):
+            refine_matched_tables(str(tables_file), matched_idxs, pdf_path=str(pdf_path))
+
+    with timed("extract"):
+        results = extract_parallel(stem, user_keys)
+
+    _log_run_summary(stem, timings, results, matches)
+    logger.info(f"=== Done with {pdf_path.name} ===")
+    return results
+
+
+def _log_run_summary(stem, timings, results, matches):
+    """End-of-run legibility: one line per stage with its wall-clock time, then a record count
+    broken down by the section type each record came from (table/text/list)."""
+    by_source, total = {}, 0
+    for res in results:
+        for records in res.values():
+            for record in records:
+                total += 1
+                src = (record.get("metadata") or {}).get("section_type") or "unknown"
+                by_source[src] = by_source.get(src, 0) + 1
+    sources = ", ".join(f"{k}={v}" for k, v in sorted(by_source.items())) or "none"
+    matched = sum(1 for hits in matches.values() if hits)
+    stages = "".join(f"\n  {label:<11}{secs:6.1f}s" for label, secs in timings.items())
+    logger.info(
+        f"=== Run summary: {stem} ==={stages}"
+        f"\n  {'TOTAL':<11}{sum(timings.values()):6.1f}s"
+        f"\n  records: {total} ({sources}) from {len(results)} model(s); "
+        f"keys matched: {matched}/{len(matches)}"
+    )
+
+
+def finalize_outputs(levels=None):
+    """Fuse per-model extractions into the final deliverables.
+
+    Reads every output/<stem>_extraction_<model>.json written by extract_parallel,
+    clusters matching records across models into facts, and writes three files to
+    OUTPUT_DIR:
+
+      fused.json
+        The primary output. Structure: {stem: {field: [fact, ...]}}.
+        Each fact is a dict with the extracted fields (value, unit, qualifier,
+        condition, ...) plus:
+          - "confidence": float in (0, 1] — fraction of configured models that
+            independently produced this fact (e.g. 2 of 3 → 0.67). With a single
+            model every fact is 1.0.
+          - "models": list[str] — which models agreed on this fact.
+          - "metadata": {"section_type": str, "specified": bool}.
+        Facts are sorted highest-confidence first within each field.
+
+      rating_report.json
+        Inter-model agreement statistics: total fact count, full-agreement count,
+        pairwise Cohen's kappa and raw agreement for each model pair, and the list
+        of facts where models disagreed. Useful for auditing model consistency.
+
+      records.json
+        Provenance-stamped envelopes wrapping fused.json: one entry per document,
+        adding source_file, document_type, the models used, and the user's
+        original keys. This is the archival output.
+
+    After writing, if multiple extraction models are configured, the user is
+    offered (or given, via `levels`) a confidence filter to keep only trusted
+    tiers.
+
+    Args:
+        levels: Confidence tiers to keep (e.g. [1.0, 0.75]) or ["all"].
+                None triggers an interactive prompt (multi-model only).
+    """
+    by_file = rater.load_outputs()
+    if not by_file:
+        logger.warning("No extraction outputs found — skipping rating + records step.")
+        return
+    rater.rate(by_file)
+    out_path, n = records.write_records(config.OUTPUT_DIR)
+    logger.info(f"Wrote records for {n} document(s) → {out_path}")
+    _offer_filter(out_path, levels)
+
+
+def _offer_filter(records_path, levels):
+    """user filter records by confidence at the end of a run. If confidence came in as an argument
+    (--levels) uses them, otherwise offers the interactive menu. One model means every record
+    is confidence 1.0, so there's nothing to filter and we bail."""
+    if len(config.EXTRACTION_MODELS) < 2:
+        return
+    filter_records.run(records_path, levels=levels)
+
+
+def process_batch(pdf_dir, user_keys):
+    pdf_dir = Path(pdf_dir)
+    if not pdf_dir.exists():
+        logger.error(f"Directory not found: {pdf_dir}")
+        return []
+
+    pdfs = sorted(pdf_dir.glob("*.pdf"))
+    if not pdfs:
+        logger.error(f"No PDFs found in {pdf_dir}")
+        return []
+
+    logger.info(f"Found {len(pdfs)} PDFs")
+    results = []
+    for pdf in pdfs:
+        try:
+            result = process_single_pdf(str(pdf), user_keys)
+            if result:
+                results.append(result)
+        except Exception as e:
+            logger.error(f"Failed {pdf.name}: {e}")
+
+    logger.info(f"Processed {len(results)}/{len(pdfs)} files")
+    return results
+
+
+def _dump_templates(dest):
+    """Copy the packaged prompt templates to `dest` for editing, then the user points
+    TEMPLATES_DIR at that folder to use them. Only the built-in .j2 defaults are written;
+    read_text/write_text keeps this working even if the package is imported from a zip."""
+    src = resources.files("structured_data_extraction") / "templates"
+    dest = Path(dest)
+    dest.mkdir(parents=True, exist_ok=True)
+    for t in src.iterdir():
+        if t.name.endswith(".j2"):
+            (dest / t.name).write_text(t.read_text(encoding="utf-8"), encoding="utf-8")
+    return dest
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Document extraction pipeline")
+    parser.add_argument("path", nargs="?", help="Path to a PDF file or a directory of PDFs")
+    parser.add_argument("--keys", nargs="+", help="Keys to extract")
+    parser.add_argument("--output", default=config.OUTPUT_DIR)
+    parser.add_argument("--force", action="store_true", help="Recompute every step, ignoring cached files")
+    parser.add_argument("--levels", nargs="+", metavar="LVL",
+                        help="confidence tiers to keep (e.g. 1.0 0.75, or 'all'); "
+                             "omit to be prompted after extraction")
+    parser.add_argument("--dump-templates", metavar="DIR",
+                        help="Write the built-in prompt templates to DIR for editing, then exit; "
+                             "point TEMPLATES_DIR at that folder to use your edits")
+    args = parser.parse_args()
+
+    if args.dump_templates:
+        dest = _dump_templates(args.dump_templates)
+        print(f"Wrote templates to {dest}\nSet TEMPLATES_DIR={dest} to use your edits.")
+        return
+    if not args.path:
+        parser.error("a PDF file or directory is required (unless using --dump-templates)")
+
+    config.FORCE = args.force
+    user_keys = args.keys
+    if not user_keys:
+        print("Enter keys to extract (comma-separated):")
+        raw = input("> ").strip()
+        user_keys = [k.strip() for k in raw.split(",") if k.strip()]
+
+    try:
+        user_keys = matching.validate_keys(user_keys or [])
+    except ValueError as e:
+        print(f"Invalid keys: {e}")
+        sys.exit(1)
+
+    config.OUTPUT_DIR = args.output
+
+    if Path(args.path).is_dir():
+        results = process_batch(args.path, user_keys)
+        if results:
+            finalize_outputs(args.levels)
+        print(f"Extracted from {len(results)} files → {config.OUTPUT_DIR}")
+        if not results:
+            print("No documents produced any extraction.")
+            sys.exit(1)
+    else:
+        result = process_single_pdf(args.path, user_keys)
+        if result:
+            finalize_outputs(args.levels)
+            print(json.dumps(result, indent=2, ensure_ascii=False))  #document them well
+        else:
+            print("Extraction failed.")
+            sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

structured_data_extraction/matching.py

@@ -0,0 +1,153 @@
+"""Same-fact equivalence and field comparison.
+
+Decides whether two extracted records describe the same real-world fact so that
+rater.py can cluster them for consensus scoring. The rule: two records match
+when every non-metadata field either agrees or is empty on at least one side
+(empty = wildcard, so a sparse record still matches a fuller one).
+
+Numeric fields must match exactly (after float parsing). Unit fields match after
+Unicode/dash normalisation and synonym expansion (e.g. "millimeters" == "mm").
+Descriptor text (conditions, qualifiers) uses a rapidfuzz token-set similarity
+threshold of 85 — fuzzy because models phrase conditions differently
+("TA = −40°C" vs "T_A = -40°C"), but exact on the number so distinct specs are
+never merged.
+
+Also provides normalize_key (user key → valid Pydantic field name) and
+validate_keys (trust-boundary cleaning of user input).
+"""
+import logging
+import re
+import unicodedata
+from rapidfuzz import fuzz
+
+logger = logging.getLogger(__name__)
+
+FUZZ_THRESHOLD = 85
+
+MATCH, MISMATCH, GAP = "match", "mismatch", "gap"
+_DASHES = str.maketrans("–—−‐‑‒", "------")
+
+
+def _norm(value):
+    return unicodedata.normalize("NFKC", str(value)).translate(_DASHES).strip().casefold()
+
+
+def _empty(value):
+    return value is None or _norm(value) == ""
+
+
+def _as_number(value):
+    if isinstance(value, bool):
+        return None
+    if isinstance(value, (int, float)):
+        return float(value)
+    try:
+        return float(str(value).strip())
+    except ValueError:
+        return None
+
+
+_LEADING_NUM_RE = re.compile(r"^-?\d+(?:[.\-]\d+)*\s*")
+_UNIT_SYNONYMS = {"millimeters": "mm", "millimeter": "mm", "inches": "in", "inch": "in"}
+
+
+def _norm_unit(value):
+    # drop a leading value number so a baked-in "5 µM" matches a standalone unit "µM"
+    u = _LEADING_NUM_RE.sub("", _norm(value)).replace(" ", "").replace("μ", "u")
+    return _UNIT_SYNONYMS.get(u, u)
+
+
+_NUM_RE = re.compile(r"-?\d+(?:\.\d+)?")
+_NONALNUM_RE = re.compile(r"[^0-9a-z]+")
+
+
+def _numbers(value):
+    return frozenset(round(float(x), 6) for x in _NUM_RE.findall(str(value).translate(_DASHES)))
+
+_SYNONYMS = {"typical": "typ", "minimum": "min", "maximum": "max"}
+
+
+def _norm_text(value):
+    # strip underscores first so subscripts collapse ("R_L" == "RL"), then tokenise
+    s = unicodedata.normalize("NFKC", str(value)).translate(_DASHES).casefold().replace("_", "")
+    return " ".join(_SYNONYMS.get(t, t) for t in _NONALNUM_RE.sub(" ", s).split())
+
+
+def normalize_key(key):
+    # strip("_") so edge punctuation (a quote, a leading "+") can't produce a leading/trailing
+    # underscore — Pydantic rejects field names that start with "_".
+    return re.sub(r"[^a-z0-9_]", "_", unicodedata.normalize("NFKC", key).lower()).strip("_")
+
+
+def validate_keys(keys):
+    """Clean + validate user keys at the trust boundary, returning the cleaned key list.
+
+    Repairs the common cases (strips surrounding quotes/whitespace, drops blanks and exact
+    duplicates) and rejects what can't become a valid, unique Pydantic field name:
+    a key with no alphanumerics (empty field name) or two distinct keys that normalize to the
+    same field name (a silent collision that would otherwise drop one key's data). Raises
+    ValueError with a user-facing message on unrepairable input.
+    """
+    cleaned = []
+    for k in keys:
+        c = k.strip().strip("\"'").strip()
+        if not c:
+            if k.strip():
+                logger.info(f"Dropped key {k!r} (empty after cleaning)")
+            continue
+        if c != k:
+            logger.info(f"Cleaned key {k!r} -> {c!r}")
+        cleaned.append(c)
+
+    cleaned = list(dict.fromkeys(cleaned))  # drop exact duplicates, keep order
+    if not cleaned:
+        raise ValueError("no usable keys provided")
+
+    by_field = {}
+    for k in cleaned:
+        field = normalize_key(k)
+        if not field:
+            raise ValueError(f"key {k!r} has no letters or digits to form a field name — rename it")
+        if field in by_field:
+            raise ValueError(f"keys {by_field[field]!r} and {k!r} both map to field name "
+                             f"{field!r} — make them distinct")
+        by_field[field] = k
+    return cleaned
+
+
+def model_slug(model_name):
+    return re.sub(r"[^a-z0-9]+", "_", model_name.lower()).strip("_")
+
+
+def compare(name, a, b):
+    if _empty(a) or _empty(b):
+        return GAP
+
+    na, nb = _as_number(a), _as_number(b)
+    if na is not None and nb is not None:
+        return MATCH if na == nb else MISMATCH
+
+    if name == "unit" or name.endswith("_unit"):
+        return MATCH if _norm_unit(a) == _norm_unit(b) else MISMATCH
+
+    # one side a number, the other a string carrying exactly that number ("10 mL" vs 10)
+    if na is not None and _numbers(b) == {na}:
+        return MATCH
+    if nb is not None and _numbers(a) == {nb}:
+        return MATCH
+
+    numa, numb = _numbers(a), _numbers(b)
+    if not (numa <= numb or numb <= numa):
+        return MISMATCH
+    return MATCH if fuzz.token_set_ratio(_norm_text(a), _norm_text(b)) >= FUZZ_THRESHOLD else MISMATCH
+
+
+_META_FIELDS = {"metadata", "confidence", "models"}
+
+
+def same_fact(a, b):
+    fields = (set(a) | set(b)) - _META_FIELDS
+    for f in fields:
+        if compare(f, a.get(f), b.get(f)) == MISMATCH:
+            return False
+    return True

structured_data_extraction/paths.py

@@ -0,0 +1,34 @@
+"""Path helpers and a simple JSON cache.
+
+"""
+import json
+from pathlib import Path
+from . import config
+
+
+def inter(stem, name):
+    return Path(config.INTERMEDIATE_DIR) / f"{stem}_{name}"
+
+
+def out(stem, name):
+    return Path(config.OUTPUT_DIR) / f"{stem}_{name}"
+
+
+def read_json(path):
+    return json.loads(Path(path).read_text(encoding="utf-8"))
+
+
+def write_json(path, obj):
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(obj, indent=2, ensure_ascii=False), encoding="utf-8")
+    return obj
+
+
+def cached(path, build):
+    """Return the JSON at `path` if it exists, else `build()` then write + return it.
+    Recomputes regardless of an existing file when config.FORCE is set."""
+    path = Path(path)
+    if path.exists() and not config.FORCE:
+        return read_json(path)
+    return write_json(path, build())