Upload 5 files

api.py

@@ -0,0 +1,54 @@
+from flask import Blueprint, request, jsonify
+from transformers import DebertaV2Tokenizer, DebertaV2ForSequenceClassification
+import torch
+from utils import mask_pii
+
+api_bp = Blueprint("api", __name__)
+
+# Repo of Hugging Face Model Hub where Model is Pushed 
+REPO_ID = "Nikpatil/Email_classifier"
+MAX_LENGTH = 256
+
+tokenizer = DebertaV2Tokenizer.from_pretrained(REPO_ID)
+model = DebertaV2ForSequenceClassification.from_pretrained(REPO_ID)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+model.eval()
+
+id2label = {0: "Incident", 1: "Request", 2: "Problem", 3: "Change"}
+
+@api_bp.route("/classify", methods=["POST"])
+def classify_email():
+    data = request.get_json()
+    email_body = data.get("email_body", "")
+
+    if not email_body:
+        return jsonify({"Error": "Email body field is required"}), 400
+
+    masked_email, entities = mask_pii(email_body)
+
+    inputs = tokenizer(
+        masked_email,
+        add_special_tokens=True,
+        max_length=MAX_LENGTH,
+        padding='max_length',
+        truncation=True,
+        return_tensors='pt'
+    )
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
+    predicted_class_id = torch.argmax(probs).item()
+    predicted_class = id2label[predicted_class_id]
+
+    return jsonify({
+        "input_email_body": email_body,
+        "list_of_masked_entities": entities,
+        "masked_email": masked_email,
+        "category_of_the_email": predicted_class
+    }), 200
+
+

app.py

@@ -0,0 +1,8 @@
+from flask import Flask
+from api import api_bp
+
+app = Flask(__name__)
+app.register_blueprint(api_bp)
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=7860)

model.py

@@ -0,0 +1,265 @@
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+from sklearn.metrics import classification_report, confusion_matrix
+from sklearn.model_selection import train_test_split
+from transformers import (
+    DebertaV2Tokenizer,
+    DebertaV2ForSequenceClassification,
+    TrainingArguments,
+    Trainer,
+    EarlyStoppingCallback
+)
+import matplotlib.pyplot as plt
+import seaborn as sns
+from torch.utils.data import Dataset
+
+# Set seed for reproducibility
+SEED = 42
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+# ------------------------ Data Preprocessing ------------------------
+
+class EmailClassification(Dataset):
+    def __init__(self, texts, labels, tokenizer, max_length):
+        """Initialize dataset with texts, labels and tokenizer settings."""
+        self.texts = texts
+        self.labels = labels
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+
+    def __len__(self):
+        """Return the number of sample in the dataset."""
+        return len(self.texts)
+
+    def __getitem__(self, idx):
+        """Get a single item from the dataset."""
+        text = str(self.texts[idx])
+        label = self.labels[idx]
+
+        encoding = self.tokenizer(
+            text,
+            add_special_tokens=True,
+            max_length=self.max_length,
+            truncation=True,
+            return_attention_mask=True,
+            return_tensors='pt',
+            padding='max_length'
+        )
+
+        return {
+            'input_ids': encoding['input_ids'].flatten(),
+            'attention_mask': encoding['attention_mask'].flatten(),
+            'labels': torch.tensor(label, dtype=torch.long)
+        }
+
+def compute_class_weights(labels):
+    """Compute class weights inversely proportional to class frequencies."""
+    class_counts = np.bincount(labels)
+    total_samples = len(labels)
+    class_weights = total_samples / (len(class_counts) * class_counts)
+    return torch.tensor(class_weights, dtype=torch.float)
+
+
+# ------------------------ Model Setup ------------------------
+
+class WeightedTrainer(Trainer):
+    """Custom trainer that uses weighted loss function for imbalanced data."""
+    def __init__(self, *args, class_weights=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.class_weights = class_weights.to(self.model.device)
+        self.loss_fn = nn.CrossEntropyLoss(
+            weight=self.class_weights,
+            label_smoothing=0.1
+        )
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        """Compute loss using the weighted loss function."""
+        labels = inputs.get("labels")
+        outputs = model(**inputs)
+        logits = outputs.get("logits")
+        loss = self.loss_fn(logits, labels)
+        return (loss, outputs) if return_outputs else loss
+
+def load_model_and_tokenizer(model_name, label2id, id2label):
+    """Load the model and tokenizer."""
+    tokenizer = DebertaV2Tokenizer.from_pretrained(model_name)
+    model = DebertaV2ForSequenceClassification.from_pretrained(
+        model_name,
+        num_labels=len(label2id),
+        id2label=id2label,
+        label2id=label2id,
+    )
+    model.gradient_checkpointing_enable()
+    return model, tokenizer
+
+# ------------------------ Metrics ------------------------
+
+def compute_metrics(pred):
+    """Compute metrics for evaluation."""
+    labels = pred.label_ids
+    preds = pred.predictions.argmax(-1)
+
+    report = classification_report(
+        labels,
+        preds,
+        target_names=list(label2id.keys()),
+        output_dict=True
+    )
+
+    results = {
+        'accuracy': report['accuracy'],
+        'f1_macro': report['macro avg']['f1-score'],
+        'f1_weighted': report['weighted avg']['f1-score'],
+        'precision_weighted': report['weighted avg']['precision'],
+        'recall_weighted': report['weighted avg']['recall']
+    }
+
+    for cls_name, cls_id in label2id.items():
+        results[f'f1_{cls_name}'] = report[cls_name]['f1-score']
+
+    return results
+
+
+# ------------------------ Training and Evaluation ------------------------
+
+def train_model(df, model_name, label2id, id2label, max_length, batch_size, epochs, seed):
+    """Train and evaluate model using train/val/test split."""
+    df['label_id'] = df['type'].map(label2id)
+
+    # Split the data
+    train_texts, temp_texts, train_labels, temp_labels = train_test_split(
+        df['email_processed'].values,
+        df['label_id'].values,
+        test_size=0.3,
+        stratify=df['label_id'].values,
+        random_state=seed
+    )
+    val_texts, test_texts, val_labels, test_labels = train_test_split(
+        temp_texts,
+        temp_labels,
+        test_size=0.5,
+        stratify=temp_labels,
+        random_state=seed
+    )
+
+    # Create datasets
+    tokenizer = DebertaV2Tokenizer.from_pretrained(model_name)
+    train_dataset = EmailClassification(train_texts, train_labels, tokenizer, max_length)
+    val_dataset = EmailClassification(val_texts, val_labels, tokenizer, max_length)
+    test_dataset = EmailClassification(test_texts, test_labels, tokenizer, max_length)
+
+    # Compute class weights
+    class_weights = compute_class_weights(train_labels)
+
+    # Load the model
+    model, tokenizer = load_model_and_tokenizer(model_name, label2id, id2label)
+
+    # Training arguments
+    training_args = TrainingArguments(
+        output_dir="./results",
+        num_train_epochs=epochs,
+        per_device_train_batch_size=batch_size,
+        per_device_eval_batch_size=batch_size,
+        learning_rate=5e-5,
+        lr_scheduler_type="cosine",
+        warmup_steps=120,
+        weight_decay=0.01,
+        logging_dir="./logs",
+        logging_steps=100,
+        eval_strategy="epoch",
+        save_strategy="epoch",
+        load_best_model_at_end=True,
+        metric_for_best_model="f1_weighted",
+        greater_is_better=True,
+        save_total_limit=1,
+        seed=seed,
+        fp16=True,
+        gradient_accumulation_steps=4
+    )
+
+    # Trainer with Weighted Loss
+    trainer = WeightedTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=val_dataset,
+        compute_metrics=compute_metrics,
+        callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
+        class_weights=class_weights
+    )
+
+    # Train the model
+    trainer.train()
+
+    # Evaluate on validation and test set
+    val_results = trainer.evaluate()
+    test_results = trainer.predict(test_dataset)
+    test_metrics = compute_metrics(test_results)
+
+    # Print results
+    print("\nValidation Results:")
+    for metric_name, value in val_results.items():
+        print(f"{metric_name}: {value:.4f}")
+
+    print("\nTest Results:")
+    for metric_name, value in test_metrics.items():
+        print(f"{metric_name}: {value:.4f}")
+
+    # Get predictions for test set
+    test_preds = np.argmax(test_results.predictions, axis=1)
+
+    print("\nTest Set Classification Report:")
+    print(classification_report(test_labels, test_preds, target_names=list(label2id.keys())))
+
+    # Plot confusion matrix
+    plt.figure(figsize=(10, 8))
+    cm = confusion_matrix(test_labels, test_preds)
+    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=list(label2id.keys()), yticklabels=list(label2id.keys()))
+    plt.xlabel('Predicted')
+    plt.ylabel('True')
+    plt.title('Confusion Matrix - Test Set')
+    plt.tight_layout()
+    plt.savefig('confusion_matrix_test_set.png')
+    plt.close()
+
+    model.save_pretrained('./model')
+    tokenizer.save_pretrained('./model')
+
+    return model, tokenizer
+
+
+# ------------------------ Main ------------------------
+
+if __name__ == "__main__":
+    label2id = {
+        'Incident': 0,
+        'Request': 1,
+        'Problem': 2,
+        'Change': 3
+    }
+    id2label = {v: k for k, v in label2id.items()}
+
+    # Model and Training Configuration
+    MODEL_NAME = "microsoft/mdeberta-v3-base"
+    MAX_LENGTH = 128
+    BATCH_SIZE = 4
+    EPOCHS = 5
+    SEED = 42
+
+    # Load Dataset
+    df = pd.read_csv("D:\OG Project\Data\combined_emails_with_natural_pii.csv")
+
+    # Train and evaluate the model
+    model, tokenizer = train_model(
+        df=df,
+        model_name=MODEL_NAME,
+        label2id=label2id,
+        id2label=id2label,
+        max_length=MAX_LENGTH,
+        batch_size=BATCH_SIZE,
+        epochs=EPOCHS,
+        seed=SEED
+    )

requirements.txt

@@ -0,0 +1,11 @@
+numpy>=1.21.0
+pandas>=1.3.0
+matplotlib>=3.4.0
+seaborn>=0.11.0
+scikit-learn>=1.0.0
+torch>=1.12.0
+transformers>=4.26.0
+sentencepiece>=0.1.95
+protobuf>=3.20.0 
+spacy>=3.2.0
+

utils.py

@@ -0,0 +1,288 @@
+import spacy
+import re
+import logging
+
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Define regex patterns for specific PII fields
+FULL_NAME_PATTERN = r'My name is ([A-Za-z\s\.\-]+)[\.,]|Name\s*:\s*([A-Za-z\s\.\-]+)[\.,]'
+EMAIL_PATTERN = r'You can reach me at ([A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,})|Email\s*:\s*([A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,})'
+PHONE_PATTERN = r'My [Cc]ontact number is\s*([+\d\s\-\(\)\.]+)|Phone\s*:\s*([+\d\s\-\(\)\.]+)|(\+?\d{1,4}?[-.\s]?\(?\d{1,3}?\)?[-.\s]?\d{1,4}[-.\s]?\d{1,4}[-.\s]?\d{1,9})'
+DOB_PATTERN = r'[Dd]ate of [Bb]irth\s*:?\s*(\d{1,2}[-/\.]\d{1,2}[-/\.]\d{2,4}|\d{2,4}[-/\.]\d{1,2}[-/\.]\d{1,2})'
+AADHAR_PATTERN = r'[Aa]adhar(?:\s*[Cc]ard)?\s*(?:[Nn]umber)?\s*:?\s*(\d{4}\s*\d{4}\s*\d{4}|\d{12})'
+CREDIT_DEBIT_PATTERN = r'[Cc](?:redit|ard)\s*(?:[Nn]umber)?\s*:?\s*(\d{4}[-\s]?\d{4}[-\s]?\d{4}[-\s]?\d{4}|\d{16})'
+CVV_PATTERN = r'[Cc][Vv][Vv]\s*(?:[Nn]umber)?\s*:?\s*(\d{3,4})'
+EXPIRY_PATTERN = r'[Ee]xpir(?:y|ation)\s*[Dd]ate\s*:?\s*(\d{1,2}[-/\.]\d{2,4}|\d{2}[-/\.]\d{2})'
+
+def preprocess_text(text):
+    """Clean and normalize text before processing. So that It impoves Model Consistency.
+    Args :
+        text: Email Text to Preprocess before Masking.
+    """
+    # Handle multiple consecutive newlines
+    text = re.sub(r'\n{2,}', '\n', text)
+    # Replace all remaining newlines with space
+    text = text.replace('\n', ' ')
+    # Remove excessive whitespaces So It can reduce Number of tokens
+    text = re.sub(r'\s+', ' ', text)
+    # Replace common unicode whitespace variants
+    text = text.replace('\xa0', ' ')
+    # Remove quotes that might be present in pasted emails
+    text = text.replace('"', '')
+
+    return text.strip()
+
+def extract_entities(email_text):
+    """
+    Extract PII entities from text using regex patterns and spaCy
+
+    Args:
+        email_text (str): The email text to extract entities from
+
+    Returns:
+        list: List of dictionaries containing entity information
+    """
+    if not email_text or len(email_text.strip()) == 0:
+        return []
+
+    # Load spaCy model
+    try:
+        nlp = spacy.load("en_core_web_sm")
+        logger.info("Successfully loaded Spacy model")
+    except OSError as e:
+        logger.error(f"Error loading Spacy model: {e}")
+        raise
+
+    preprocessed_text = preprocess_text(email_text)
+    entities = []
+
+    #Extract full names and also Postitons of Entity
+    name_matches = re.finditer(FULL_NAME_PATTERN, preprocessed_text)
+    for match in name_matches:
+        name = next((g for g in match.groups() if g), "")
+        if name:
+            start_idx = email_text.find(name)
+            if start_idx != -1:
+                entities.append({
+                    "start": start_idx,
+                    "end": start_idx + len(name),
+                    "text": name,
+                    "type": "full_name"
+                })
+
+    #Extract email addresses and also Postitons of Entity
+    email_matches = re.finditer(EMAIL_PATTERN, preprocessed_text)
+    for match in email_matches:
+        email = next((g for g in match.groups() if g), "")
+        if email:
+            start_idx = email_text.find(email)
+            if start_idx != -1:
+                entities.append({
+                    "start": start_idx,
+                    "end": start_idx + len(email),
+                    "text": email,
+                    "type": "email"
+                })
+
+    #Extract phone numbers and also Postitons of Entity
+    phone_matches = re.finditer(PHONE_PATTERN, preprocessed_text)
+    for match in phone_matches:
+        phone = next((g for g in match.groups() if g), "")
+        if phone:
+            start_idx = email_text.find(phone)
+            if start_idx != -1:
+                entities.append({
+                    "start": start_idx,
+                    "end": start_idx + len(phone),
+                    "text": phone,
+                    "type": "phone_number"
+                })
+
+    #Extract date of birth and also Postitons of Entity
+    dob_matches = re.finditer(DOB_PATTERN, preprocessed_text)
+    for match in dob_matches:
+        dob = match.group(1)
+        if dob:
+            start_idx = email_text.find(dob)
+            if start_idx != -1:
+                entities.append({
+                    "start": start_idx,
+                    "end": start_idx + len(dob),
+                    "text": dob,
+                    "type": "dob"
+                })
+
+    #Extract Aadhar numbers and also Postitons of Entity
+    aadhar_matches = re.finditer(AADHAR_PATTERN, preprocessed_text)
+    for match in aadhar_matches:
+        aadhar = match.group(1)
+        if aadhar:
+            start_idx = email_text.find(aadhar)
+            if start_idx != -1:
+                entities.append({
+                    "start": start_idx,
+                    "end": start_idx + len(aadhar),
+                    "text": aadhar,
+                    "type": "aadhar_num"
+                })
+
+    #Extract credit/debit card numbers and also Postitons of Entity
+    card_matches = re.finditer(CREDIT_DEBIT_PATTERN, preprocessed_text)
+    for match in card_matches:
+        card = match.group(1)
+        if card:
+            start_idx = email_text.find(card)
+            if start_idx != -1:
+                entities.append({
+                    "start": start_idx,
+                    "end": start_idx + len(card),
+                    "text": card,
+                    "type": "credit_debit_no"
+                })
+
+    #Extract CVV numbers and also Postitons of Entity
+    cvv_matches = re.finditer(CVV_PATTERN, preprocessed_text)
+    for match in cvv_matches:
+        cvv = match.group(1)
+        if cvv:
+            start_idx = email_text.find(cvv)
+            if start_idx != -1:
+                entities.append({
+                    "start": start_idx,
+                    "end": start_idx + len(cvv),
+                    "text": cvv,
+                    "type": "cvv_no"
+                })
+
+    #Extract card expiry dates and also Postitons of Entity
+    expiry_matches = re.finditer(EXPIRY_PATTERN, preprocessed_text)
+    for match in expiry_matches:
+        expiry = match.group(1)
+        if expiry:
+            start_idx = email_text.find(expiry)
+            if start_idx != -1:
+                entities.append({
+                    "start": start_idx,
+                    "end": start_idx + len(expiry),
+                    "text": expiry,
+                    "type": "expiry_no"
+                })
+
+    # Use spaCy for additional PII detection
+    doc = nlp(preprocessed_text)
+
+    # Used spaCy's more advanced NER capabilities
+    # Created a mapping from character indices in preprocessed_text to original email_text
+    char_mapping = {}
+    preprocessed_idx = 0
+    original_idx = 0
+    
+    while preprocessed_idx < len(preprocessed_text) and original_idx < len(email_text):
+        if preprocessed_text[preprocessed_idx] == email_text[original_idx]:
+            char_mapping[preprocessed_idx] = original_idx
+            preprocessed_idx += 1
+            original_idx += 1
+        else:
+            # Skip characters in original text that were removed during preprocessing
+            original_idx += 1
+    
+    # Process spaCy entities with proper position mapping
+    for ent in doc.ents:
+        entity_type = None
+        
+        # Map spaCy entity types to our entity types
+        if ent.label_ == "PERSON":
+            entity_type = "full_name"
+        elif ent.label_ == "DATE":
+            # Check if it looks like a date of birth or expiry date
+            if re.search(DOB_PATTERN, ent.text):
+                entity_type = "dob"
+            elif re.search(EXPIRY_PATTERN, ent.text):
+                entity_type = "expiry_no"
+        elif ent.label_ == "CARDINAL" or ent.label_ == "MONEY":
+            # Check if it looks like a credit card, aadhar, or other sensitive numbers
+            if re.search(CREDIT_DEBIT_PATTERN, ent.text):
+                entity_type = "credit_debit_no"
+            elif re.search(AADHAR_PATTERN, ent.text):
+                entity_type = "aadhar_num"
+            elif re.search(CVV_PATTERN, ent.text):
+                entity_type = "cvv_no"
+        elif ent.label_ == "ORG" and "@" in ent.text:
+            # Sometimes spaCy identifies emails as organizations
+            entity_type = "email"
+            
+        if entity_type:
+            # Check if this entity overlaps with any existing entity
+            already_captured = False
+            ent_start_idx = char_mapping.get(ent.start_char, -1)
+            
+            if ent_start_idx != -1:
+                ent_end_idx = char_mapping.get(min(ent.end_char, len(preprocessed_text)-1), 
+                                              ent_start_idx + len(ent.text))
+                
+                # Check for overlap with existing entities
+                for entity in entities:
+                    if (entity["start"] <= ent_end_idx and 
+                        entity["end"] >= ent_start_idx):
+                        already_captured = True
+                        break
+                
+                if not already_captured:
+                    # Get the actual text from the original email
+                    original_text = email_text[ent_start_idx:ent_end_idx]
+                    
+                    entities.append({
+                        "start": ent_start_idx,
+                        "end": ent_end_idx,
+                        "text": original_text,
+                        "type": entity_type
+                    })
+
+    # Sort entities by start position
+    entities.sort(key=lambda x: x["start"])
+
+    return entities
+
+def mask_pii(email_text):
+    """
+    Mask PII in email text and return both masked text and entity list.
+    
+    Args:
+        email_text (str): Original email text
+        
+    Returns:
+        tuple: (masked_text, list_of_masked_entities)
+            - masked_text: Text with PII replaced by entity type tags
+            - list_of_masked_entities: List of dictionaries with entity information
+    """
+    if not email_text or len(email_text.strip()) == 0:
+        return "", []
+    
+    # Extract all entities from text
+    entities = extract_entities(email_text)
+    
+    # Create a copy of the original text
+    masked_text = email_text
+    
+    # Process entities in reverse order to avoid index shifting when replacing
+    for entity in sorted(entities, key=lambda x: x["start"], reverse=True):
+        start, end = entity["start"], entity["end"]
+        entity_type = entity["type"]
+        original_text = entity["text"]
+        
+        # Replace the entity with a tag
+        masked_text = masked_text[:start] + f"<{entity_type}>" + masked_text[end:]
+    
+    # Format entities for output 
+    formatted_entities = []
+    for entity in entities:
+        formatted_entities.append({
+            "position": [entity["start"], entity["end"]],
+            "classification": entity["type"],
+            "entity": entity["text"]
+        })
+    
+    return masked_text, formatted_entities
+