Update classifier_manager/deberta_model.py

classifier_manager/deberta_model.py

@@ -1,34 +1,48 @@
 import os
-from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
+import torch
+from transformers import pipeline
 
 class PiiDebertaAnalyzer:
+    """
+    Implements the DeBERTa V3 model, widely recognized for winning the Kaggle PII Detection competition.
+    It uses a token-classification pipeline to detect PII entities.
+    """
     def __init__(self, model_name="lakshyakh93/deberta-large-finetuned-pii"):
-        self.model = None
-        self.available = False
+        self.device = 0 if torch.cuda.is_available() else -1
+        print(f"Loading DeBERTa Model on device: {'GPU' if self.device == 0 else 'CPU'}...")
         
         try:
-            print(f"Loading DeBERTa Model on device: CPU...")
-            
-            # Get HuggingFace token from environment
+            # Get HuggingFace token from environment (for private/gated models)
             hf_token = os.getenv('HF_TOKEN')
             
-            # Load with token if available
-            self.model = pipeline(
-                "token-classification",
-                model=model_name,
-                tokenizer=model_name,
-                device=-1,  # CPU
-                use_auth_token=hf_token,  # Add this line
+            # Aggregation strategy 'simple' merges B-TAG and I-TAG into a single entity automatically
+            self.pipe = pipeline(
+                "token-classification", 
+                model=model_name, 
+                device=self.device, 
+                token=hf_token,  # Use 'token' parameter (use_auth_token is deprecated)
                 aggregation_strategy="simple"
             )
-            
-            self.available = True
+            self.model_loaded = True
             print(f"✅ DeBERTa model '{model_name}' loaded successfully.")
-            
         except Exception as e:
-            print(f"Failed to load DeBERTa model: {e}")
+            print(f"❌ Failed to load DeBERTa model: {e}")
+            self.model_loaded = False
+
+        # Map Kaggle/DeBERTa labels to your App's standard labels
+        self.label_mapping = {
+            "NAME_STUDENT": "FIRST_NAME",
+            "EMAIL": "EMAIL",
+            "PHONE_NUM": "PHONE",
+            "STREET_ADDRESS": "LOCATION",
+            "ID_NUM": "SSN",
+            "USERNAME": "FIRST_NAME",
+            "URL_PERSONAL": "URL",
+            "PER": "FIRST_NAME",  # Generic NER label
+            "LOC": "LOCATION",    # Generic NER label
+            "ORG": "LOCATION"     # Mapping ORG to Location or ignore based on preference
+        }
 
-    
     def scan(self, text: str):
         if not self.model_loaded or not text:
             return []
@@ -56,4 +70,4 @@ class PiiDebertaAnalyzer:
             
         except Exception as e:
             print(f"DeBERTa scan error: {e}")
-            return []    
+            return []