Update README.md

README.md

@@ -101,89 +101,197 @@ To use the model for spell correction:
 
 ```python
 import torch
-from transformers import AutoTokenizer, AutoModelForTokenClassification
-
-model_id = "LocalDoc/private_ner_azerbaijani_v2"
-
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForTokenClassification.from_pretrained(model_id)
-
-test_text = (
-    "Salam, mənim adım Əli Hüseynovdur. Doğum tarixim 15.05.1990-dır. Bakı şəhərində, Nizami küçəsində, 25/31 ünvanında yaşayıram. Telefon nömrəm +994552345678-dir."
-)
-
-inputs = tokenizer(test_text, return_tensors="pt", return_offsets_mapping=True)
-
-offset_mapping = inputs.pop("offset_mapping")
-
-with torch.no_grad():
-    outputs = model(**inputs)
-
-predictions = torch.argmax(outputs.logits, dim=2)
-
-tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
-offset_mapping = offset_mapping[0].tolist()
-predicted_labels = [model.config.id2label[pred.item()] for pred in predictions[0]]
-word_ids = inputs.word_ids(batch_index=0)
-
-aggregated = []
-prev_word_id = None
-for idx, word_id in enumerate(word_ids):
-    if word_id is None:
-        continue
-    if word_id != prev_word_id:
-        aggregated.append({
-            "word_id": word_id,
-            "tokens": [tokens[idx]],
-            "offsets": [offset_mapping[idx]],
-            "label": predicted_labels[idx]
-        })
-    else:
-        aggregated[-1]["tokens"].append(tokens[idx])
-        aggregated[-1]["offsets"].append(offset_mapping[idx])
-    prev_word_id = word_id
-
-entities = []
-current_entity = None
-for word in aggregated:
-    if word["label"] == "O":
-        if current_entity is not None:
-            entities.append(current_entity)
-            current_entity = None
-    else:
-        if current_entity is None:
-            current_entity = {
-                "type": word["label"],
-                "start": word["offsets"][0][0],
-                "end": word["offsets"][-1][1]
-            }
-        else:
-            if word["label"] == current_entity["type"]:
-                current_entity["end"] = word["offsets"][-1][1]
-            else:
-                entities.append(current_entity)
+from transformers import AutoModelForTokenClassification, XLMRobertaTokenizerFast
+import numpy as np
+from typing import List, Dict, Tuple
+
+class AzerbaijaniNER:
+    def __init__(self, model_name_or_path="LocalDoc/private_ner_azerbaijani_v2"):
+        self.model = AutoModelForTokenClassification.from_pretrained(model_name_or_path)
+        self.tokenizer = XLMRobertaTokenizerFast.from_pretrained("xlm-roberta-base")
+        
+        self.model.eval()
+        
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device)
+        
+        self.id_to_label = {
+            0: "O",
+            1: "B-AGE", 2: "B-BUILDINGNUM", 3: "B-CITY", 4: "B-CREDITCARDNUMBER",
+            5: "B-DATE", 6: "B-DRIVERLICENSENUM", 7: "B-EMAIL", 8: "B-GIVENNAME",
+            9: "B-IDCARDNUM", 10: "B-PASSPORTNUM", 11: "B-STREET", 12: "B-SURNAME",
+            13: "B-TAXNUM", 14: "B-TELEPHONENUM", 15: "B-TIME", 16: "B-ZIPCODE",
+            17: "I-AGE", 18: "I-BUILDINGNUM", 19: "I-CITY", 20: "I-CREDITCARDNUMBER",
+            21: "I-DATE", 22: "I-DRIVERLICENSENUM", 23: "I-EMAIL", 24: "I-GIVENNAME", 
+            25: "I-IDCARDNUM", 26: "I-PASSPORTNUM", 27: "I-STREET", 28: "I-SURNAME",
+            29: "I-TAXNUM", 30: "I-TELEPHONENUM", 31: "I-TIME", 32: "I-ZIPCODE"
+        }
+        
+        self.entity_types = {
+            "AGE": "Age",
+            "BUILDINGNUM": "Building Number",
+            "CITY": "City",
+            "CREDITCARDNUMBER": "Credit Card Number",
+            "DATE": "Date",
+            "DRIVERLICENSENUM": "Driver License Number",
+            "EMAIL": "Email",
+            "GIVENNAME": "Given Name",
+            "IDCARDNUM": "ID Card Number",
+            "PASSPORTNUM": "Passport Number",
+            "STREET": "Street",
+            "SURNAME": "Surname",
+            "TAXNUM": "Tax ID Number",
+            "TELEPHONENUM": "Phone Number",
+            "TIME": "Time",
+            "ZIPCODE": "Zip Code"
+        }
+    
+    def predict(self, text: str, max_length: int = 512) -> List[Dict]:
+        text = text.lower()
+        
+        inputs = self.tokenizer(
+            text,
+            return_tensors="pt",
+            max_length=max_length,
+            padding="max_length",
+            truncation=True,
+            return_offsets_mapping=True
+        )
+        
+        offset_mapping = inputs.pop("offset_mapping").numpy()[0]
+        
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+        
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            predictions = outputs.logits.argmax(dim=2)
+        
+        predictions = predictions[0].cpu().numpy()
+        
+        entities = []
+        current_entity = None
+        
+        for idx, (offset, pred_id) in enumerate(zip(offset_mapping, predictions)):
+            if offset[0] == 0 and offset[1] == 0:
+                continue
+                
+            pred_label = self.id_to_label[pred_id]
+            
+            if pred_label.startswith("B-"):
+                if current_entity:
+                    entities.append(current_entity)
+                
+                entity_type = pred_label[2:]
                 current_entity = {
-                    "type": word["label"],
-                    "start": word["offsets"][0][0],
-                    "end": word["offsets"][-1][1]
+                    "label": entity_type,
+                    "name": self.entity_types.get(entity_type, entity_type),
+                    "start": int(offset[0]),
+                    "end": int(offset[1]),
+                    "value": text[offset[0]:offset[1]]
                 }
-if current_entity is not None:
-    entities.append(current_entity)
-
-for entity in entities:
-    entity["text"] = test_text[entity["start"]:entity["end"]]
-
-for entity in entities:
-    print(entity)
+            
+            elif pred_label.startswith("I-") and current_entity is not None:
+                entity_type = pred_label[2:]
+                
+                if entity_type == current_entity["label"]:
+                    current_entity["end"] = int(offset[1])
+                    current_entity["value"] = text[current_entity["start"]:current_entity["end"]]
+                else:
+                    entities.append(current_entity)
+                    current_entity = None
+            
+            elif pred_label == "O" and current_entity is not None:
+                entities.append(current_entity)
+                current_entity = None
+        
+        if current_entity:
+            entities.append(current_entity)
+        
+        return entities
+    
+    def anonymize_text(self, text: str, replacement_char: str = "X") -> Tuple[str, List[Dict]]:
+        entities = self.predict(text)
+        
+        if not entities:
+            return text, []
+        
+        entities.sort(key=lambda x: x["start"], reverse=True)
+        
+        anonymized_text = text
+        for entity in entities:
+            start = entity["start"]
+            end = entity["end"]
+            length = end - start
+            anonymized_text = anonymized_text[:start] + replacement_char * length + anonymized_text[end:]
+        
+        entities.sort(key=lambda x: x["start"])
+        
+        return anonymized_text, entities
+
+    def highlight_entities(self, text: str) -> str:
+        entities = self.predict(text)
+        
+        if not entities:
+            return text
+        
+        entities.sort(key=lambda x: x["start"], reverse=True)
+        
+        highlighted_text = text
+        for entity in entities:
+            start = entity["start"]
+            end = entity["end"]
+            entity_value = entity["value"]
+            entity_type = entity["name"]
+            
+            highlighted_text = (
+                highlighted_text[:start] + 
+                f"[{entity_type}: {entity_value}]" + 
+                highlighted_text[end:]
+            )
+        
+        return highlighted_text
+
+if __name__ == "__main__":
+    ner = AzerbaijaniNER()
+    
+    test_text = """Salam, mənim adım Əli Hüseynovdu. Doğum tarixim 15.05.1990-dır. Bakı şəhərində, 28 may küçəsi 4 ünvanında yaşayıram. Telefon nömrəm +994552345678-dir. Mən 4169741358254152 nömrəli kartdan ödəniş etmişəm. Sifarişim nə vaxt çatdırılcaq ?"""
+    
+    print("=== Original Text ===")
+    print(test_text)
+    print("\n=== Found Entities ===")
+    
+    entities = ner.predict(test_text)
+    for entity in entities:
+        print(f"{entity['name']}: {entity['value']} (positions {entity['start']}-{entity['end']})")
+    
+    print("\n=== Text with Highlighted Entities ===")
+    highlighted_text = ner.highlight_entities(test_text)
+    print(highlighted_text)
+    
+    print("\n=== Anonymized Text ===")
+    anonymized_text, _ = ner.anonymize_text(test_text)
+    print(anonymized_text)
 ```
 
-```json
-{'type': 'FIRSTNAME', 'start': 18, 'end': 21, 'text': 'Əli'}
-{'type': 'LASTNAME', 'start': 22, 'end': 34, 'text': 'Hüseynovdur.'}
-{'type': 'DOB', 'start': 49, 'end': 64, 'text': '15.05.1990-dır.'}
-{'type': 'STREET', 'start': 81, 'end': 87, 'text': 'Nizami'}
-{'type': 'BUILDINGNUMBER', 'start': 99, 'end': 104, 'text': '25/31'}
-{'type': 'PHONENUMBER', 'start': 141, 'end': 159, 'text': '+994552345678-dir.'}
+```
+=== Original Text ===
+Salam, mənim adım Əli Hüseynovdu. Doğum tarixim 15.05.1990-dır. Bakı şəhərində, 28 may küçəsi 4 ünvanında yaşayıram. Telefon nömrəm +994552345678-dir. Mən 4169741358254152 nömrəli kartdan ödəniş etmişəm. Sifarişim nə vaxt çatdırılcaq ?
+
+=== Found Entities ===
+Given Name: əli (positions 18-21)
+Surname: hüseynov (positions 22-30)
+Date: 15.05.1990 (positions 48-58)
+City: bakı (positions 64-68)
+Street: 28 may küçəsi (positions 80-93)
+Building Number: 4 (positions 94-95)
+Phone Number: +994552345678 (positions 132-145)
+Credit Card Number: 4169741358254152 (positions 155-171)
+
+=== Text with Highlighted Entities ===
+Salam, mənim adım [Given Name: əli] [Surname: hüseynov]du. Doğum tarixim [Date: 15.05.1990]-dır. [City: bakı] şəhərində, [Street: 28 may küçəsi] [Building Number: 4] ünvanında yaşayıram. Telefon nömrəm [Phone Number: +994552345678]-dir. Mən [Credit Card Number: 4169741358254152] nömrəli kartdan ödəniş etmişəm. Sifarişim nə vaxt çatdırılcaq ?
+
+=== Anonymized Text ===
+Salam, mənim adım XXX XXXXXXXXdu. Doğum tarixim XXXXXXXXXX-dır. XXXX şəhərində, XXXXXXXXXXXXX X ünvanında yaşayıram. Telefon nömrəm XXXXXXXXXXXXX-dir. Mən XXXXXXXXXXXXXXXX nömrəli kartdan ödəniş etmişəm. Sifarişim nə vaxt çatdırılcaq ?
 ```