Update app.py

app.py

@@ -2,6 +2,7 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 import joblib
 import re
+from transformers import pipeline
 
 # Initialize FastAPI app
 app = FastAPI(
@@ -15,56 +16,100 @@ app = FastAPI(
 # Load pre-trained model
 model = joblib.load("model.joblib")
 
+# Initialize NER pipeline
+ner = pipeline('ner', model='Davlan/xlm-roberta-base-ner-hrl', grouped_entities=True)
+
+# Map NER entity labels to token names
+NER_TO_TOKEN = {
+    'PER': 'full_name',
+    'EMAIL': 'email',
+    'DATE': 'dob'
+}
+
+# Regex patterns for PII detection
+EMAIL_REGEX = r'\b[\w\.-]+@[\w\.-]+\.\w{2,}\b'
+AADHAAR_REGEX = r'\b\d{4}\s?\d{4}\s?\d{4}\b'
+CARD_REGEX = r'\b(?:\d[ -]*?){13,19}\b'
+CVV_REGEX = r'(?i)\b(?:cvv[:\s\-]*)?(\d{3,4})\b'
+EXPIRY_REGEX = r'\b(0[1-9]|1[0-2])[\/\-]\d{2,4}\b'
+PHONE_REGEX = r'\+?\d[\d\s\-]{7,14}\d'
+DOB_REGEX = r'\b\d{1,2}[\/\-\.\s]\d{1,2}[\/\-\.\s]\d{2,4}\b'
+
 # Input schema
 class EmailInput(BaseModel):
     input_email_body: str
 
-# PII Masking Function
+# Updated PII Masking Function with NER and regex
 def mask_and_store_all_pii(text):
     text = str(text)
-    pii_map = {}
+    mapping = {}
+    counter = {
+        'full_name': 0,
+        'email': 0,
+        'phone_number': 0,
+        'dob': 0,
+        'aadhar_num': 0,
+        'credit_debit_no': 0,
+        'cvv_no': 0,
+        'expiry_no': 0
+    }
     entity_list = []
 
-    patterns = {
-        "email": r"\b[a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+\b",
-        "phone_number": r"(?<!\d)(\+?\d[\d\s\-]{7,14}\d)(?!\d)",
-        "dob": r"\b\d{1,2}[\/\-\.\s]\d{1,2}[\/\-\.\s]\d{2,4}\b",
-        "aadhar_num": r"\b\d{4}[ -]?\d{4}[ -]?\d{4}\b(?![\d])",
-        "credit_debit_no": r"\b(?:\d[ -]*?){13,19}\b",
-        "cvv_no": r"(?i)\b(?:CVV[:\s]*)?(\d{3,4})\b",
-        "expiry_no": r"\b(0[1-9]|1[0-2])[\/\-]\d{2,4}\b",
-    }
+    # NER masking
+    entities = ner(text)
+    for ent in entities:
+        label = ent['entity_group']
+        if label in NER_TO_TOKEN:
+            token_name = NER_TO_TOKEN[label]
+            original = ent['word'].replace('##', '')
+            token = f"[{token_name}_{counter[token_name]:03d}]"
+            if original in text:
+                start = text.index(original)
+                end = start + len(original)
+                text = text.replace(original, token, 1)
+                mapping[token] = original
+                counter[token_name] += 1
+                entity_list.append({
+                    "position": [start, start + len(token)],
+                    "classification": token_name,
+                    "entity": original
+                })
+
+    # Regex masking
+    regex_map = [
+        (CARD_REGEX, 'credit_debit_no'),
+        (AADHAAR_REGEX, 'aadhar_num'),
+        (PHONE_REGEX, 'phone_number'),
+        (CVV_REGEX, 'cvv_no'),
+        (EXPIRY_REGEX, 'expiry_no'),
+        (EMAIL_REGEX, 'email'),
+        (DOB_REGEX, 'dob')
+    ]
 
-    # Track masked spans to prevent overlapping matches
-    masked_spans = []
-
-    def is_overlapping(start, end):
-        return any(s <= start < e or s < end <= e for s, e in masked_spans)
-
-    for label, pattern in patterns.items():
-        for match in re.finditer(pattern, text):
-            original = match.group()
-            start, end = match.start(), match.end()
-            if is_overlapping(start, end):
-                continue
-            placeholder = f"[{label}_{len(pii_map):03d}]"
-            pii_map[placeholder] = original
-            entity_list.append({
-                "position": [start, end],
-                "classification": label,
-                "entity": original
-            })
-            text = text[:start] + placeholder + text[end:]
-            masked_spans.append((start, start + len(placeholder)))
-
-    return text, pii_map, entity_list
+    for regex, token_name in regex_map:
+        for match in re.finditer(regex, text):
+            original = match.group(0)
+            token = f"[{token_name}_{counter[token_name]:03d}]"
+            start = match.start()
+            end = match.end()
+            if original in text:
+                text = text.replace(original, token, 1)
+                mapping[token] = original
+                counter[token_name] += 1
+                entity_list.append({
+                    "position": [start, start + len(token)],
+                    "classification": token_name,
+                    "entity": original
+                })
+
+    return text, mapping, entity_list
 
 # Restore PII
+
 def restore_pii(masked_text, pii_map):
-    restored = masked_text
     for placeholder, original in pii_map.items():
-        restored = restored.replace(placeholder, original)
-    return restored
+        masked_text = masked_text.replace(placeholder, original)
+    return masked_text
 
 # Classification Endpoint
 @app.post("/classify")
@@ -77,7 +122,6 @@ def classify_email(data: EmailInput):
     # Prediction
     predicted_category = model.predict([masked_text])[0]
 
-    # Response format
     return {
         "input_email_body": raw_text,
         "list_of_masked_entities": entity_list,