Fix PII entity numbering order in reconstruction

🔧 Problem: Entity flags were numbered in reverse order (X_4, X_3, X_2, X_1)
✅ Solution: Now correctly numbered in appearance order (X_1, X_2, X_3, X_4)

Changes:
- Modified reconstruct_masked_text() in text_processing.py
- First pass: assign numbers in order of appearance
- Second pass: replace in reverse order to maintain text positions
- Added test script to verify correct numbering

This ensures PII entities are numbered intuitively based on their order in the text.

src/pii_masking/text_processing.py

@@ -178,17 +178,28 @@ def reconstruct_masked_text(text: str, pii_dict: Dict[str, List[str]]) -> str:
         Text with PII replaced by [ENTITY_TYPE_X] placeholders
     """
     masked_text = text
-    entity_counters = defaultdict(int)
     
-    # Get all spans and sort by position (reverse order to maintain positions)
+    # Get all spans and sort by position (normal order for numbering)
     all_spans = json_to_spans(text, pii_dict)
+    all_spans.sort(key=lambda x: x.start)
+    
+    # First pass: assign numbers to spans in order of appearance
+    entity_counters = defaultdict(int)
+    span_numbers = {}
+    
+    for i, span in enumerate(all_spans):
+        entity_counters[span.entity_type] += 1
+        # Create a unique key for this span
+        span_key = (span.start, span.end, span.entity_type)
+        span_numbers[span_key] = entity_counters[span.entity_type]
+    
+    # Second pass: replace spans in reverse order to maintain text positions
     all_spans.sort(key=lambda x: x.start, reverse=True)
     
-    # Replace each span with placeholder
     for span in all_spans:
-        entity_counters[span.entity_type] += 1
-        placeholder = f"[{span.entity_type}_{entity_counters[span.entity_type]}]"
-        
+        span_key = (span.start, span.end, span.entity_type)
+        number = span_numbers[span_key]
+        placeholder = f"[{span.entity_type}_{number}]"
         # Replace the span in the text
         masked_text = masked_text[:span.start] + placeholder + masked_text[span.end:]
     

static/index.html

@@ -451,7 +451,7 @@
     <div class="container">
         <div class="header">
             <h1>🔒 PII Masking Demo</h1>
-            <p>Detect and mask Personal Identifiable Information using AI</p>
+            <p>Detect and mask Personal Identifiable Information in your documents.</p>
         </div>
 
         <div class="content">
@@ -460,10 +460,10 @@
                     <label>Choose input method:</label>
                     <div class="input-tabs">
                         <button type="button" class="input-tab active" onclick="switchInputMethod('text')">
-                            📝 Text Input
+                            Text Input
                         </button>
                         <button type="button" class="input-tab" onclick="switchInputMethod('pdf')">
-                            📄 PDF Upload
+                            PDF Upload
                         </button>
                     </div>
                     

test_reconstruction.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+"""
+Test script to verify that PII reconstruction works with correct numbering.
+"""
+
+import sys
+from pathlib import Path
+
+# Add the src directory to Python path
+sys.path.insert(0, str(Path(__file__).parent / "src"))
+
+from pii_masking.text_processing import reconstruct_masked_text
+
+def test_reconstruction():
+    """Test the reconstruction function with multiple entities of the same type."""
+    
+    # Test case with multiple entities of the same type
+    text = "John T. Smith lives at 4872 Willow Creek Drive Apartment 14B Springfield, IL 62701. Phone: 555-1234. Email: john@email.com. On January 15, 2024 Ms. Alice A. Johnson from Records Office at 9135 Westfield Parkway, Suite 320 Hartford, CT 06101. Dear Ms. Johnson, I am writing regarding my claim from December 2024. As of February 1, 2024, my address is 4872 Willow Creek Drive, Apartment 14B, Springfield, IL 62701."
+    
+    # Simulate PII entities found in the text
+    pii_dict = {
+        "FIRSTNAME": ["John", "Alice"],
+        "LASTNAME": ["Smith", "Johnson"],
+        "CITY": ["Springfield", "Hartford"],
+        "STATE": ["IL", "CT"],
+        "ZIPCODE": ["62701", "06101"],
+        "PHONENUMBER": ["555-1234"],
+        "EMAIL": ["john@email.com"],
+        "DATE": ["January 15, 2024", "December 2024", "February 1, 2024"]
+    }
+    
+    print("🧪 Testing PII reconstruction with correct numbering")
+    print("=" * 60)
+    print(f"Original text: {text[:100]}...")
+    print()
+    
+    masked_text = reconstruct_masked_text(text, pii_dict)
+    
+    print("🎭 Masked text:")
+    print(masked_text)
+    print()
+    
+    # Check if numbering is correct (first occurrence should be _1, second _2, etc.)
+    print("🔍 Checking numbering order:")
+    
+    # Check FIRSTNAMEs
+    john_pos = masked_text.find("[FIRSTNAME_1]")
+    alice_pos = masked_text.find("[FIRSTNAME_2]")
+    print(f"  FIRSTNAME_1 position: {john_pos}")
+    print(f"  FIRSTNAME_2 position: {alice_pos}")
+    print(f"  ✅ Correct order: {john_pos < alice_pos}")
+    
+    # Check LASTNAMEs
+    smith_pos = masked_text.find("[LASTNAME_1]")
+    johnson_pos = masked_text.find("[LASTNAME_2]")
+    print(f"  LASTNAME_1 position: {smith_pos}")
+    print(f"  LASTNAME_2 position: {johnson_pos}")
+    print(f"  ✅ Correct order: {smith_pos < johnson_pos}")
+    
+    # Check CITYs
+    springfield_pos = masked_text.find("[CITY_1]")
+    hartford_pos = masked_text.find("[CITY_2]")
+    print(f"  CITY_1 position: {springfield_pos}")
+    print(f"  CITY_2 position: {hartford_pos}")
+    print(f"  ✅ Correct order: {springfield_pos < hartford_pos}")
+    
+    # Check DATEs
+    date1_pos = masked_text.find("[DATE_1]")
+    date2_pos = masked_text.find("[DATE_2]")
+    date3_pos = masked_text.find("[DATE_3]")
+    print(f"  DATE_1 position: {date1_pos}")
+    print(f"  DATE_2 position: {date2_pos}")
+    print(f"  DATE_3 position: {date3_pos}")
+    print(f"  ✅ Correct order: {date1_pos < date2_pos < date3_pos}")
+
+if __name__ == "__main__":
+    test_reconstruction()