Create use.py

use.py

@@ -0,0 +1,77 @@
+import torch
+from transformers import AutoTokenizer, AutoModelForTokenClassification
+
+class ShieldFilter:
+    def __init__(self, model_path="LH-Tech-AI/Shield-82M"):
+        print(f"Loading Shield-82M from {model_path}...")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModelForTokenClassification.from_pretrained(model_path)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device)
+        self.model.eval()
+
+        self.group_map = {
+            "FIRSTNAME": "PERSON", "MIDDLENAME": "PERSON", "LASTNAME": "PERSON",
+            "BUILDINGNUMBER": "ADDRESS", "STREET": "ADDRESS", "CITY": "ADDRESS", 
+            "STATE": "ADDRESS", "ZIPCODE": "ADDRESS", "SECONDARYADDRESS": "ADDRESS",
+            "EMAIL": "EMAIL", "PHONENUMBER": "PHONE", "PHONEIMEI": "PHONE",
+            "DATE": "DOB", "TIME": "DOB"
+        }
+
+    def protect(self, text):
+        inputs = self.tokenizer(
+            text, 
+            return_tensors="pt", 
+            truncation=True, 
+            max_length=512, 
+            return_offsets_mapping=True
+        ).to(self.device)
+        
+        offsets = inputs.pop("offset_mapping")[0].cpu().numpy()
+        
+        with torch.no_grad():
+            outputs = self.model(**inputs).logits
+            
+        predictions = torch.argmax(outputs, dim=2)[0].cpu().numpy()
+        id2label = self.model.config.id2label
+        
+        spans_to_replace = []
+        current_group = None
+        start_char = -1
+        last_char = -1
+        
+        for idx, (pred_id, offset) in enumerate(zip(predictions, offsets)):
+            if offset[0] == 0 and offset[1] == 0:
+                continue
+                
+            label = id2label[pred_id]
+            
+            if label == "O":
+                if current_group is not None:
+                    spans_to_replace.append((start_char, last_char, current_group))
+                    current_group = None
+            else:
+                group_tag = self.group_map.get(label, label)
+                
+                if current_group != group_tag:
+                    if current_group is not None:
+                        spans_to_replace.append((start_char, last_char, current_group))
+                    current_group = group_tag
+                    start_char = offset[0]
+                    
+                last_char = offset[1]
+                
+        if current_group is not None:
+            spans_to_replace.append((start_char, last_char, current_group))
+            
+        filtered_text = text
+        for start, end, tag in sorted(spans_to_replace, key=lambda x: x[0], reverse=True):
+            filtered_text = filtered_text[:start] + f"[{tag}]" + filtered_text[end:]
+            
+        return filtered_text
+
+if __name__ == "__main__":
+    shield = ShieldFilter()
+    sample = "My name is John Doe. Email: john@example.com. Phone: +49 123 45678."
+    print(f"Original: {sample}")
+    print(f"Protected: {shield.protect(sample)}")