add filter

inference.py

@@ -17,7 +17,6 @@ Quick start:
         "客戶王小明來電諮詢，身分證A123456789，手機0912345678"
     )
     print(result["masked_text"])
-    # 客戶王OO來電諮詢，身分證A1******89，手機0912******
 
 Environment variables (optional, for local development):
     PII_BASE_MODEL    Override base model id or path
@@ -61,7 +60,6 @@ SYSTEM_PROMPT = """你是一個台灣個人資料（PII）偵測專家。分析
 # JSON parsing (minimal defensive)
 # ============================================================
 def _extract_json(text):
-    """Pull the first JSON object out of the model output and parse it."""
     text = text.strip()
     start = text.find("{")
     end = text.rfind("}") + 1
@@ -73,13 +71,84 @@ def _extract_json(text):
         return None
 
 
+# ============================================================
+# Digit-character detection (halfwidth / fullwidth / Chinese numerals)
+# ============================================================
+_CHINESE_DIGITS = "零一二三四五六七八九"
+
+
+def _is_digit_char(c):
+    return c.isdigit() or c in _CHINESE_DIGITS
+
+
+# ============================================================
+# Filter 1: structural validation
+# ------------------------------------------------------------
+# Verify the entity's value matches its claimed type's structural pattern.
+# Filters out misclassifications produced by the model.
+# ============================================================
+def _is_valid_entity(entity):
+    etype = entity.get("type", "")
+    value = entity.get("value", "")
+    if not value:
+        return False
+
+    if etype == "NAME":
+        return 2 <= len(value) <= 4
+
+    if etype == "EMAIL":
+        if "@" not in value:
+            return False
+        _, _, domain = value.partition("@")
+        return "." in domain
+
+    if etype == "ADDRESS":
+        # System prompt requires "號" or "樓" for valid full address
+        return "號" in value or "樓" in value
+
+    raw = value.replace("-", "").replace(" ", "")
+
+    if etype == "ROC_ID":
+        if len(raw) != 10:
+            return False
+        return raw[0].isalpha() and all(_is_digit_char(c) for c in raw[1:])
+
+    if etype == "PHONE":
+        return len(raw) == 10 and all(_is_digit_char(c) for c in raw)
+
+    if etype == "NHI_ID":
+        return len(raw) == 12 and all(_is_digit_char(c) for c in raw)
+
+    if etype == "BANK_ACCOUNT":
+        return 10 <= len(raw) <= 16 and all(_is_digit_char(c) for c in raw)
+
+    if etype == "CREDIT_CARD":
+        return len(raw) == 16 and all(_is_digit_char(c) for c in raw)
+
+    return True
+
+
+# ============================================================
+# Filter 2: provenance check (guards against hallucination)
+# ------------------------------------------------------------
+# The value must verbatim appear in the input text. Tolerant of
+# whitespace/separator differences.
+# ============================================================
+def _value_in_text(value, text):
+    if not value:
+        return False
+    if value in text:
+        return True
+    norm = lambda s: s.replace(" ", "").replace("-", "").replace("　", "")
+    return norm(value) in norm(text)
+
+
 # ============================================================
 # Mask validation and correction
 # ============================================================
 def _validate_and_fix_mask(entity):
     """Recompute the masked field from the value field using deterministic
-    rules. Compensates for the model's occasional miscount in star quantity.
-    """
+    rules. Compensates for the model's occasional miscount in star quantity."""
     etype = entity.get("type", "")
     value = entity.get("value", "")
     masked = entity.get("masked", "")
@@ -118,7 +187,9 @@ def _validate_and_fix_mask(entity):
                 if idx >= 0 and idx > best_idx:
                     best_idx = idx
                     prefix = value[:idx + 1]
-        if prefix:
+        # Guard: prefix must be strictly shorter than value
+        # (avoids cases like "社區" where 區 is at the end)
+        if prefix and len(prefix) < len(value):
             correct_mask = prefix + "*" * (len(value) - len(prefix))
 
     elif etype in ("NHI_ID", "BANK_ACCOUNT", "CREDIT_CARD"):
@@ -158,14 +229,17 @@ class PIIDetector:
     def detect(self, text, fix_masks=True):
         """Detect PII in `text`.
 
-        Args:
-            text: input text (Chinese, fullwidth, or plain ASCII all work)
-            fix_masks: if True, recompute `masked` field using deterministic
-                rules to guarantee correct star count.
+        Pipeline:
+            1. Model inference
+            2. JSON parse
+            3. Filter 1: structural validation
+            4. Filter 2: provenance check (rejects hallucinations)
+            5. validate_and_fix_mask: deterministic mask recomputation
+            6. Recompute pii_found (False if all entities filtered out)
 
         Returns:
-            dict with keys: pii_found (bool), entities (list of {type, value,
-            masked, optionally mask_corrected}), raw_response (str), parse_error (bool).
+            dict: {"pii_found": bool, "entities": list, "raw_response": str,
+                   "parse_error": bool}
         """
         messages = [
             {"role": "system", "content": SYSTEM_PROMPT},
@@ -196,28 +270,30 @@ class PIIDetector:
                 "parse_error": True,
             }
 
-        if fix_masks and parsed.get("entities"):
-            parsed["entities"] = [_validate_and_fix_mask(e) for e in parsed["entities"]]
-
+        # ===== Two-stage filtering =====
+        entities = parsed.get("entities", [])
+        if entities:
+            # Filter 1: structural validation
+            entities = [e for e in entities if _is_valid_entity(e)]
+            # Filter 2: provenance check (rejects hallucinated/translated values)
+            entities = [e for e in entities if _value_in_text(e.get("value", ""), text)]
+            # Mask correction (only for entities passing both filters)
+            if fix_masks:
+                entities = [_validate_and_fix_mask(e) for e in entities]
+
+        parsed["entities"] = entities
+        parsed["pii_found"] = bool(entities)
         parsed["raw_response"] = raw_response
         parsed["parse_error"] = False
         return parsed
 
     def detect_and_replace(self, text):
         """Detect PII and return the input text with all detected PII replaced
-        by their masked forms.
-
-        Returns:
-            dict with keys: original (str), masked_text (str), entities (list),
-            pii_found (bool).
-        """
+        by their masked forms."""
         result = self.detect(text)
         masked_text = text
 
         if result.get("entities"):
-            # Sort by value length descending so longer strings are replaced
-            # first; otherwise a short value that is a substring of a longer
-            # one would corrupt the longer match.
             sorted_ents = sorted(
                 result["entities"],
                 key=lambda e: len(e.get("value", "")),
@@ -227,7 +303,7 @@ class PIIDetector:
                 value = e.get("value", "")
                 masked = e.get("masked", "")
                 if value and masked and value in masked_text:
-                    # Replace all occurrences (e.g., a name appearing twice).
+                    # Replace all occurrences
                     masked_text = masked_text.replace(value, masked)
 
         return {