Update app.py

app.py

@@ -1,7 +1,6 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 import joblib
-import pandas as pd
 import re
 from transformers import pipeline
 
@@ -14,28 +13,13 @@ app = FastAPI(
     redoc_url="/redoc"
 )
 
-# Load the combined model pipeline (includes vectorizer)
+# Load pre-trained model
 model = joblib.load("model.joblib")
 
-# Initialize NER pipeline
+# Initialize NER pipeline (multilingual)
 ner = pipeline('ner', model='Davlan/xlm-roberta-base-ner-hrl', grouped_entities=True)
 
-# Input schemas
-class EmailInput(BaseModel):
-    input_email_body: str
-
-class TrainingExample(BaseModel):
-    email_body: str
-    label: str
-
-# Map NER labels to types
-NER_TO_TOKEN = {
-    'PER': 'full_name',
-    'EMAIL': 'email',
-    'DATE': 'dob'
-}
-
-# Regex patterns for PII
+# 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'
@@ -44,17 +28,28 @@ 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'
 
-# Masking function
-def mask_and_store_all_pii(text):
-    text = str(text)
-    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 = []
+NER_TO_TOKEN = {
+    'PER': 'full_name',
+    'EMAIL': 'email',
+    'DATE': 'dob'
+}
 
-    # NER-based masking
+def mask_pii(text, mapping=None, counter=None):
+    if mapping is None:
+        mapping = {}
+    if counter is None:
+        counter = {
+            'full_name': 0,
+            'email': 0,
+            'phone_number': 0,
+            'dob': 0,
+            'aadhar_num': 0,
+            'credit_debit_no': 0,
+            'cvv_no': 0,
+            'expiry_no': 0
+        }
+
+    # Mask NER entities first
     entities = ner(text)
     for ent in entities:
         label = ent['entity_group']
@@ -63,17 +58,11 @@ def mask_and_store_all_pii(text):
             original = ent['word'].replace('##', '')
             token = f"[{token_name}_{counter[token_name]:03d}]"
             if original in text:
-                start = text.index(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-based masking
+    # Mask regex patterns
     regex_map = [
         (CARD_REGEX, 'credit_debit_no'),
         (AADHAAR_REGEX, 'aadhar_num'),
@@ -83,35 +72,37 @@ def mask_and_store_all_pii(text):
         (EMAIL_REGEX, 'email'),
         (DOB_REGEX, 'dob')
     ]
+
     for regex, token_name in regex_map:
-        for match in re.finditer(regex, text):
+        def replacer(match):
             original = match.group(0)
             token = f"[{token_name}_{counter[token_name]:03d}]"
-            if original in text:
-                start = text.index(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
-                })
+            counter[token_name] += 1
+            mapping[token] = original
+            return token
+        text = re.sub(regex, replacer, text)
 
-    return text, mapping, entity_list
+    return text, mapping
 
-# Restore PII (optional use)
-def restore_pii(masked_text, pii_map):
-    for placeholder, original in pii_map.items():
-        masked_text = masked_text.replace(placeholder, original)
-    return masked_text
+# Input schema
+class EmailInput(BaseModel):
+    input_email_body: str
 
-# Prediction endpoint
+# Classification Endpoint
 @app.post("/classify")
 def classify_email(data: EmailInput):
     raw_text = data.input_email_body
-    masked_text, pii_map, entity_list = mask_and_store_all_pii(raw_text)
+
+    # Masking using your advanced function
+    masked_text, pii_map = mask_pii(raw_text)
+
+    # Convert pii_map to a list for easier frontend use (optional)
+    entity_list = [{"placeholder": k, "original": v} for k, v in pii_map.items()]
+
+    # Prediction
     predicted_category = model.predict([masked_text])[0]
+
+    # Response format
     return {
         "input_email_body": raw_text,
         "list_of_masked_entities": entity_list,
@@ -119,7 +110,7 @@ def classify_email(data: EmailInput):
         "category_of_the_email": predicted_category
     }
 
-# Health check
+# Health check endpoint
 @app.get("/")
 def root():
     return {"message": "Email Classification API is running."}