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

structured_data_extraction/rater.py

@@ -0,0 +1,188 @@
+"""When multiple extraction models run on the same document they produce
+overlapping records. This module clusters those records into "facts" using the
+same-fact equivalence from matching.py, then scores each fact's confidence as
+the fraction of models that found it.
+
+Outputs:
+  - fused.json: every fact, grouped by document and key, sorted by confidence.
+    This is the single complete deliverable that downstream consumers read.
+  - rating_report.json: inter-rater agreement stats (Cohen's kappa and raw
+    agreement percentage for each model pair).
+"""
+from collections import defaultdict
+from pathlib import Path
+from sklearn.metrics import cohen_kappa_score
+from . import config
+from .paths import read_json, write_json
+from .matching import same_fact, _empty, _META_FIELDS, model_slug
+
+
+def _filled_fields(rec):
+    return sum(1 for k, v in rec.items() if k not in _META_FIELDS and not _empty(v))
+
+
+def fuse(records_by_model):
+    pool = [(m, r) for m, recs in records_by_model.items() for r in recs]
+    used = set()
+    clusters = []
+    for i, (model, rec) in enumerate(pool):
+        if i in used:
+            continue
+        used.add(i)
+        cluster = {model: rec}
+        for j in range(i + 1, len(pool)):
+            mj, rj = pool[j]
+            if j in used or mj in cluster:
+                continue
+            if same_fact(rec, rj):
+                used.add(j)
+                cluster[mj] = rj
+        clusters.append(cluster)
+    return clusters
+
+
+def _best_section(records: list[dict]) -> str:
+    """
+    Return the highest-priority section_type found across the records.
+    Priority is table > text > list; an unknown/missing type loses to all of them.
+
+    Arguments:
+    records: A list of records, each a dict that may hold a "metadata" dict
+        with a "section_type" key.
+
+    Returns:
+    The winning section_type as a string ("table", "text", "list", or "unknown").
+    """
+    types = [r.get("metadata", {}).get("section_type", "unknown") for r in records]
+    return min(types, key=lambda t: config.SECTION_TYPE_PRIORITY.get(t, 99))
+
+
+def _fused_metadata(records: list[dict]) -> dict:
+    """
+    Build the metadata for a fused record.
+
+    Arguments
+    records: The records in one cluster, each a dict that may hold a "metadata" dict.
+
+    Returns
+    A metadata dict with "section_type" and "specified" (defaulting to True).
+    """
+    rep = max(records, key=_filled_fields)
+    return {"section_type": _best_section(records),
+            "specified": rep.get("metadata", {}).get("specified", True)}
+
+
+def consensus(cluster, n_models):
+    """Collapse a cluster (model -> record) into one fact: field values taken from the
+    most-filled record, confidence = agreeing models / n_models, the sorted list of agreeing
+    models, and the fused metadata."""
+    recs = list(cluster.values())
+    rep = max(recs, key=_filled_fields)
+    fields = {k: v for k, v in rep.items() if k != "metadata"}
+    return {"confidence": round(len(cluster) / n_models, 2), "models": sorted(cluster),
+            **fields, "metadata": _fused_metadata(recs)}
+
+
+def load_outputs():
+    by_file = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
+    slug_to_model = {model_slug(m): m for m in config.EXTRACTION_MODELS}
+    for path in sorted(Path(config.OUTPUT_DIR).glob("*_extraction_*.json")):
+        stem, _, slug = path.stem.rpartition("_extraction_")
+        model = slug_to_model.get(slug)
+        if model is None:
+            continue
+        data = read_json(path)
+
+        for field, records in data.items():
+            if isinstance(records, list):
+                by_file[stem][field][model].extend(records)
+    return by_file
+
+
+def fuse_outputs(by_file, models):
+    n = len(models)
+    facts = []
+    for stem, fields in sorted(by_file.items()):
+        for field, recs_by_model in sorted(fields.items()):
+            clusters = fuse({m: recs_by_model.get(m, []) for m in models})
+            for c in clusters:
+                facts.append({"source_file": f"{stem}.pdf", "field": field, **consensus(c, n)})
+    return facts
+
+
+def _group_by_doc(facts):
+    grouped = defaultdict(lambda: defaultdict(list))
+    for f in facts:
+        stem = Path(f["source_file"]).stem
+        grouped[stem][f["field"]].append({k: v for k, v in f.items()
+                                          if k not in ("source_file", "field")})
+    # highest-confidence facts first within each key; stable, so ties keep fuse order
+    for fields in grouped.values():
+        for recs in fields.values():
+            recs.sort(key=lambda r: r.get("confidence", 0), reverse=True)
+    return {stem: dict(fields) for stem, fields in grouped.items()}
+
+
+def majority_cutoff(n):
+    return n // 2 + 1
+
+
+def rate(by_file, models=None): 
+    models = config.EXTRACTION_MODELS if models is None else models
+    n = len(models)
+    Path(config.OUTPUT_DIR).mkdir(parents=True, exist_ok=True)
+    all_facts = fuse_outputs(by_file, models)
+
+    # fused.json = the end result 
+    fused_path = Path(config.OUTPUT_DIR) / "fused.json"
+    write_json(fused_path, _group_by_doc(all_facts))
+
+    rows = agreement(all_facts, models)
+    report = {
+        "models": models,
+        "facts": len(all_facts),
+        "full_agreement": sum(1 for f in all_facts if len(f["models"]) == n),
+        "agreement": rows,
+        "partial_agreement": [f for f in all_facts if len(f["models"]) < n],
+    }
+    out = Path(config.OUTPUT_DIR) / "rating_report.json"
+    write_json(out, report)
+    _print_agreement(len(all_facts), rows)
+    print(f"\nwrote {out}\nwrote {fused_path}")
+
+
+def agreement(facts, models):
+    if not facts:
+        return []
+    present = {m: [1 if m in f["models"] else 0 for f in facts] for m in models}
+    rows = []
+    for i, a in enumerate(models):
+        for b in models[i + 1:]:
+            agree = sum(x == y for x, y in zip(present[a], present[b])) / len(facts)
+            if len(set(present[a] + present[b])) < 2:
+                # cohen_kappa_score requires at least two label classes; skip if all labels are identical
+                kappa = None
+            else:
+                kappa = round(cohen_kappa_score(present[a], present[b], labels=[0, 1]), 3)
+            rows.append({"raters": [a, b], "kappa": kappa, "agreement": round(agree, 4)})
+    return rows
+
+
+def _print_agreement(n_facts, rows):
+    print(f"\ninter-rater agreement across {n_facts} facts (all files):")
+    for row in rows:
+        a, b = row["raters"]
+        kappa = "n/a" if row["kappa"] is None else row["kappa"]
+        print(f"  {a} <-> {b}: kappa={kappa}, agreement={row['agreement']:.1%}")
+
+
+def main():
+    by_file = load_outputs()
+    if not by_file:
+        print(f"No extraction outputs found in {config.OUTPUT_DIR}")
+        return
+    rate(by_file)
+
+
+if __name__ == "__main__":
+    main()

structured_data_extraction/records.py

@@ -0,0 +1,71 @@
+"""Provenance-stamped record envelopes.
+
+Wraps each document's fused facts (from fused.json) in a metadata envelope
+that records the source file, document type, which models were used (schema,
+vision, embedding), the user's original keys, and the normalised key mapping.
+The result is records.json.
+"""
+import os
+from pathlib import Path
+from . import config, matching
+from .paths import inter, read_json, write_json
+
+FUSED = Path(config.OUTPUT_DIR) / "fused.json"
+
+
+def _run(stem):
+    """Load the run stamp extract.py wrote (extracted_at, user_keys) for the envelope's
+    provenance. Returns {} when it's missing so make_record falls back to its defaults."""
+    path = inter(stem, "run.json")
+    return read_json(path) if path.exists() else {}
+
+
+def _document_type(stem):
+    path = inter(stem, "schema.json")
+    return read_json(path).get("document_type", "unknown") if path.exists() else "unknown"
+
+
+def make_record(stem, by_key):
+    """Wrap one document's fused facts in a metadata envelope: a type tag, a timestamp, and
+    provenance (source file, document type, the models used, and the user's keys)."""
+    run = _run(stem)
+    return {
+        "type": "ExtractionResult",
+        "generated_at": run.get("extracted_at"),
+        "provenance": {
+            "source_file": f"{stem}.pdf",
+            "document_type": _document_type(stem),
+            "models": {
+                "schema_discovery": os.environ["SCHEMA_MODEL"],
+                "vision": os.environ["VISION_MODEL"],
+                "embedding": os.environ["EMBEDDING_MODEL"],
+            },
+            "user_keys": run.get("user_keys", []),
+            "key_mapping": {matching.normalize_key(k): k for k in run.get("user_keys", [])},
+        },
+        **by_key,
+    }
+
+
+def build_records(path=FUSED):
+    report = read_json(path)
+    return {stem: make_record(stem, by_key) for stem, by_key in report.items()}
+
+
+def write_records(output_dir):
+    """Wrap every document's fused facts in a provenance envelope and write records.json.
+    Keeps all records; consumers filter by the `confidence` field (see filter_records.py)."""
+    output_dir = Path(output_dir)
+    envelopes = build_records(output_dir / "fused.json")
+    out_path = output_dir / "records.json"
+    write_json(out_path, envelopes)
+    return out_path, len(envelopes)
+
+
+def main():
+    out_path, n = write_records(config.OUTPUT_DIR)
+    print(f"wrote {out_path} ({n} documents)")
+
+
+if __name__ == "__main__":
+    main()

structured_data_extraction/search.py

@@ -0,0 +1,110 @@
+"""Semantic search : match user keys to document sections and table rows.
+
+Computes cosine similarity between the key embeddings and the section/row
+embeddings produced by embed.py, then returns the top-K hits above
+SIMILARITY_THRESHOLD for each key.
+
+The embedding matrix has sections first, then rows; a hit index
+>= num_sections is a row-level match and carries the originating section_idx,
+row_text, and group_header so downstream extraction knows which table and row
+the hit came from.
+
+Results are cached to matches.json and reused on reruns unless --force is set.
+"""
+import argparse
+import logging
+import numpy as np
+from . import config
+from .paths import inter, read_json, cached
+
+logger = logging.getLogger(__name__)
+
+
+def cosine_similarity(a, b):
+    # floor the norm so a zero-norm vector stays zero instead of 0/0 = nan, which would
+    a_norm = a / np.maximum(np.linalg.norm(a, axis=1, keepdims=True), 1e-12)
+    b_norm = b / np.maximum(np.linalg.norm(b, axis=1, keepdims=True), 1e-12)
+    return a_norm @ b_norm.T
+
+
+def _build_match_dict(idx, score, sec_texts, sections, num_secs, rows):
+    preview = sec_texts[idx][:150] + "..." if len(sec_texts[idx]) > 150 else sec_texts[idx]
+
+    is_row = idx >= num_secs and len(rows) > 0
+    if is_row:
+        row_idx = idx - num_secs
+        row_info = rows[row_idx] if row_idx < len(rows) else {}
+        sec_idx = row_info.get("section_idx", idx)
+    else:
+        sec_idx = idx
+
+    match = {
+        "section_idx": int(sec_idx),
+        "similarity": round(score, 4),
+        "content_preview": preview,
+    }
+    if is_row:
+        match.update({
+            "is_row": True,
+            "row_text": row_info.get("row_text", ""),
+            "row_number": row_info.get("row_number"),
+            "group_header": row_info.get("group_header", ""),
+        })
+    if sec_idx < len(sections):
+        match["type"] = sections[sec_idx].get("type", "unknown")
+
+    return match
+
+
+def _select_hits_for_key(sims, sec_texts, sections, num_secs, rows, threshold):
+    key_matches = []
+    for idx in np.argsort(sims)[::-1]:
+        score = float(sims[idx])
+        if score < threshold or len(key_matches) >= config.SEARCH_TOP_K:
+            break
+        key_matches.append(_build_match_dict(idx, score, sec_texts, sections, num_secs, rows))
+    return key_matches
+
+
+def search(stem, original_keys=None, threshold=None):
+    threshold = config.SIMILARITY_THRESHOLD if threshold is None else threshold
+
+    def build():
+        key_embs = np.load(inter(stem, "key_embeddings.npy"))
+        if key_embs.size == 0:
+            logger.warning("No key embeddings. Nothing to search.")
+            return {}
+        sec_embs = np.load(inter(stem, "section_embeddings.npy"))
+        sec_texts = read_json(inter(stem, "section_texts.json"))
+        sections = read_json(inter(stem, "sections.json"))
+        key_texts = read_json(inter(stem, "key_texts.json"))
+        row_map = read_json(inter(stem, "row_map.json"))
+        # num_secs is the boundary: idx >= num_secs means a row hit
+        num_secs = row_map.get("num_sections", len(sections))
+        rows = row_map.get("rows", [])
+        keys = original_keys if original_keys is not None else read_json(inter(stem, "original_keys.json"))
+
+        sims = cosine_similarity(key_embs, sec_embs)
+        matches = {}
+        for i, key_text in enumerate(key_texts):
+            orig_key = keys[i] if i < len(keys) else key_text
+            matches[orig_key] = _select_hits_for_key(
+                sims[i], sec_texts, sections, num_secs, rows, threshold)
+            logger.info(f"Key '{orig_key}': {len(matches[orig_key])} matches above {threshold}")
+        return matches
+
+    return cached(inter(stem, "matches.json"), build)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Semantic search")
+    parser.add_argument("--stem", required=True)
+    parser.add_argument("--threshold", type=float)
+    parser.add_argument("--original-keys", nargs="+")
+    args = parser.parse_args()
+    result = search(args.stem, original_keys=args.original_keys, threshold=args.threshold)
+    print(f"Found matches for {len(result)} keys")
+
+
+if __name__ == "__main__":
+    main()