Match drivers correctly across documents and keep highlights on the value

Mirrors GitHub commit ef18c8f: pack-level masking tokens, exact matching for masked driver names, provenance ranking by value coverage / heading / driver row, spaCy lg installed at build time.

Dockerfile

@@ -22,7 +22,7 @@ RUN apt-get update \
 
 COPY requirements.txt ./
 RUN pip install --no-cache-dir -r requirements.txt \
-    && python -m spacy download en_core_web_sm
+    && python -m spacy download en_core_web_lg
 
 COPY config/ ./config/
 COPY src/ ./src/

docs/hugging-face.md

@@ -8,7 +8,7 @@ The root `Dockerfile` does this:
 
 1. Builds the React UI with Vite.
 2. Installs the Python backend dependencies.
-3. Downloads the small spaCy English model used by Presidio.
+3. Downloads the spaCy English model (en_core_web_lg) used by Presidio.
 4. Copies `ui/dist` into the image.
 5. Starts FastAPI on port `7860`.
 6. Lets FastAPI serve both `/api/*` and the React app.

requirements.txt

@@ -8,7 +8,7 @@ pydantic>=2.7.0
 presidio-analyzer>=2.2.354
 presidio-anonymizer>=2.2.354
 # spaCy model (download separately after install):
-#   python -m spacy download en_core_web_lg
+#   python -m spacy download en_core_web_lg   (en_core_web_sm also works, with weaker name detection)
 spacy>=3.7.0
 
 # ── Utilities ──────────────────────────────────────────────────────────────

src/arbiter.py

@@ -166,6 +166,11 @@ def _check_conflict(
     return _first(sched_val, cert_val)  # schedule wins (default)
 
 
+def _is_masked(name: str) -> bool:
+    """True for a placeholder token produced by PIIMasker."""
+    return name.strip().startswith("MASKED_")
+
+
 def _find_matching_driver(name: str, candidates: list[Driver]) -> Driver | None:
     """
     Find the best-matching driver from *candidates* using fuzzy name matching.
@@ -174,6 +179,12 @@ def _find_matching_driver(name: str, candidates: list[Driver]) -> Driver | None:
     differences (e.g. "JOHN A SMITH" vs "SMITH JOHN") still match.
     Returns None when the best score is below ``_DRIVER_NAME_MATCH_THRESHOLD``.
     """
+    # Masked names (MASKED_PERSON_1, MASKED_PERSON_2, ...) differ by one
+    # character, so any two score ~93 on token_sort_ratio and would "match".
+    # Tokens are consistent across a pack, so they must match exactly.
+    if _is_masked(name):
+        return next((d for d in candidates if d.name.strip() == name.strip()), None)
+
     try:
         from rapidfuzz import fuzz as rfuzz
     except ImportError:
@@ -184,6 +195,8 @@ def _find_matching_driver(name: str, candidates: list[Driver]) -> Driver | None:
     best_score = 0
     best_driver: Driver | None = None
     for candidate in candidates:
+        if _is_masked(candidate.name):
+            continue
         score = rfuzz.token_sort_ratio(name.strip(), candidate.name.strip())
         if score > best_score:
             best_score = score

src/privacy.py

@@ -21,7 +21,6 @@ from typing import Optional
 from presidio_analyzer import AnalyzerEngine, RecognizerResult
 from presidio_analyzer.nlp_engine import NlpEngineProvider
 from presidio_anonymizer import AnonymizerEngine
-from presidio_anonymizer.entities import OperatorConfig
 
 from settings import settings
 
@@ -79,23 +78,33 @@ class PIIMasker:
         self._language = language
         self._score_threshold = score_threshold
 
-        # Build NLP engine (spaCy en_core_web_lg preferred; falls back to sm)
+        # Build NLP engine: spaCy en_core_web_lg if it is installed, otherwise
+        # en_core_web_sm. Checking first matters: Presidio silently downloads a
+        # missing model (400 MB for lg) instead of raising, so a try/except
+        # fallback never triggers.
         nlp_config = {
             "nlp_engine_name": "spacy",
-            "models": [{"lang_code": "en", "model_name": "en_core_web_lg"}],
+            "models": [{"lang_code": "en", "model_name": _pick_spacy_model()}],
         }
-        try:
-            provider = NlpEngineProvider(nlp_configuration=nlp_config)
-            nlp_engine = provider.create_engine()
-        except OSError:
-            # Fall back to the small model if lg is not installed
-            nlp_config["models"][0]["model_name"] = "en_core_web_sm"
-            provider = NlpEngineProvider(nlp_configuration=nlp_config)
-            nlp_engine = provider.create_engine()
+        provider = NlpEngineProvider(nlp_configuration=nlp_config)
+        nlp_engine = provider.create_engine()
 
         self._analyzer = AnalyzerEngine(nlp_engine=nlp_engine, supported_languages=[language])
         self._anonymizer = AnonymizerEngine()
 
+        # Token assignments persist across mask() calls, so every document in a
+        # pack masked by this instance gets the same token for the same value
+        # ("Alex Morgan" is MASKED_PERSON_1 in the Schedule and the Certificate).
+        # Without this, tokens were numbered per document and the arbiter's
+        # driver matching paired the wrong people. Call reset() between packs.
+        self._original_to_token: dict[tuple[str, str], str] = {}
+        self._counters: dict[str, int] = {}
+
+    def reset(self) -> None:
+        """Forget token assignments, e.g. before masking an unrelated pack."""
+        self._original_to_token.clear()
+        self._counters.clear()
+
     # ------------------------------------------------------------------
     # Public API
     # ------------------------------------------------------------------
@@ -125,31 +134,27 @@ class PIIMasker:
         if not results:
             return text, {}
 
-        # Build per-entity-type counters for unique token names
-        counters: dict[str, int] = {}
         token_map: dict[str, str] = {}
-        operators: dict[str, OperatorConfig] = {}
 
         # Sort by position so token numbering is left-to-right and deterministic
         results_sorted = sorted(results, key=lambda r: r.start)
 
-        # We need custom lambda operators to generate named tokens.
-        # Presidio's "replace" operator uses a fixed `new_value`; we work
-        # around this by building a value map keyed on (entity_type, original).
-        original_to_token: dict[tuple[str, str], str] = {}
-
+        # Presidio's "replace" operator uses a fixed `new_value`, so tokens are
+        # assigned here, keyed on (entity_type, normalised original) and shared
+        # across calls on this instance.
         for r in results_sorted:
             original = text[r.start : r.end]
-            key = (r.entity_type, original)
-            if key not in original_to_token:
-                counters[r.entity_type] = counters.get(r.entity_type, 0) + 1
-                token = f"{_TOKEN_PREFIX}{r.entity_type}_{counters[r.entity_type]}"
-                original_to_token[key] = token
-                token_map[token] = original
+            key = _token_key(r.entity_type, original)
+            token = self._original_to_token.get(key)
+            if token is None:
+                self._counters[r.entity_type] = self._counters.get(r.entity_type, 0) + 1
+                token = f"{_TOKEN_PREFIX}{r.entity_type}_{self._counters[r.entity_type]}"
+                self._original_to_token[key] = token
+            token_map.setdefault(token, original)
 
         # Perform replacement manually (Presidio replace operator doesn't
         # support per-occurrence dynamic values in a single pass).
-        masked_text = _replace_spans(text, results_sorted, original_to_token)
+        masked_text = _replace_spans(text, results_sorted, self._original_to_token)
         return masked_text, token_map
 
     def restore(self, text: str, token_map: dict[str, str]) -> str:
@@ -169,6 +174,23 @@ class PIIMasker:
 # ---------------------------------------------------------------------------
 
 
+def _pick_spacy_model() -> str:
+    """Return en_core_web_lg if installed, else en_core_web_sm."""
+    try:
+        import spacy.util
+
+        if spacy.util.is_package("en_core_web_lg"):
+            return "en_core_web_lg"
+    except ImportError:
+        pass
+    return "en_core_web_sm"
+
+
+def _token_key(entity_type: str, original: str) -> tuple[str, str]:
+    """Token lookup key: same entity type, case and runs of whitespace ignored."""
+    return entity_type, " ".join(original.split()).casefold()
+
+
 def _replace_spans(
     text: str,
     results: list[RecognizerResult],
@@ -181,6 +203,6 @@ def _replace_spans(
     chars = list(text)
     for r in sorted(results, key=lambda r: r.start, reverse=True):
         original = text[r.start : r.end]
-        token = original_to_token.get((r.entity_type, original), original)
+        token = original_to_token.get(_token_key(r.entity_type, original), original)
         chars[r.start : r.end] = list(token)
     return "".join(chars)

src/provenance.py

@@ -43,6 +43,9 @@ _SKIP_LEAF_NAMES = {
 # they don't contain verbatim PDF values so matching against them is meaningless.
 _SKIP_SECTION_NAMES = {"source_document", "field_citations"}
 
+# Docling labels that mark headings rather than values.
+_HEADING_LABELS = {"title", "section_header", "page_header", "page_footer"}
+
 # Document types whose corpora are unreliable for field-level matching.
 # Policy Booklets contain generic boilerplate — matching against them produces
 # false positives for almost every field ("Full", "UK", date digits, etc.).
@@ -69,6 +72,7 @@ class CorpusItem:
     page: int
     bbox: list[float]       # [x0%, y0%, x1%, y1%] — top-left origin, 0–100
     source_filename: str
+    label: str = ""         # Docling item label, e.g. "text", "section_header"
 
 
 class ProvenanceCorpus:
@@ -118,7 +122,7 @@ class ProvenanceCorpus:
             if not text or len(text) < 2:
                 continue
             for prov in getattr(item, "prov", []):
-                self._add_prov_item(prov, text, filename, page_sizes)
+                self._add_prov_item(prov, text, filename, page_sizes, _item_label(item))
 
     def _add_prov_item(
         self,
@@ -126,6 +130,7 @@ class ProvenanceCorpus:
         text: str,
         filename: str,
         page_sizes: dict[int, tuple[float, float]],
+        label: str = "",
     ) -> None:
         page_no = getattr(prov, "page_no", None)
         if page_no is None:
@@ -155,6 +160,7 @@ class ProvenanceCorpus:
             page=page_no,
             bbox=[round(x0, 3), round(y0, 3), round(x1, 3), round(y1, 3)],
             source_filename=filename,
+            label=label,
         ))
 
     @staticmethod
@@ -208,6 +214,17 @@ def _item_text(item: Any) -> str:
     return ""
 
 
+def _item_label(item: Any) -> str:
+    """Docling item label as a plain string ("" when unavailable)."""
+    label = getattr(item, "label", "")
+    return str(getattr(label, "value", label) or "")
+
+
+def _norm(text: str) -> str:
+    """Lowercase and collapse whitespace for containment checks."""
+    return " ".join(str(text).split()).lower()
+
+
 def _clamp(v: float) -> float:
     return max(0.0, min(100.0, v))
 
@@ -268,6 +285,7 @@ def build_provenance(
     # Track assigned positions to avoid two fields pointing to the same corpus item.
     # Key: (source_filename, page, x0, y0) — unpadded, original corpus position.
     used_positions: set[tuple] = set()
+    driver_values = _driver_values(record)
 
     for field_path, value_str in _walk_record(record):
         leaf = field_path.split(".")[-1].strip("[]0123456789")
@@ -291,10 +309,34 @@ def build_provenance(
         score_fn = rfuzz.partial_ratio if using_citation else rfuzz.WRatio
         threshold = _CITATION_THRESHOLD if using_citation else _MATCH_THRESHOLD
 
-        # Find best match, preferring positions not yet assigned to another field.
-        best_score = 0
+        # Rank candidates above the threshold by, in order:
+        #   1. does the item contain the extracted value (or its UK date form)?
+        #   2. is it ordinary text rather than a heading?
+        #   3. for driver fields, how many of the same driver's other values it
+        #      contains, so the field lands on that driver's row, not a label,
+        #      a section heading or another driver's row;
+        #   4. the fuzzy score.
+        # A short label inside the quote ("Name" in "Named Driver") can score as
+        # high as the real row, so the score alone is not enough.
+        value_norm = _norm(value_str)
+        uk_norm = _norm(_iso_to_uk_date(value_str) or "")
+        driver_idx = _driver_index(field_path)
+        sibling_values = [
+            v for v in driver_values.get(driver_idx, []) if v and v != value_norm
+        ] if driver_idx is not None else []
+
+        def rank(item: CorpusItem, score: float) -> tuple:
+            text_norm = _norm(item.text)
+            covers = (len(value_norm) >= _MIN_VALUE_LEN and value_norm in text_norm) or (
+                bool(uk_norm) and uk_norm in text_norm
+            )
+            not_heading = item.label not in _HEADING_LABELS
+            row_bonus = sum(1 for v in sibling_values if v in text_norm)
+            return (covers, not_heading, row_bonus, score)
+
+        best_key: tuple | None = None
         best_item: CorpusItem | None = None
-        best_unused_score = 0
+        best_unused_key: tuple | None = None
         best_unused_item: CorpusItem | None = None
 
         for corpus in trusted_corpora:
@@ -305,22 +347,24 @@ def build_provenance(
                     alt_score = rfuzz.partial_ratio(alt_search, item.text.lower())
                     if alt_score > score:
                         score = alt_score
+                if score < threshold:
+                    continue
+                key = rank(item, score)
                 pos_key = (item.source_filename, item.page, item.bbox[0], item.bbox[1])
-                if score > best_score:
-                    best_score = score
-                    best_item = item
-                if score > best_unused_score and pos_key not in used_positions:
-                    best_unused_score = score
-                    best_unused_item = item
-
-        # Prefer an unused position if it scores above threshold,
-        # otherwise fall back to best overall (may share a location).
-        if best_unused_item is not None and best_unused_score >= threshold:
-            chosen_item = best_unused_item
-            chosen_score = best_unused_score
-        elif best_item is not None and best_score >= threshold:
-            chosen_item = best_item
-            chosen_score = best_score
+                # Strictly greater: on a full tie the first item in reading order wins.
+                if best_key is None or key > best_key:
+                    best_key, best_item = key, item
+                if pos_key not in used_positions and (best_unused_key is None or key > best_unused_key):
+                    best_unused_key, best_unused_item = key, item
+
+        # Prefer an unused position, but only if it is as good on everything
+        # except the raw score. A driver's row is shared by several fields
+        # (date of birth, relationship, occupation), and an unused heading
+        # nearby must not win just because the row is already taken.
+        if best_unused_item is not None and best_unused_key[:3] >= best_key[:3]:
+            chosen_item, chosen_score = best_unused_item, best_unused_key[-1]
+        elif best_item is not None:
+            chosen_item, chosen_score = best_item, best_key[-1]
         else:
             continue
 
@@ -392,6 +436,34 @@ def _count_total_fields(record: Any) -> int:
     return sum(1 for _ in _walk_dict(data, ""))
 
 
+_DRIVER_PATH_RE = re.compile(r"^driver_details\[(\d+)\]\.")
+
+
+def _driver_index(field_path: str) -> int | None:
+    """Index of the driver a field belongs to, or None for non-driver fields."""
+    m = _DRIVER_PATH_RE.match(field_path)
+    return int(m.group(1)) if m else None
+
+
+def _driver_values(record: Any) -> dict[int, list[str]]:
+    """Normalised values per driver (UK date form for dates), used to keep a
+    driver's fields on that driver's row. Names and booleans are left out:
+    names are usually masked tokens and booleans match everywhere."""
+    out: dict[int, list[str]] = {}
+    for i, driver in enumerate(getattr(record, "driver_details", None) or []):
+        vals: list[str] = []
+        for key, val in driver.model_dump(exclude_none=True).items():
+            if key == "name" or isinstance(val, bool):
+                continue
+            text = str(val)
+            uk = _iso_to_uk_date(text)
+            norm = _norm(uk or text)
+            if len(norm) >= _MIN_VALUE_LEN:
+                vals.append(norm)
+        out[i] = vals
+    return out
+
+
 # ISO 8601 date/datetime patterns → UK DD/MM/YYYY
 _ISO_DATE_RE = re.compile(r'^(\d{4})-(\d{2})-(\d{2})')
 

tests/test_arbiter.py

@@ -368,3 +368,24 @@ class TestLiveRunRegressions:
         for c in conflicts:
             expected = c.schedule_value if c.winner == "schedule" else c.certificate_value
             assert str(_value_at(golden, c.field)) == expected, c.field
+
+
+class TestMaskedDriverNames:
+    """Masked names differ by one character, so they must match exactly."""
+
+    def test_different_tokens_do_not_match(self):
+        sched = _make_schedule(drivers=[{"name": "MASKED_PERSON_4", "is_main_driver": False}])
+        cert = _make_certificate(
+            drivers=[{"name": "MASKED_PERSON_1", "is_main_driver": True, "relationship": "Proposer"}]
+        )
+        golden, conflicts = PolicyArbiter().merge_records(sched, "s.pdf", cert, "c.pdf")
+        assert golden.driver_details[0].relationship is None
+        assert not any("is_main_driver" in c.field for c in conflicts)
+
+    def test_same_token_matches(self):
+        sched = _make_schedule(drivers=[{"name": "MASKED_PERSON_2", "is_main_driver": True}])
+        cert = _make_certificate(
+            drivers=[{"name": "MASKED_PERSON_2", "is_main_driver": True, "relationship": "Proposer"}]
+        )
+        golden, _ = PolicyArbiter().merge_records(sched, "s.pdf", cert, "c.pdf")
+        assert golden.driver_details[0].relationship == "Proposer"

tests/test_masking.py

@@ -0,0 +1,44 @@
+"""
+Masking must give the same person the same token in every document of a pack.
+
+Needs Presidio and a spaCy English model; skipped when they are not installed.
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import pytest
+
+sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
+
+pytest.importorskip("presidio_analyzer")
+spacy_util = pytest.importorskip("spacy.util")
+if not (spacy_util.is_package("en_core_web_sm") or spacy_util.is_package("en_core_web_lg")):
+    pytest.skip("no spaCy English model installed", allow_module_level=True)
+
+from privacy import PIIMasker  # noqa: E402
+
+
+def test_same_person_same_token_across_documents():
+    masker = PIIMasker()
+    schedule, s_map = masker.mask(
+        "Policyholder: Alex Morgan. Named drivers: Priya Shah and Jordan Reed."
+    )
+    certificate, c_map = masker.mask(
+        "Persons entitled to drive: Jordan Reed, Alex  Morgan, Priya Shah."
+    )
+    s_tokens = {v: k for k, v in s_map.items()}
+    c_tokens = {" ".join(v.split()): k for k, v in c_map.items()}
+    for name in ("Alex Morgan", "Priya Shah", "Jordan Reed"):
+        if name in s_tokens and name in c_tokens:
+            assert s_tokens[name] == c_tokens[name], name
+    assert "Alex Morgan" not in certificate
+
+
+def test_reset_starts_numbering_again():
+    masker = PIIMasker()
+    first, _ = masker.mask("My name is Alex Morgan.")
+    masker.reset()
+    second, _ = masker.mask("My name is Priya Shah.")
+    assert first.strip() == second.strip()

tests/test_provenance.py

@@ -0,0 +1,91 @@
+"""
+Tests for provenance matching (src/provenance.py).
+
+The corpus below mirrors the synthetic demo Schedule: a policyholder section on
+page 1 and a driver table on page 2 whose rows Docling returns as single text
+items. These cases used to land on headings, labels or another driver's row.
+"""
+from __future__ import annotations
+
+import sys
+from datetime import date
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
+
+from provenance import CorpusItem, ProvenanceCorpus, build_provenance  # noqa: E402
+from schema import Driver, PolicyHeader, UKMotorGoldenRecord  # noqa: E402
+
+_F = "Schedule of Insurance - Demo.pdf"
+
+
+def _corpus() -> ProvenanceCorpus:
+    corpus = ProvenanceCorpus(source_filename=_F, doc_type="Schedule")
+    rows = [
+        (1, "Policyholder details", "section_header"),
+        (1, "Policy number", "text"),
+        (1, "NBM-DEMO-0427", "text"),
+        (1, "Name", "text"),
+        (1, "Alex Morgan", "text"),
+        (1, "14/03/1991", "text"),
+        (1, "Product", "text"),
+        (2, "Driver details", "section_header"),
+        (2, "Driver name", "text"),
+        (2, "14/03/1991 Policyholder Product Manager Full Licence UK / 2/1 / No Yes", "text"),
+        (2, "07/08/1995 Named Driver Business Analyst UK Provisional / 1/4 / No No GBP 200.00", "text"),
+        (2, "11/10/1985 Named Driver Data Administrator Full Licence UK / 5/0 / No No", "text"),
+    ]
+    for i, (page, text, label) in enumerate(rows):
+        corpus.items.append(CorpusItem(
+            text=text, page=page, bbox=[10.0, float(i * 5), 60.0, float(i * 5 + 2)],
+            source_filename=_F, label=label,
+        ))
+    return corpus
+
+
+def _record() -> UKMotorGoldenRecord:
+    record = UKMotorGoldenRecord(
+        policy_header=PolicyHeader(policy_number="NBM-DEMO-0427"),
+        driver_details=[
+            Driver(name="MASKED_PERSON_1", dob=date(1991, 3, 14), relationship="Policyholder",
+                   occupation="Product Manager", license_type="Full UK", is_main_driver=True),
+            Driver(name="MASKED_PERSON_2", dob=date(1995, 8, 7), relationship="Named Driver",
+                   occupation="Business Analyst", license_type="UK Provisional"),
+            Driver(name="MASKED_PERSON_3", dob=date(1985, 10, 11), relationship="Named Driver",
+                   occupation="Data Administrator", license_type="Full UK"),
+        ],
+    )
+    record.field_citations = {
+        "driver_details[0].license_type": "Full Licence UK / 2/1 / No",
+        "driver_details[2].license_type": "Full Licence UK / 5/0 / No",
+    }
+    return record
+
+
+def _by_path() -> dict:
+    return {p.field_path: p for p in build_provenance(_record(), [_corpus()])}
+
+
+def test_policy_number_lands_on_value():
+    assert _by_path()["policy_header.policy_number"].matched_text == "NBM-DEMO-0427"
+
+
+def test_relationship_not_on_heading_or_label():
+    prov = _by_path()
+    assert prov["driver_details[0].relationship"].matched_text.startswith("14/03/1991 Policyholder")
+    assert prov["driver_details[1].relationship"].matched_text.startswith("07/08/1995")
+    assert prov["driver_details[2].relationship"].matched_text.startswith("11/10/1985")
+
+
+def test_occupation_not_on_fragment():
+    assert _by_path()["driver_details[0].occupation"].matched_text.startswith("14/03/1991 Policyholder")
+
+
+def test_repeated_value_lands_on_the_right_drivers_row():
+    prov = _by_path()
+    assert prov["driver_details[0].license_type"].matched_text.startswith("14/03/1991")
+    assert prov["driver_details[2].license_type"].matched_text.startswith("11/10/1985")
+
+
+def test_masked_names_are_not_located():
+    assert not any(p.endswith(".name") for p in _by_path())