refactored the code

.gitignore

@@ -33,7 +33,7 @@ env/
 .vscode/
 *.swp
 *.swo
-
+.flake8
 # Jupyter Notebook
 .ipynb_checkpoints
 

database.py

@@ -4,7 +4,7 @@ Database module for handling email storage operations with SQLite.
 import os
 import json
 import sqlite3
-from typing import Dict, Any, Optional, List, Tuple
+from typing import Dict, Any, Optional, List
 from datetime import datetime
 import uuid
 
@@ -14,29 +14,32 @@ class EmailDatabase:
     Database class for storing and retrieving email data with PII masking information.
     Uses SQLite for storage in Hugging Face's persistent directory.
     """
-    
+
     def __init__(self, connection_string: str = None):
         """
         Initialize the database connection.
-        
+
         Args:
-            connection_string: Database connection string or path. 
-                              For SQLite, this will be treated as a file path.
+            connection_string: Database connection string or path.
+                             For SQLite, this will be treated as a file path.
         """
         # Hugging Face Spaces has a /data directory that persists between restarts
         self.db_path = connection_string or os.environ.get(
-            "DATABASE_PATH", 
+            "DATABASE_PATH",
             "/data/emails.db"  # This path persists in Hugging Face Spaces
         )
-        
+
         # Get the global access key from environment variables
-        self.access_key = os.environ.get("EMAIL_ACCESS_KEY", "default_secure_access_key")
-        
+        self.access_key = os.environ.get(
+            "EMAIL_ACCESS_KEY",
+            "default_secure_access_key"
+        )
+
         # Ensure the data directory exists
         self._ensure_data_directory()
-        
+
         self._create_tables()
-    
+
     def _ensure_data_directory(self):
         """Ensure the data directory exists, and use a fallback if needed."""
         try:
@@ -47,74 +50,77 @@ class EmailDatabase:
             # If we can't write to /data, fall back to the current directory
             self.db_path = "emails.db"
             print(f"Warning: Using fallback database path: {self.db_path}")
-    
+
     def _get_connection(self):
         """Get a database connection."""
         return sqlite3.connect(self.db_path)
-    
+
     def _create_tables(self):
         """Create the necessary tables if they don't exist."""
         conn = self._get_connection()
         try:
             cursor = conn.cursor()
-            
+
             # Create the emails table to store original emails and their masked versions
             cursor.execute('''
             CREATE TABLE IF NOT EXISTS emails (
                 id TEXT PRIMARY KEY,
                 original_email TEXT NOT NULL,
                 masked_email TEXT NOT NULL,
-                masked_entities TEXT NOT NULL, 
+                masked_entities TEXT NOT NULL,
                 category TEXT,
                 created_at TEXT NOT NULL
             )
             ''')
-            
+
             conn.commit()
         except Exception as e:
             conn.rollback()
             raise e
         finally:
             conn.close()
-    
+
     def _generate_id(self) -> str:
         """Generate a unique ID for the email record."""
         return str(uuid.uuid4())
-    
-    def store_email(self, original_email: str, masked_email: str, 
-                   masked_entities: List[Dict[str, Any]], category: Optional[str] = None) -> str:
+
+    def store_email(
+        self, original_email: str, masked_email: str,
+        masked_entities: List[Dict[str, Any]], category: Optional[str] = None
+    ) -> str:
         """
         Store the original email along with its masked version and related information.
-        
+
         Args:
             original_email: The original email with PII
             masked_email: The masked version of the email
             masked_entities: List of entities that were masked
             category: Optional category of the email
-            
+
         Returns:
             email_id for future reference
         """
         conn = self._get_connection()
         try:
             cursor = conn.cursor()
-            
+
             email_id = self._generate_id()
-            
+
             # Store the email data
             cursor.execute(
-                'INSERT INTO emails (id, original_email, masked_email, masked_entities, category, created_at) '
+                'INSERT INTO emails '
+                '(id, original_email, masked_email, masked_entities, category, created_at) '
                 'VALUES (?, ?, ?, ?, ?, ?)',
                 (
                     email_id,
                     original_email,
                     masked_email,
-                    json.dumps(masked_entities),  # Convert to JSON string for SQLite
+                    json.dumps(masked_entities),  # JSON string for SQLite
                     category,
                     datetime.now().isoformat()
                 )
             )
-            
+
             conn.commit()
             return email_id
         except Exception as e:
@@ -122,112 +128,112 @@ class EmailDatabase:
             raise e
         finally:
             conn.close()
-    
+
     def get_original_email(self, email_id: str, access_key: str) -> Optional[Dict[str, Any]]:
         """
         Retrieve the original email with PII using the access key.
-        
+
         Args:
             email_id: The ID of the email record
             access_key: The security key required to access the original email
-            
+
         Returns:
             Dictionary with email data or None if not found or access_key is invalid
         """
         # Verify the access key matches the global access key
         if access_key != self.access_key:
             return None
-            
+
         conn = self._get_connection()
         try:
             cursor = conn.cursor()
-            
+
             cursor.execute(
-                'SELECT id, original_email, masked_email, masked_entities, category, created_at '
-                'FROM emails WHERE id = ?',
+                'SELECT id, original_email, masked_email, masked_entities, category, '
+                'created_at FROM emails WHERE id = ?',
                 (email_id,)
             )
-            
+
             row = cursor.fetchone()
             if not row:
                 return None
-                
+
             return {
                 "id": row[0],
                 "original_email": row[1],
                 "masked_email": row[2],
-                "masked_entities": json.loads(row[3]),  # Convert from JSON string back to Python dict
+                "masked_entities": json.loads(row[3]),  # Convert JSON to dict
                 "category": row[4],
                 "created_at": row[5]
             }
         finally:
             conn.close()
-    
+
     def get_email_by_id(self, email_id: str) -> Optional[Dict[str, Any]]:
         """
         Retrieve the masked email data without the original PII-containing email.
-        
+
         Args:
             email_id: The ID of the email
-            
+
         Returns:
             Dictionary with masked email data or None if not found
         """
         conn = self._get_connection()
         try:
             cursor = conn.cursor()
-            
+
             cursor.execute(
                 'SELECT id, masked_email, masked_entities, category, created_at '
                 'FROM emails WHERE id = ?',
                 (email_id,)
             )
-            
+
             row = cursor.fetchone()
             if not row:
                 return None
-                
+
             return {
                 "id": row[0],
                 "masked_email": row[1],
-                "masked_entities": json.loads(row[2]),  # Convert from JSON string back to Python dict
+                "masked_entities": json.loads(row[2]),  # Convert JSON to dict
                 "category": row[3],
                 "created_at": row[4]
             }
         finally:
             conn.close()
-    
+
     def get_email_by_masked_content(self, masked_email: str) -> Optional[Dict[str, Any]]:
         """
         Retrieve the original email using the masked email content.
-        
+
         Args:
             masked_email: The masked version of the email to search for
-            
+
         Returns:
             Dictionary with full email data or None if not found
         """
         conn = self._get_connection()
         try:
             cursor = conn.cursor()
-            
+
             cursor.execute(
-                'SELECT id, original_email, masked_email, masked_entities, category, created_at '
-                'FROM emails WHERE masked_email = ?',
+                'SELECT id, original_email, masked_email, masked_entities, category, '
+                'created_at FROM emails WHERE masked_email = ?',
                 (masked_email,)
             )
-            
+
             row = cursor.fetchone()
             if not row:
                 return None
-                
+
             return {
                 "id": row[0],
                 "original_email": row[1],
                 "masked_email": row[2],
-                "masked_entities": json.loads(row[3]),  # Convert from JSON string back to Python dict
+                "masked_entities": json.loads(row[3]),  # Convert JSON to dict
                 "category": row[4],
                 "created_at": row[5]
             }
         finally:
-            conn.close()
+            conn.close()

main.py

@@ -1,7 +1,7 @@
 import os
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from typing import Dict, Any, List, Tuple, Optional
+from typing import Dict, Any, List, Tuple
 import uvicorn
 from dotenv import load_dotenv
 
@@ -21,7 +21,7 @@ else:
     db_path = "emails.db"  # Fallback to local directory
 
 # Initialize the FastAPI application
-app = FastAPI(title="Email Classification API", 
+app = FastAPI(title="Email Classification API",
               description="API for classifying support emails and masking PII",
               version="1.0.0")
 
@@ -29,16 +29,19 @@ app = FastAPI(title="Email Classification API",
 pii_masker = PIIMasker(db_path=db_path)
 email_classifier = EmailClassifier()
 
+
 class EmailInput(BaseModel):
     """Input model for the email classification endpoint"""
     input_email_body: str
 
+
 class EntityInfo(BaseModel):
     """Model for entity information"""
     position: Tuple[int, int]
-    classification: str  
+    classification: str
     entity: str
 
+
 class EmailOutput(BaseModel):
     """Output model for the email classification endpoint"""
     input_email_body: str
@@ -46,29 +49,30 @@ class EmailOutput(BaseModel):
     masked_email: str
     category_of_the_email: str
 
+
 class MaskedEmailInput(BaseModel):
     """Input model for retrieving original email by masked email content"""
     masked_email: str
     access_key: str
 
+
 @app.post("/classify", response_model=EmailOutput)
 async def classify_email(email_input: EmailInput) -> Dict[str, Any]:
     """
     Classify an email into a support category while masking PII
-    
+
     Args:
         email_input: The input email data
-        
+
     Returns:
         The classified email data with masked PII
     """
     try:
         # Process the email to mask PII and store original in database
         processed_data = pii_masker.process_email(email_input.input_email_body)
-        
         # Classify the masked email
         classified_data = email_classifier.process_email(processed_data)
-        
+
         # Make sure we return only the fields expected in the response model
         return {
             "input_email_body": email_input.input_email_body,
@@ -79,28 +83,35 @@ async def classify_email(email_input: EmailInput) -> Dict[str, Any]:
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error processing email: {str(e)}")
 
+
 @app.post("/api/v1/unmask-email", response_model=Dict[str, Any])
 async def unmask_email(masked_email_input: MaskedEmailInput) -> Dict[str, Any]:
     """
-    Retrieve the original unmasked email using the masked email content from the classify response.
-    
+    Retrieve the original unmasked email.
+
     Args:
         masked_email_input: Contains the masked email and access key
-        
+
     Returns:
         The original email data with PII information
     """
     try:
         # Verify access key matches the global access key
-        if masked_email_input.access_key != os.environ.get("EMAIL_ACCESS_KEY", "default_secure_access_key"):
+        if masked_email_input.access_key != os.environ.get(
+                "EMAIL_ACCESS_KEY", "default_secure_access_key"):
             raise HTTPException(status_code=401, detail="Invalid access key")
-        
+
         # Retrieve the original email using the masked content
-        email_data = pii_masker.get_original_by_masked_email(masked_email_input.masked_email)
-        
+        email_data = pii_masker.get_original_by_masked_email(
+            masked_email_input.masked_email
+        )
+
         if not email_data:
-            raise HTTPException(status_code=404, detail="Original email not found for the provided masked email")
-        
+            raise HTTPException(
+                status_code=404,
+                detail="Original email not found for the provided masked email"
+            )
+
         return {
             "status": "success",
             "data": {
@@ -116,19 +127,24 @@ async def unmask_email(masked_email_input: MaskedEmailInput) -> Dict[str, Any]:
     except Exception as e:
         if isinstance(e, HTTPException):
             raise e
-        raise HTTPException(status_code=500, detail=f"Error retrieving original email: {str(e)}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error retrieving original email: {str(e)}"
+        )
+
 
 @app.get("/health")
 async def health_check():
     """
     Health check endpoint
-    
+
     Returns:
         Status message indicating the API is running
     """
     return {"status": "healthy", "message": "Email classification API is running"}
 
+
 # For local development and testing
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 8000))
-    uvicorn.run("main:app", host="0.0.0.0", port=port, reload=True)
+    uvicorn.run("main:app", host="0.0.0.0", port=port, reload=True)

models.py

@@ -3,79 +3,81 @@ import torch
 from transformers import XLMRobertaForSequenceClassification, XLMRobertaTokenizer
 from typing import Dict, Any
 
+
 class EmailClassifier:
     """
     Email classification model to categorize emails into different support categories
     """
-    
     CATEGORIES = ['Change', 'Incident', 'Problem', 'Request']
-    
+
     def __init__(self, model_path: str = None):
         """
         Initialize the email classifier with a pre-trained model
-        
+
         Args:
             model_path: Path or Hugging Face Hub model ID
         """
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        
+
         # Use environment variable for model path or fall back to Hugging Face Hub model
         # This allows for flexibility in deployment
-        model_path = model_path or os.environ.get("MODEL_PATH", "Sparkonix11/email-classifier-model")
-        
+        model_path = model_path or os.environ.get(
+            "MODEL_PATH", "Sparkonix11/email-classifier-model"
+        )
+
         # Load the tokenizer and model from Hugging Face Hub or local path
         self.tokenizer = XLMRobertaTokenizer.from_pretrained(model_path)
         self.model = XLMRobertaForSequenceClassification.from_pretrained(model_path)
         self.model.to(self.device)
         self.model.eval()
-    
+
     def classify(self, masked_email: str) -> str:
         """
         Classify a masked email into one of the predefined categories
-        
+
         Args:
             masked_email: The email content with PII masked
-            
+
         Returns:
             The predicted category as a string
         """
         # Tokenize the masked email
         inputs = self.tokenizer(
-            masked_email, 
+            masked_email,
             return_tensors="pt",
             padding="max_length",
             truncation=True,
             max_length=512
         )
-        
+
         inputs = {key: val.to(self.device) for key, val in inputs.items()}
-        
+
         # Perform inference
         with torch.no_grad():
             outputs = self.model(**inputs)
             logits = outputs.logits
             predicted_class_idx = torch.argmax(logits, dim=1).item()
-        
+
         # Map the predicted class index to the category
         return self.CATEGORIES[predicted_class_idx]
-    
+
     def process_email(self, masked_email_data: Dict[str, Any]) -> Dict[str, Any]:
         """
         Process an email by classifying it into a category
-        
+
         Args:
             masked_email_data: Dictionary containing the masked email and other data
-            
+
         Returns:
             The input dictionary with the classification added
         """
         # Extract masked email content
         masked_email = masked_email_data["masked_email"]
-        
+
         # Classify the masked email
         category = self.classify(masked_email)
-        
+
         # Add the classification to the data
         masked_email_data["category_of_the_email"] = category
-        
-        return masked_email_data
+
+        return masked_email_data

requirements.in

@@ -1,381 +0,0 @@
-# This file was autogenerated by uv via the following command:
-#    uv pip compile requirements.in -o requirements.txt
-annotated-types==0.7.0
-    # via
-    #   -r requirements.in
-    #   pydantic
-anyio==4.9.0
-    # via
-    #   -r requirements.in
-    #   starlette
-blis==1.3.0
-    # via
-    #   -r requirements.in
-    #   thinc
-catalogue==2.0.10
-    # via
-    #   -r requirements.in
-    #   spacy
-    #   srsly
-    #   thinc
-certifi==2025.4.26
-    # via
-    #   -r requirements.in
-    #   requests
-charset-normalizer==3.4.2
-    # via
-    #   -r requirements.in
-    #   requests
-click==8.2.0
-    # via
-    #   -r requirements.in
-    #   typer
-    #   uvicorn
-cloudpathlib==0.21.1
-    # via
-    #   -r requirements.in
-    #   weasel
-confection==0.1.5
-    # via
-    #   -r requirements.in
-    #   thinc
-    #   weasel
-cymem==2.0.11
-    # via
-    #   -r requirements.in
-    #   preshed
-    #   spacy
-    #   thinc
-en-core-web-sm @ https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl
-    # via -r requirements.in
-exceptiongroup==1.3.0
-    # via
-    #   -r requirements.in
-    #   anyio
-fastapi==0.115.12
-    # via -r requirements.in
-filelock==3.18.0
-    # via
-    #   -r requirements.in
-    #   huggingface-hub
-    #   torch
-    #   transformers
-fsspec==2025.3.2
-    # via
-    #   -r requirements.in
-    #   huggingface-hub
-    #   torch
-h11==0.16.0
-    # via
-    #   -r requirements.in
-    #   uvicorn
-huggingface-hub==0.31.2
-    # via
-    #   -r requirements.in
-    #   tokenizers
-    #   transformers
-idna==3.10
-    # via
-    #   -r requirements.in
-    #   anyio
-    #   requests
-jinja2==3.1.6
-    # via
-    #   -r requirements.in
-    #   spacy
-    #   torch
-langcodes==3.5.0
-    # via
-    #   -r requirements.in
-    #   spacy
-language-data==1.3.0
-    # via
-    #   -r requirements.in
-    #   langcodes
-marisa-trie==1.2.1
-    # via
-    #   -r requirements.in
-    #   language-data
-markdown-it-py==3.0.0
-    # via
-    #   -r requirements.in
-    #   rich
-markupsafe==3.0.2
-    # via
-    #   -r requirements.in
-    #   jinja2
-mdurl==0.1.2
-    # via
-    #   -r requirements.in
-    #   markdown-it-py
-mpmath==1.3.0
-    # via
-    #   -r requirements.in
-    #   sympy
-murmurhash==1.0.12
-    # via
-    #   -r requirements.in
-    #   preshed
-    #   spacy
-    #   thinc
-networkx==3.4.2
-    # via
-    #   -r requirements.in
-    #   torch
-numpy==2.2.5
-    # via
-    #   -r requirements.in
-    #   blis
-    #   spacy
-    #   spacy-transformers
-    #   thinc
-    #   transformers
-
-# SQLite is included in Python standard library
-python-dotenv
-    # for environment variable management
-
-nvidia-cublas-cu12==12.6.4.1
-    # via
-    #   -r requirements.in
-    #   nvidia-cudnn-cu12
-    #   nvidia-cusolver-cu12
-    #   torch
-nvidia-cuda-cupti-cu12==12.6.80
-    # via
-    #   -r requirements.in
-    #   torch
-nvidia-cuda-nvrtc-cu12==12.6.77
-    # via
-    #   -r requirements.in
-    #   torch
-nvidia-cuda-runtime-cu12==12.6.77
-    # via
-    #   -r requirements.in
-    #   torch
-nvidia-cudnn-cu12==9.5.1.17
-    # via
-    #   -r requirements.in
-    #   torch
-nvidia-cufft-cu12==11.3.0.4
-    # via
-    #   -r requirements.in
-    #   torch
-nvidia-cufile-cu12==1.11.1.6
-    # via
-    #   -r requirements.in
-    #   torch
-nvidia-curand-cu12==10.3.7.77
-    # via
-    #   -r requirements.in
-    #   torch
-nvidia-cusolver-cu12==11.7.1.2
-    # via
-    #   -r requirements.in
-    #   torch
-nvidia-cusparse-cu12==12.5.4.2
-    # via
-    #   -r requirements.in
-    #   nvidia-cusolver-cu12
-    #   torch
-nvidia-cusparselt-cu12==0.6.3
-    # via
-    #   -r requirements.in
-    #   torch
-nvidia-nccl-cu12==2.26.2
-    # via
-    #   -r requirements.in
-    #   torch
-nvidia-nvjitlink-cu12==12.6.85
-    # via
-    #   -r requirements.in
-    #   nvidia-cufft-cu12
-    #   nvidia-cusolver-cu12
-    #   nvidia-cusparse-cu12
-    #   torch
-nvidia-nvtx-cu12==12.6.77
-    # via
-    #   -r requirements.in
-    #   torch
-packaging==25.0
-    # via
-    #   -r requirements.in
-    #   huggingface-hub
-    #   spacy
-    #   thinc
-    #   transformers
-    #   weasel
-preshed==3.0.9
-    # via
-    #   -r requirements.in
-    #   spacy
-    #   thinc
-pydantic==2.11.4
-    # via
-    #   -r requirements.in
-    #   confection
-    #   fastapi
-    #   spacy
-    #   thinc
-    #   weasel
-pydantic-core==2.33.2
-    # via
-    #   -r requirements.in
-    #   pydantic
-pygments==2.19.1
-    # via
-    #   -r requirements.in
-    #   rich
-python-multipart==0.0.20
-    # via -r requirements.in
-pyyaml==6.0.2
-    # via
-    #   -r requirements.in
-    #   huggingface-hub
-    #   transformers
-regex==2024.11.6
-    # via
-    #   -r requirements.in
-    #   transformers
-requests==2.32.3
-    # via
-    #   -r requirements.in
-    #   huggingface-hub
-    #   spacy
-    #   transformers
-    #   weasel
-rich==14.0.0
-    # via
-    #   -r requirements.in
-    #   typer
-safetensors==0.5.3
-    # via
-    #   -r requirements.in
-    #   transformers
-sentencepiece==0.2.0
-    # via -r requirements.in
-setuptools==80.7.1
-    # via
-    #   -r requirements.in
-    #   marisa-trie
-    #   spacy
-    #   thinc
-    #   triton
-shellingham==1.5.4
-    # via
-    #   -r requirements.in
-    #   typer
-smart-open==7.1.0
-    # via
-    #   -r requirements.in
-    #   weasel
-sniffio==1.3.1
-    # via
-    #   -r requirements.in
-    #   anyio
-spacy==3.8.5
-    # via
-    #   -r requirements.in
-    #   spacy-transformers
-spacy-alignments==0.9.1
-    # via
-    #   -r requirements.in
-    #   spacy-transformers
-spacy-legacy==3.0.12
-    # via
-    #   -r requirements.in
-    #   spacy
-spacy-loggers==1.0.5
-    # via
-    #   -r requirements.in
-    #   spacy
-spacy-transformers==1.3.8
-    # via -r requirements.in
-srsly==2.5.1
-    # via
-    #   -r requirements.in
-    #   confection
-    #   spacy
-    #   spacy-transformers
-    #   thinc
-    #   weasel
-starlette==0.46.2
-    # via
-    #   -r requirements.in
-    #   fastapi
-sympy==1.14.0
-    # via
-    #   -r requirements.in
-    #   torch
-thinc==8.3.6
-    # via
-    #   -r requirements.in
-    #   spacy
-tokenizers==0.21.1
-    # via
-    #   -r requirements.in
-    #   transformers
-torch==2.7.0
-    # via
-    #   -r requirements.in
-    #   spacy-transformers
-tqdm==4.67.1
-    # via
-    #   -r requirements.in
-    #   huggingface-hub
-    #   spacy
-    #   transformers
-transformers==4.49.0
-    # via
-    #   -r requirements.in
-    #   spacy-transformers
-triton==3.3.0
-    # via
-    #   -r requirements.in
-    #   torch
-typer==0.15.3
-    # via
-    #   -r requirements.in
-    #   spacy
-    #   weasel
-typing-extensions==4.13.2
-    # via
-    #   -r requirements.in
-    #   anyio
-    #   cloudpathlib
-    #   exceptiongroup
-    #   fastapi
-    #   huggingface-hub
-    #   pydantic
-    #   pydantic-core
-    #   rich
-    #   torch
-    #   typer
-    #   typing-inspection
-    #   uvicorn
-typing-inspection==0.4.0
-    # via
-    #   -r requirements.in
-    #   pydantic
-urllib3==2.4.0
-    # via
-    #   -r requirements.in
-    #   requests
-uvicorn==0.34.2
-    # via -r requirements.in
-wasabi==1.1.3
-    # via
-    #   -r requirements.in
-    #   spacy
-    #   thinc
-    #   weasel
-weasel==0.4.1
-    # via
-    #   -r requirements.in
-    #   spacy
-wrapt==1.17.2
-    # via
-    #   -r requirements.in
-    #   smart-open
-xx-ent-wiki-sm @ https://github.com/explosion/spacy-models/releases/download/xx_ent_wiki_sm-3.8.0/xx_ent_wiki_sm-3.8.0-py3-none-any.whl
-    # via -r requirements.in

upload_model.py

@@ -2,27 +2,30 @@
 Script to upload the email classification model to Hugging Face Hub
 """
 
-import os
 import sys
 import argparse
 import subprocess
 import pkg_resources
 
+
 def check_and_install_dependencies():
     """Check for required libraries and install if missing"""
     required_packages = ['torch', 'transformers', 'sentencepiece']
     installed_packages = {pkg.key for pkg in pkg_resources.working_set}
-    
+
     missing_packages = [pkg for pkg in required_packages if pkg not in installed_packages]
-    
+
     if missing_packages:
-        print(f"Installing missing dependencies: {', '.join(missing_packages)}")
-        subprocess.check_call([sys.executable, "-m", "pip", "install"] + missing_packages)
+        missing_packages_str = ", ".join(missing_packages)
+        print(f"Installing missing dependencies: {missing_packages_str}")
+        subprocess.check_call([sys.executable, "-m", "pip", "install"]
+                              + missing_packages)
         print("Dependencies installed. You may need to restart the script.")
         return False
-    
+
     return True
 
+
 def get_huggingface_username(token=None):
     """Get the username for the authenticated user"""
     try:
@@ -34,51 +37,57 @@ def get_huggingface_username(token=None):
         print(f"Error getting Hugging Face username: {e}")
         return None
 
+
 def main():
     """Upload model to Hugging Face Hub"""
     # Check dependencies first
     if not check_and_install_dependencies():
         return
-    
+
     # Import dependencies after installation check
     from transformers import XLMRobertaForSequenceClassification, XLMRobertaTokenizer
     from huggingface_hub import login
-    
-    parser = argparse.ArgumentParser(description="Upload email classification model to Hugging Face Hub")
+
+    parser = argparse.ArgumentParser(
+        description="Upload email classification model to Hugging Face Hub")
     parser.add_argument("--model_path", type=str, default="classification_model",
                         help="Local path to the model files")
     parser.add_argument("--hub_model_id", type=str,
-                        help="Hugging Face Hub model ID (e.g., 'username/email-classifier-model')")
+                        help="Hugging Face Hub model ID (e.g., "
+                             "'username/email-classifier-model')")
     parser.add_argument("--model_name", type=str, default="email-classifier-model",
-                        help="Name for the model repository (default: email-classifier-model)")
+                        help="Name for the model repository "
+                             "(default: email-classifier-model)")
     parser.add_argument("--token", type=str,
-                        help="Hugging Face API token (optional, can use environment variable or huggingface-cli login)")
-    
+                        help="Hugging Face API token (optional, can use "
+                             "environment variable or huggingface-cli login)")
+
     args = parser.parse_args()
-    
+
     # Login if token is provided
     if args.token:
         login(token=args.token)
-    
+
     # If hub_model_id is not provided, try to get username and construct it
     if not args.hub_model_id:
         username = get_huggingface_username(args.token)
         if not username:
-            print("Could not determine Hugging Face username. Please provide --hub_model_id explicitly.")
+            print("Could not determine Hugging Face username. "
+                  "Please provide --hub_model_id explicitly.")
             return
         args.hub_model_id = f"{username}/{args.model_name}"
-    
+
     print(f"Loading model from {args.model_path}...")
     # Load the local model and tokenizer
     model = XLMRobertaForSequenceClassification.from_pretrained(args.model_path)
     tokenizer = XLMRobertaTokenizer.from_pretrained(args.model_path)
-    
+
     print(f"Uploading model to {args.hub_model_id}...")
     try:
         # Push to Hugging Face Hub
         model.push_to_hub(args.hub_model_id)
         tokenizer.push_to_hub(args.hub_model_id)
-        
+
         print("Model successfully uploaded to Hugging Face Hub!")
         print(f"You can now use the model with the ID: {args.hub_model_id}")
         print(f"Update the MODEL_PATH in Dockerfile to: {args.hub_model_id}")
@@ -86,8 +95,11 @@ def main():
         print(f"Error uploading model: {e}")
         print("\nPossible solutions:")
         print("1. Make sure you're logged in with 'huggingface-cli login'")
-        print("2. Check that you have permission to create repos in the specified namespace")
-        print("3. Try using your own username: --hub_model_id yourusername/email-classifier-model")
+        print("2. Check that you have permission to create repos in the "
+              "specified namespace")
+        print("3. Try using your own username: "
+              "--hub_model_id yourusername/email-classifier-model")
+
 
 if __name__ == "__main__":
-    main()
+    main()

utils.py

@@ -1,8 +1,10 @@
 import re
 import spacy
 from typing import List, Dict, Tuple, Any, Optional
+
 from database import EmailDatabase
 
+
 class Entity:
     def __init__(self, start: int, end: int, entity_type: str, value: str):
         self.start = start
@@ -17,11 +19,19 @@ class Entity:
             "entity": self.value
         }
 
-    def __repr__(self): # Added for easier debugging
-        return f"Entity(type='{self.entity_type}', value='{self.value}', start={self.start}, end={self.end})"
+    def __repr__(self):  # Added for easier debugging
+        return (
+            f"Entity(type='{self.entity_type}', value='{self.value}', "
+            f"start={self.start}, end={self.end})"
+        )
+
 
 class PIIMasker:
-    def __init__(self, spacy_model_name: str = "xx_ent_wiki_sm", db_path: str = None): # Allow model choice
+    def __init__(
+        self,
+        spacy_model_name: str = "xx_ent_wiki_sm",
+        db_path: str = None
+    ):  # Allow model choice
         # Load SpaCy model
         try:
             self.nlp = spacy.load(spacy_model_name)
@@ -42,7 +52,7 @@ class PIIMasker:
 
         # Initialize database connection with SQLite path
         self.db = EmailDatabase(connection_string=db_path)
-        
+
         # Initialize regex patterns
         self._initialize_patterns()
 
@@ -51,7 +61,11 @@ class PIIMasker:
         self.patterns = {
             "email": r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b',
             # Simplified phone regex to capture both standard and international formats
-            "phone_number": r'\b(?:(?:\+|00)[1-9]\d{0,3}[-\s.]?)?(?:\(?\d{1,5}\)?[-\s.]?)?\d{1,5}(?:[-\s.]\d{1,5}){1,4}\b',
+            "phone_number": (
+                r'\b(?:(?:\+|00)[1-9]\d{0,3}[-\s.]?)?'
+                r'(?:\(?\d{1,5}\)?[-\s.]?)?\d{1,5}'
+                r'(?:[-\s.]\d{1,5}){1,4}\b'
+            ),
             # Card number regex: common formats, allows optional spaces/hyphens
             "credit_debit_no": r'\b(?:(?:\d{4}[\s-]?){3}\d{4}|\d{13,19})\b',
             # CVV: 3 or 4 digits, ensuring it's a standalone number (word boundary)
@@ -60,7 +74,10 @@ class PIIMasker:
             "expiry_no": r'\b(0[1-9]|1[0-2])[/\s-]([0-9]{2}|20[0-9]{2})\b',
             "aadhar_num": r'\b\d{4}\s?\d{4}\s?\d{4}\b',
             # DOB: DD/MM/YYYY or DD-MM-YYYY etc.
-            "dob": r'\b(0[1-9]|[12][0-9]|3[01])[/\s-](0[1-9]|1[0-2])[/\s-](?:19|20)\d\d\b'
+            "dob": (
+                r'\b(0[1-9]|[12][0-9]|3[01])[/\s-]'
+                r'(0[1-9]|1[0-2])[/\s-](?:19|20)\d\d\b'
+            )
         }
 
     def detect_regex_entities(self, text: str) -> List[Entity]:
@@ -83,14 +100,19 @@ class PIIMasker:
                     if not self.verify_phone_number(text, match):
                         continue
                 elif entity_type == "dob":
-                    if not self._verify_with_context(text, start, end, ["birth", "dob", "born"]):
+                    if not self._verify_with_context(
+                        text, start, end, ["birth", "dob", "born"]
+                    ):
                         continue
 
-                # Avoid detecting parts of already matched longer entities (e.g. year within a DOB)
+                # Avoid detecting parts of already matched longer entities
+                # (e.g. year within a DOB)
                 # This is a simple check; more robust overlap handling is done later
                 is_substring_of_existing = False
                 for existing_entity in entities:
-                    if existing_entity.start <= start and existing_entity.end >= end and existing_entity.value != value:
+                    if (existing_entity.start <= start
+                            and existing_entity.end >= end  # W504 corrected
+                            and existing_entity.value != value):  # W504 corrected
                         is_substring_of_existing = True
                         break
                 if is_substring_of_existing:
@@ -99,7 +121,9 @@ class PIIMasker:
                 entities.append(Entity(start, end, entity_type, value))
         return entities
 
-    def _verify_with_context(self, text: str, start: int, end: int, keywords: List[str], window: int = 50) -> bool:
+    def _verify_with_context(
+        self, text: str, start: int, end: int, keywords: List[str], window: int = 50
+    ) -> bool:
         """Verify an entity match using surrounding context"""
         context_before = text[max(0, start - window):start].lower()
         context_after = text[end:min(len(text), end + window)].lower()
@@ -117,7 +141,10 @@ class PIIMasker:
         context_before = text[max(0, start - context_window):start].lower()
         context_after = text[end:min(len(text), end + context_window)].lower()
 
-        card_keywords = ["card", "credit", "debit", "visa", "mastercard", "payment", "amex", "account no", "card no"]
+        card_keywords = [
+            "card", "credit", "debit", "visa", "mastercard",
+            "payment", "amex", "account no", "card no"
+        ]
         for keyword in card_keywords:
             if keyword in context_before or keyword in context_after:
                 return True
@@ -125,19 +152,19 @@ class PIIMasker:
         # For simplicity, we'll rely on context here. If needed, Luhn can be added.
         return False
 
-
     def verify_cvv(self, text: str, match: re.Match) -> bool:
         """Verify if a 3-4 digit number is actually a CVV using contextual clues"""
         context_window = 50
         start, end = match.span()
         value = match.group()
 
-        # If it's part of a longer number sequence (like a phone number or ID), it's likely not a CVV
+        # If it's part of a longer number sequence (like a phone number or ID),
+        # it's likely not a CVV
         # Check character immediately before and after
-        char_before = text[start-1:start] if start > 0 else ""
-        char_after = text[end:end+1] if end < len(text) else ""
+        char_before = text[start - 1:start] if start > 0 else ""
+        char_after = text[end:end + 1] if end < len(text) else ""
         if char_before.isdigit() or char_after.isdigit():
-            return False # It's part of a larger number
+            return False  # It's part of a larger number
 
         # Only consider 3-4 digit numbers
         if not value.isdigit() or len(value) < 3 or len(value) > 4:
@@ -148,14 +175,16 @@ class PIIMasker:
 
         # Expanded list of CVV-related keywords to improve detection
         cvv_keywords = [
-            "cvv", "cvc", "csc", "security code", "card verification", "verification no",
-            "security", "security number", "cv2", "card code", "security value"
+            "cvv", "cvc", "csc", "security code", "card verification",
+            "verification no", "security", "security number", "cv2",
+            "card code", "security value"
         ]
-        
-        date_keywords = ["date", "year", "/", "born", "age", "since", "established"]
 
         # Look for CVV context clues
-        is_cvv_context = any(keyword in context_before or keyword in context_after for keyword in cvv_keywords)
+        is_cvv_context = any(
+            keyword in context_before or keyword in context_after
+            for keyword in cvv_keywords
+        )
 
         # If explicitly mentioned as a CVV, immediately return true
         if is_cvv_context:
@@ -163,17 +192,23 @@ class PIIMasker:
 
         # If it looks like a year, reject it
         if len(value) == 4 and 1900 <= int(value) <= 2100:
-            if any(k in context_before or k in context_after for k in ["year", "born", "established", "since"]):
+            if any(
+                k in context_before or k in context_after
+                for k in ["year", "born", "established", "since"]
+            ):
                 return False
 
         # If in expiry date context, reject it
         if re.search(r'\b(0[1-9]|1[0-2])[/\s-]$', context_before.strip()):
             return False
-            
-        # If no context clues but we have a credit card mention nearby, it could be a CVV
-        card_context = any(k in context_before or k in context_after 
-                          for k in ["card", "credit", "visa", "mastercard", "amex", "discover"])
-        
+
+        # If no context clues but we have a credit card mention nearby,
+        # it could be a CVV
+        card_context = any(
+            k in context_before or k in context_after for k in
+            ["card", "credit", "visa", "mastercard", "amex", "discover"]
+        )
+
         return is_cvv_context or (card_context and len(value) in [3, 4])
 
     def verify_phone_number(self, text: str, match: re.Match) -> bool:
@@ -182,56 +217,62 @@ class PIIMasker:
         """
         value = match.group()
         start, end = match.span()
-        
+
         # Extract only digits to count them
         digits = ''.join(c for c in value if c.isdigit())
         digit_count = len(digits)
-        
+
         # Most phone numbers worldwide have between 7 and 15 digits
         if digit_count < 7 or digit_count > 15:
             return False
-        
+
         # Check for common phone number indicators
         context_window = 50
         context_before = text[max(0, start - context_window):start].lower()
         context_after = text[end:min(len(text), end + context_window)].lower()
-        
+
         # Expanded phone keywords
         phone_keywords = [
-            "phone", "call", "tel", "telephone", "contact", "dial", "mobile", "cell", 
-            "number", "direct", "office", "fax", "reach me at", "call me", "contact me",
-            "line", "extension", "ext", "phone number"
+            "phone", "call", "tel", "telephone", "contact", "dial", "mobile",
+            "cell", "number", "direct", "office", "fax", "reach me at",
+            "call me", "contact me", "line", "extension", "ext", "phone number"
         ]
-        
+
         # Check for phone context
-        has_phone_context = any(kw in context_before or kw in context_after for kw in phone_keywords)
-        
+        has_phone_context = any(
+            kw in context_before or kw in context_after for kw in phone_keywords
+        )
+
         # Check for formatting that indicates a phone number
-        has_phone_formatting = bool(re.search(r'[-\s.()\+]', value))
-        
+        has_phone_formatting = bool(re.search(r'[-\s.()\\+]', value))
+
         # Check for international prefix
         has_intl_prefix = value.startswith('+') or value.startswith('00')
-        
+
         # Return true if any of these conditions are met:
         # 1. Has explicit phone context
         # 2. Has phone-like formatting AND reasonable digit count
         # 3. Has international prefix AND reasonable digit count
         # 4. Has 10 digits exactly (common in many countries) with formatting
-        return has_phone_context or \
-               (has_phone_formatting and digit_count >= 7) or \
-               (has_intl_prefix) or \
-               (digit_count == 10 and has_phone_formatting)
+        return (
+            has_phone_context
+            or (has_phone_formatting and digit_count >= 7)
+            or (has_intl_prefix)
+            or (digit_count == 10 and has_phone_formatting)
+        )
 
     def detect_name_entities(self, text: str) -> List[Entity]:
         """Detect name entities using SpaCy NER"""
         entities = []
         doc = self.nlp(text)
-        
+
         for ent in doc.ents:
             # Use PER for person, common in many models like xx_ent_wiki_sm
             # Also checking for PERSON as some models might use it.
             if ent.label_ in ["PER", "PERSON"]:
-                entities.append(Entity(ent.start_char, ent.end_char, "full_name", ent.text))
+                entities.append(
+                    Entity(ent.start_char, ent.end_char, "full_name", ent.text)
+                )
         return entities
 
     def detect_all_entities(self, text: str) -> List[Entity]:
@@ -265,58 +306,74 @@ class PIIMasker:
         # A simple greedy approach: iterate and remove/adjust overlaps
         # This can be made more sophisticated
         resolved_entities: List[Entity] = []
-        for current_entity in sorted(entities, key=lambda e: (e.start, -(e.end - e.start))): # Process by start, then by longest
+        # Process by start, then by longest
+        for current_entity in sorted(
+            entities, key=lambda e: (e.start, -(e.end - e.start))
+        ):
             is_overlapped_or_contained = False
             temp_resolved = []
             for i, res_entity in enumerate(resolved_entities):
                 # Check for overlap:
                 # Current: |----|
                 # Res:         |----|   or |----| or   |--|  or  |------|
-                overlap = max(0, min(current_entity.end, res_entity.end) - max(current_entity.start, res_entity.start))
+                overlap = max(
+                    0,
+                    min(current_entity.end, res_entity.end)  # Fixed W504 line break
+                    - max(current_entity.start, res_entity.start)
+                )
 
                 if overlap > 0:
                     is_overlapped_or_contained = True
                     # Preference:
-                    # 1. NER names often trump regex if they are the ones causing overlap
+                    # 1. NER often trump regex if they are the ones causing overlap
                     # 2. Longer entity wins
                     current_len = current_entity.end - current_entity.start
                     res_len = res_entity.end - res_entity.start
 
-                    # If current is a name and overlaps, and previous is not a name, prefer current if it's not fully contained
-                    if current_entity.entity_type == "full_name" and res_entity.entity_type != "full_name":
-                        if not (res_entity.start <= current_entity.start and res_entity.end >= current_entity.end): # current not fully contained by res
-                             # remove res_entity, current will be added later
-                            continue # go to next res_entity, this one is marked for removal
-                    elif res_entity.entity_type == "full_name" and current_entity.entity_type != "full_name":
-                         # res_entity is a name, current is not. Prefer res_entity if it's not fully contained
-                        if not (current_entity.start <= res_entity.start and current_entity.end >= res_entity.end):
-                            # current entity is subsumed or less important, so don't add current
-                            # and keep res_entity
+                    # If current is a name and overlaps, and previous is not a name,
+                    # prefer current if it's not fully contained
+                    if (current_entity.entity_type == "full_name"  # E501 corrected
+                            and res_entity.entity_type != "full_name"):
+                        # current not fully contained by res
+                        if not (res_entity.start <= current_entity.start
+                                and res_entity.end >= current_entity.end):
+                            # remove res_entity, current will be added later
+                            continue  # go to next res_entity, marked for removal
+                    elif (res_entity.entity_type == "full_name"
+                          and current_entity.entity_type != "full_name"):
+                        # res_entity is a name, current is not. Prefer res_entity
+                        # if it's not fully contained
+                        if not (current_entity.start <= res_entity.start
+                                and current_entity.end >= res_entity.end):
+                            # current entity is subsumed or less important,
+                            # so don't add current and keep res_entity
                             temp_resolved.append(res_entity)
-                            is_overlapped_or_contained = True # Mark current as handled
-                            break # Current is dominated
+                            is_overlapped_or_contained = True  # Mark current as handled
+                            break  # Current is dominated
 
                     # General case: longer entity wins
                     if current_len > res_len:
-                        # current is longer, res_entity is removed from consideration for this current_entity
-                        pass # res_entity will not be added to temp_resolved if it's fully replaced
+                        # current is longer, res_entity is removed from
+                        # consideration for this current_entity
+                        pass  # res_entity not added to temp_resolved if fully replaced
                     elif res_len > current_len:
                         # res is longer, current is dominated
                         temp_resolved.append(res_entity)
-                        is_overlapped_or_contained = True # Mark current as handled
+                        is_overlapped_or_contained = True  # Mark current as handled
                         break
-                    else: # Same length, keep existing one (res_entity)
+                    else:  # Same length, keep existing one (res_entity)
                         temp_resolved.append(res_entity)
-                        is_overlapped_or_contained = True # Mark current as handled
+                        is_overlapped_or_contained = True  # Mark current as handled
                         break
-                else: # No overlap
+                else:  # No overlap
                     temp_resolved.append(res_entity)
 
             if not is_overlapped_or_contained:
                 temp_resolved.append(current_entity)
 
-            resolved_entities = sorted(temp_resolved, key=lambda e: (e.start, -(e.end - e.start)))
-
+            resolved_entities = sorted(
+                temp_resolved, key=lambda e: (e.start, -(e.end - e.start))
+            )
 
         # Final pass to remove fully contained entities if a larger one exists
         final_entities = []
@@ -329,8 +386,10 @@ class PIIMasker:
                 if i == j:
                     continue
                 # If 'entity' is strictly contained within 'other_entity'
-                if other_entity.start <= entity.start and other_entity.end >= entity.end and \
-                   (other_entity.end - other_entity.start > entity.end - entity.start):
+                if (other_entity.start <= entity.start
+                        and other_entity.end >= entity.end
+                        and (other_entity.end - other_entity.start
+                             > entity.end - entity.start)):
                     is_contained = True
                     break
             if not is_contained:
@@ -338,7 +397,6 @@ class PIIMasker:
 
         return final_entities
 
-
     def mask_text(self, text: str) -> Tuple[str, List[Dict[str, Any]]]:
         """
         Mask PII entities in the text and return masked text and entity information
@@ -370,7 +428,6 @@ class PIIMasker:
 
         return "".join(new_text_parts), entity_info
 
-
     def process_email(self, email_text: str) -> Dict[str, Any]:
         """
         Process an email by detecting and masking PII entities.
@@ -378,56 +435,60 @@ class PIIMasker:
         """
         # Mask the email
         masked_email, entity_info = self.mask_text(email_text)
-        
+
         # Store the email in the SQLite database - only get back email_id now
         email_id = self.db.store_email(
             original_email=email_text,
             masked_email=masked_email,
             masked_entities=entity_info
         )
-        
+
         # Return the processed data with just the email_id
         return {
-            "input_email_body": email_text,  # Return original input for API compatibility
+            "input_email_body": email_text,  # Return original for API compatibility
             "list_of_masked_entities": entity_info,
             "masked_email": masked_email,
             "category_of_the_email": "",
             "email_id": email_id
         }
-    
-    def get_original_email(self, email_id: str, access_key: str) -> Optional[Dict[str, Any]]:
+
+    def get_original_email(
+        self, email_id: str, access_key: str
+    ) -> Optional[Dict[str, Any]]:
         """
         Retrieve the original email with PII using the email ID and access key.
-        
+
         Args:
             email_id: The ID of the stored email
             access_key: The security key for accessing the original email
-            
+
         Returns:
             The original email data or None if not found or access_key is invalid
         """
         return self.db.get_original_email(email_id, access_key)
-    
+
     def get_masked_email_by_id(self, email_id: str) -> Optional[Dict[str, Any]]:
         """
         Retrieve a masked email by its ID (without the original PII-containing email).
-        
+
         Args:
             email_id: The ID of the stored email
-            
+
         Returns:
             The masked email data or None if not found
         """
         return self.db.get_email_by_id(email_id)
-        
-    def get_original_by_masked_email(self, masked_email: str) -> Optional[Dict[str, Any]]:
+
+    def get_original_by_masked_email(
+        self, masked_email: str
+    ) -> Optional[Dict[str, Any]]:
         """
         Retrieve the original unmasked email using the masked email content.
-        
+
         Args:
             masked_email: The masked version of the email to search for
-            
+
         Returns:
             The original email data or None if not found
         """
-        return self.db.get_email_by_masked_content(masked_email)
+        return self.db.get_email_by_masked_content(masked_email)