1import torch
2from transformers import AutoModelForTokenClassification, XLMRobertaTokenizerFast
3import numpy as np
4from typing import List, Dict, Tuple
5
6class AzerbaijaniNER:
7 def __init__(self, model_name_or_path="LocalDoc/private_ner_azerbaijani_v2"):
8 self.model = AutoModelForTokenClassification.from_pretrained(model_name_or_path)
9 self.tokenizer = XLMRobertaTokenizerFast.from_pretrained("xlm-roberta-base")
10
11 self.model.eval()
12
13 self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
14 self.model.to(self.device)
15
16 self.id_to_label = {
17 0: "O",
18 1: "B-AGE", 2: "B-BUILDINGNUM", 3: "B-CITY", 4: "B-CREDITCARDNUMBER",
19 5: "B-DATE", 6: "B-DRIVERLICENSENUM", 7: "B-EMAIL", 8: "B-GIVENNAME",
20 9: "B-IDCARDNUM", 10: "B-PASSPORTNUM", 11: "B-STREET", 12: "B-SURNAME",
21 13: "B-TAXNUM", 14: "B-TELEPHONENUM", 15: "B-TIME", 16: "B-ZIPCODE",
22 17: "I-AGE", 18: "I-BUILDINGNUM", 19: "I-CITY", 20: "I-CREDITCARDNUMBER",
23 21: "I-DATE", 22: "I-DRIVERLICENSENUM", 23: "I-EMAIL", 24: "I-GIVENNAME",
24 25: "I-IDCARDNUM", 26: "I-PASSPORTNUM", 27: "I-STREET", 28: "I-SURNAME",
25 29: "I-TAXNUM", 30: "I-TELEPHONENUM", 31: "I-TIME", 32: "I-ZIPCODE"
26 }
27
28 self.entity_types = {
29 "AGE": "Age",
30 "BUILDINGNUM": "Building Number",
31 "CITY": "City",
32 "CREDITCARDNUMBER": "Credit Card Number",
33 "DATE": "Date",
34 "DRIVERLICENSENUM": "Driver License Number",
35 "EMAIL": "Email",
36 "GIVENNAME": "Given Name",
37 "IDCARDNUM": "ID Card Number",
38 "PASSPORTNUM": "Passport Number",
39 "STREET": "Street",
40 "SURNAME": "Surname",
41 "TAXNUM": "Tax ID Number",
42 "TELEPHONENUM": "Phone Number",
43 "TIME": "Time",
44 "ZIPCODE": "Zip Code"
45 }
46
47 def predict(self, text: str, max_length: int = 512) -> List[Dict]:
48 text = text.lower()
49
50 inputs = self.tokenizer(
51 text,
52 return_tensors="pt",
53 max_length=max_length,
54 padding="max_length",
55 truncation=True,
56 return_offsets_mapping=True
57 )
58
59 offset_mapping = inputs.pop("offset_mapping").numpy()[0]
60
61 inputs = {k: v.to(self.device) for k, v in inputs.items()}
62
63 with torch.no_grad():
64 outputs = self.model(**inputs)
65 predictions = outputs.logits.argmax(dim=2)
66
67 predictions = predictions[0].cpu().numpy()
68
69 entities = []
70 current_entity = None
71
72 for idx, (offset, pred_id) in enumerate(zip(offset_mapping, predictions)):
73 if offset[0] == 0 and offset[1] == 0:
74 continue
75
76 pred_label = self.id_to_label[pred_id]
77
78 if pred_label.startswith("B-"):
79 if current_entity:
80 entities.append(current_entity)
81
82 entity_type = pred_label[2:]
83 current_entity = {
84 "label": entity_type,
85 "name": self.entity_types.get(entity_type, entity_type),
86 "start": int(offset[0]),
87 "end": int(offset[1]),
88 "value": text[offset[0]:offset[1]]
89 }
90
91 elif pred_label.startswith("I-") and current_entity is not None:
92 entity_type = pred_label[2:]
93
94 if entity_type == current_entity["label"]:
95 current_entity["end"] = int(offset[1])
96 current_entity["value"] = text[current_entity["start"]:current_entity["end"]]
97 else:
98 entities.append(current_entity)
99 current_entity = None
100
101 elif pred_label == "O" and current_entity is not None:
102 entities.append(current_entity)
103 current_entity = None
104
105 if current_entity:
106 entities.append(current_entity)
107
108 return entities
109
110 def anonymize_text(self, text: str, replacement_char: str = "X") -> Tuple[str, List[Dict]]:
111 entities = self.predict(text)
112
113 if not entities:
114 return text, []
115
116 entities.sort(key=lambda x: x["start"], reverse=True)
117
118 anonymized_text = text
119 for entity in entities:
120 start = entity["start"]
121 end = entity["end"]
122 length = end - start
123 anonymized_text = anonymized_text[:start] + replacement_char * length + anonymized_text[end:]
124
125 entities.sort(key=lambda x: x["start"])
126
127 return anonymized_text, entities
128
129 def highlight_entities(self, text: str) -> str:
130 entities = self.predict(text)
131
132 if not entities:
133 return text
134
135 entities.sort(key=lambda x: x["start"], reverse=True)
136
137 highlighted_text = text
138 for entity in entities:
139 start = entity["start"]
140 end = entity["end"]
141 entity_value = entity["value"]
142 entity_type = entity["name"]
143
144 highlighted_text = (
145 highlighted_text[:start] +
146 f"[{entity_type}: {entity_value}]" +
147 highlighted_text[end:]
148 )
149
150 return highlighted_text
151
152if __name__ == "__main__":
153 ner = AzerbaijaniNER()
154
155 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 ?"""
156
157 print("=== Original Text ===")
158 print(test_text)
159 print("\n=== Found Entities ===")
160
161 entities = ner.predict(test_text)
162 for entity in entities:
163 print(f"{entity['name']}: {entity['value']} (positions {entity['start']}-{entity['end']})")
164
165 print("\n=== Text with Highlighted Entities ===")
166 highlighted_text = ner.highlight_entities(test_text)
167 print(highlighted_text)
168
169 print("\n=== Anonymized Text ===")
170 anonymized_text, _ = ner.anonymize_text(test_text)
171 print(anonymized_text)
