Update app.py

app.py

@@ -1,7 +1,10 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 import joblib
+import pandas as pd
 import re
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.svm import LinearSVC
 from transformers import pipeline
 
 # Initialize FastAPI app
@@ -13,20 +16,29 @@ app = FastAPI(
     redoc_url="/redoc"
 )
 
-# Load pre-trained model
+# Load model and vectorizer
 model = joblib.load("model.joblib")
+vectorizer = joblib.load("vectorizer.joblib")
 
 # Initialize NER pipeline
 ner = pipeline('ner', model='Davlan/xlm-roberta-base-ner-hrl', grouped_entities=True)
 
-# Map NER entity labels to token names
+# 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 detection
+# Regex patterns for PII
 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'
@@ -35,27 +47,17 @@ 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
-
-# Updated PII Masking Function with NER and regex
+# 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
+        '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 masking
+    # NER-based masking
     entities = ner(text)
     for ent in entities:
         label = ent['entity_group']
@@ -65,7 +67,6 @@ def mask_and_store_all_pii(text):
             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
@@ -75,7 +76,7 @@ def mask_and_store_all_pii(text):
                     "entity": original
                 })
 
-    # Regex masking
+    # Regex-based masking
     regex_map = [
         (CARD_REGEX, 'credit_debit_no'),
         (AADHAAR_REGEX, 'aadhar_num'),
@@ -85,14 +86,12 @@ 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):
             original = match.group(0)
             token = f"[{token_name}_{counter[token_name]:03d}]"
-            start = match.start()
-            end = match.end()
             if original in text:
+                start = text.index(original)
                 text = text.replace(original, token, 1)
                 mapping[token] = original
                 counter[token_name] += 1
@@ -104,24 +103,19 @@ def mask_and_store_all_pii(text):
 
     return text, mapping, entity_list
 
-# Restore PII
-
+# 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
 
-# Classification Endpoint
+# Prediction endpoint
 @app.post("/classify")
 def classify_email(data: EmailInput):
     raw_text = data.input_email_body
-
-    # Masking
     masked_text, pii_map, entity_list = mask_and_store_all_pii(raw_text)
-
-    # Prediction
-    predicted_category = model.predict([masked_text])[0]
-
+    features = vectorizer.transform([masked_text])
+    predicted_category = model.predict(features)[0]
     return {
         "input_email_body": raw_text,
         "list_of_masked_entities": entity_list,
@@ -129,7 +123,33 @@ def classify_email(data: EmailInput):
         "category_of_the_email": predicted_category
     }
 
-# Health check endpoint
+# Retraining endpoint
+@app.post("/train")
+def train_model(new_example: TrainingExample):
+    df = pd.DataFrame([{"email_body": new_example.email_body, "label": new_example.label}])
+    try:
+        df.to_csv("training_data.csv", mode='a', header=not pd.io.common.file_exists("training_data.csv"), index=False)
+    except Exception as e:
+        return {"error": f"Failed to append to dataset: {str(e)}"}
+
+    # Load dataset
+    full_df = pd.read_csv("training_data.csv")
+    full_df['masked_text'] = full_df['email_body'].apply(lambda x: mask_and_store_all_pii(x)[0])
+
+    # Vectorize and train
+    new_vectorizer = TfidfVectorizer()
+    X = new_vectorizer.fit_transform(full_df['masked_text'])
+    y = full_df['label']
+    new_model = LinearSVC()
+    new_model.fit(X, y)
+
+    # Save updated model and vectorizer
+    joblib.dump(new_model, "model.joblib")
+    joblib.dump(new_vectorizer, "vectorizer.joblib")
+
+    return {"message": "Model retrained successfully with new example."}
+
+# Health check
 @app.get("/")
 def root():
     return {"message": "Email Classification API is running."}