PII-detect / app.py
quynong's picture
add model mmBERT
8a2af25
Raw
History Blame Contribute Delete
47.5 kB
import gradio as gr
import json
import os
import re
import gc
import ipaddress
import unicodedata
from datetime import datetime
from typing import Dict, List, Optional, Tuple, Callable, Pattern
from collections import defaultdict
from dataclasses import dataclass
from urllib.parse import urlparse
import torch
from gliner2 import GLiNER2
from gliner2.processor import WhitespaceTokenSplitter
from transformers import pipeline
from huggingface_hub import snapshot_download
# ==========================================
# 1. Custom CSS
# ==========================================
custom_css = """
.pii-entity {
font-weight: bold;
background-color: rgba(250, 204, 21, 0.3);
border-bottom: 2px solid #eab308;
border-radius: 4px;
padding: 2px 4px;
position: relative;
cursor: help;
transition: background-color 0.2s;
}
.pii-entity:hover {
background-color: rgba(250, 204, 21, 0.6);
}
.pii-entity .tooltip-label {
visibility: hidden;
background-color: #1f2937;
color: #f9fafb;
text-align: center;
border-radius: 6px;
padding: 4px 8px;
position: absolute;
z-index: 50;
bottom: 125%;
left: 50%;
transform: translateX(-50%);
opacity: 0;
transition: opacity 0.2s, bottom 0.2s;
font-size: 0.75rem;
font-family: monospace;
font-weight: normal;
white-space: nowrap;
box-shadow: 0 4px 6px -1px rgba(0,0,0,0.1), 0 2px 4px -1px rgba(0,0,0,0.06);
}
.pii-entity .tooltip-label::after {
content: "";
position: absolute;
top: 100%;
left: 50%;
margin-left: -5px;
border-width: 5px;
border-style: solid;
border-color: #1f2937 transparent transparent transparent;
}
.pii-entity:hover .tooltip-label {
visibility: visible;
opacity: 1;
bottom: 135%;
}
.text-container {
font-size: 1rem;
line-height: 1.8;
padding: 1.5rem;
background: #ffffff;
border: 1px solid #e5e7eb;
border-radius: 8px;
min-height: 120px;
}
.dark .text-container {
background: #1f2937;
border-color: #374151;
color: #e5e7eb;
}
.model-header-gliner {
background: linear-gradient(135deg, #3b82f6, #1d4ed8);
color: white;
padding: 8px 16px;
border-radius: 8px 8px 0 0;
font-weight: bold;
font-size: 0.9rem;
margin-bottom: 0;
}
.model-header-xlmbert {
background: linear-gradient(135deg, #8b5cf6, #5b21b6);
color: white;
padding: 8px 16px;
border-radius: 8px 8px 0 0;
font-weight: bold;
font-size: 0.9rem;
margin-bottom: 0;
}
.model-header-mmbert {
background: linear-gradient(135deg, #f97316, #c2410c);
color: white;
padding: 8px 16px;
border-radius: 8px 8px 0 0;
font-weight: bold;
font-size: 0.9rem;
margin-bottom: 0;
}
.batch-sample {
border: 1px solid #e5e7eb;
border-radius: 10px;
margin-bottom: 24px;
overflow: hidden;
box-shadow: 0 1px 3px rgba(0,0,0,0.06);
}
.dark .batch-sample {
border-color: #374151;
}
.batch-sample-header {
background: linear-gradient(135deg, #0ea5e9, #0284c7);
color: white;
padding: 10px 16px;
font-weight: bold;
font-size: 0.95rem;
display: flex;
justify-content: space-between;
align-items: center;
}
.batch-sample-input {
padding: 12px 16px;
background: #f8fafc;
border-bottom: 1px solid #e5e7eb;
font-size: 0.85rem;
color: #475569;
white-space: pre-wrap;
max-height: 120px;
overflow-y: auto;
}
.dark .batch-sample-input {
background: #111827;
border-color: #374151;
color: #94a3b8;
}
.batch-models-row {
display: grid;
grid-template-columns: 1fr 1fr;
gap: 0;
}
.batch-model-col {
padding: 12px 16px;
}
.batch-model-col:first-child {
border-right: 1px solid #e5e7eb;
}
.dark .batch-model-col:first-child {
border-color: #374151;
}
.batch-model-label {
font-size: 0.75rem;
font-weight: 600;
text-transform: uppercase;
letter-spacing: 0.05em;
margin-bottom: 8px;
padding: 3px 8px;
border-radius: 4px;
display: inline-block;
}
.batch-model-label.gliner {
background: #dbeafe;
color: #1d4ed8;
}
.batch-model-label.bert {
background: #ede9fe;
color: #5b21b6;
}
.batch-model-label.mmbert {
background: #ffedd5;
color: #c2410c;
}
.dark .batch-model-label.gliner {
background: #1e3a5f;
color: #93c5fd;
}
.dark .batch-model-label.bert {
background: #2e1065;
color: #c4b5fd;
}
.dark .batch-model-label.mmbert {
background: #431407;
color: #fdba74;
}
.batch-entity-tags {
display: flex;
flex-wrap: wrap;
gap: 4px;
margin-top: 6px;
}
.batch-entity-tag {
font-size: 0.72rem;
padding: 2px 6px;
border-radius: 4px;
background: #fef3c7;
color: #92400e;
font-family: monospace;
border: 1px solid #fde68a;
}
.dark .batch-entity-tag {
background: #422006;
color: #fde68a;
border-color: #78350f;
}
.batch-summary {
background: linear-gradient(135deg, #f0fdf4, #dcfce7);
border: 1px solid #bbf7d0;
border-radius: 10px;
padding: 16px 20px;
margin-bottom: 20px;
font-size: 0.9rem;
}
.dark .batch-summary {
background: linear-gradient(135deg, #052e16, #064e3b);
border-color: #166534;
color: #bbf7d0;
}
.batch-progress {
text-align: center;
padding: 40px;
color: #6b7280;
font-size: 1.1rem;
}
"""
# ==========================================
# 2. Cấu hình Label & Tokenizer (Cho GLiNER)
# ==========================================
_IMPROVED_PATTERN = re.compile(
r"""(?:https?://[^\s]+|www\.[^\s]+)|[a-z0-9._%+-]+@[a-z0-9.-]+\.[a-z]{2,}|@[a-z0-9_]+|(?:[0-9a-f]{1,4}:){7}[0-9a-f]{1,4}|(?:[0-9a-f]{1,4}:){1,7}:|:(?::[0-9a-f]{1,4}){1,7}|[a-z0-9_-]{10,}\.[a-z0-9_-]{10,}\.[a-z0-9_-]{10,}|\w+(?:[-_]\w+)*|\S""",
re.VERBOSE | re.IGNORECASE,
)
WhitespaceTokenSplitter._PATTERN = _IMPROVED_PATTERN
LABELS = {
"PREFIX": "Personal titles and prefixes",
"PERSONAL_NAME": "Full, given, and family names",
"SEX": "Biological sex and gender identities",
"AGE": "Age values",
"DOB": "Dates of birth",
"PHONE_NUMBER": "Telephone and mobile numbers",
"EMAIL": "Email addresses",
"LOCATION": "Geographic regions and cities",
"STREET_ADDRESS": "Specific local and street addresses",
"ZIPCODE": "Postal and ZIP codes",
"GPS_COORDINATE": "GPS coordinates and latitude/longitude",
"USERNAME": "Usernames and digital aliases",
"PASSWORD": "Passwords and login credentials",
"PIN": "Personal Identification Numbers (PINs)",
"URL": "URLs and web addresses",
"IP_ADDRESS": "IP addresses",
"ACCOUNT_NUMBER": "Bank account numbers",
"AMOUNT": "Monetary amounts and values",
"CREDIT_CARD_ISSUER": "Credit card issuers and brands",
"CREDIT_CARD_NUMBER": "Credit card numbers",
"CREDIT_CARD_CVV": "Credit card verification values (CVV/CVC)",
"IBAN": "International Bank Account Numbers (IBAN)",
"BIC_SWIFT": "BIC and SWIFT codes",
"CRYPTO_ADDRESS": "Cryptocurrency addresses",
"OCCUPATION": "Occupations, job titles, and industries",
"SSN_CCCD": "Social Security and national identification numbers",
"PASSPORT_NUM": "Passport numbers",
"DRIVER_LICENSE": "Driver's license numbers",
"TAX_ID": "Tax identification numbers",
"DATE": "General calendar dates",
"TIME": "Specific times of day",
"MARITAL_STATUS": "Marital and legal relationship status",
"RELIGION": "Religious affiliations and beliefs",
"ETHNICITY": "Ethnicities and cultural backgrounds",
"TRADE_UNION_INFO": "Trade union membership and affiliations",
"NATIONALITY": "Nationalities and citizenships",
"HEALTH_INSURANCE": "Health insurance numbers and identifiers",
"HEALTH_STATUS": "Health status and medical conditions"
}
# ==========================================
# 2b. Preprocessing — clean_text_with_mapping
# ==========================================
_STRIP_CHARS = frozenset(["\u00ad", "\u200b", "\u200c", "\u200d", "\u2060", "\ufeff", "\ufffd"])
_C1_RE = re.compile(r"[\x80-\x9f]")
_C0_RE = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]")
_MULTI_SPACE_RE = re.compile(r"[^\S\n]+")
def clean_text_with_mapping(text: str) -> tuple:
if not text:
return text, []
text = unicodedata.normalize("NFC", text)
cleaned_chars = []
char_map = []
orig_pos = 0
for ch in text:
if ch in _STRIP_CHARS:
orig_pos += 1
continue
if _C0_RE.match(ch):
orig_pos += 1
continue
if _C1_RE.match(ch):
cleaned_chars.append(" ")
char_map.append(orig_pos)
orig_pos += 1
continue
cleaned_chars.append(ch)
char_map.append(orig_pos)
orig_pos += 1
cleaned = "".join(cleaned_chars)
cleaned = _MULTI_SPACE_RE.sub(" ", cleaned).strip()
final_map = []
i = 0
for ch in cleaned:
while i < len(char_map) and (text[char_map[i]] != ch and not (ch == " " and text[char_map[i]].isspace())):
i += 1
final_map.append(char_map[i] if i < len(char_map) else len(text))
i += 1
final_map.append(len(text))
return cleaned, final_map
# ==========================================
# 2c. Rule-based detector
# ==========================================
@dataclass(frozen=True)
class RegexRule:
entity_type: str
pattern: Pattern[str]
validator: Callable[[str], bool]
normalizer: Callable[[str], str] | None = None
MONTH_PATTERN = (
r"(?:Jan(?:uary)?|Feb(?:ruary)?|Mar(?:ch)?|Apr(?:il)?|May|"
r"Jun(?:e)?|Jul(?:y)?|Aug(?:ust)?|Sep(?:t(?:ember)?)?|"
r"Oct(?:ober)?|Nov(?:ember)?|Dec(?:ember)?)"
)
MONTH_PATTERN_VI = r"(?:tháng\s*(?:1|2|3|4|5|6|7|8|9|10|11|12|một|hai|ba|tư|năm|sáu|bảy|tám|chín|mười|mười\s*một|mười\s*hai))"
MONTH_PATTERN_DE = r"(?:Januar|Februar|März|April|Mai|Juni|Juli|August|September|Oktober|November|Dezember)"
DATE_TOKEN_PATTERN = (
rf"(?:\d{{4}}[-/]\d{{1,2}}[-/]\d{{1,2}}|"
rf"\d{{1,2}}[-/]\d{{1,2}}[-/]\d{{2,4}}|"
rf"\d{{1,2}}\.\d{{1,2}}\.\d{{2,4}}|"
rf"{MONTH_PATTERN}\s+\d{{1,2}}(?:st|nd|rd|th)?(?:,\s*\d{{2,4}})?|"
rf"\d{{1,2}}(?:st|nd|rd|th)?\s+(?:of\s+)?{MONTH_PATTERN}(?:,\s*\d{{2,4}})?|"
rf"(?:ngày\s+)?\d{{1,2}}\s+{MONTH_PATTERN_VI}(?:\s+(?:năm\s+)?\d{{2,4}})?|"
rf"\d{{1,2}}\.\s*{MONTH_PATTERN_DE}(?:\s+\d{{2,4}})?)"
)
def normalize_space(value: str) -> str:
return re.sub(r"\s+", " ", value).strip()
def strip_number_separators(value: str) -> str:
return re.sub(r"[\s-]", "", value).strip()
def normalize_phone(value: str) -> str:
compact = re.sub(r"[().\s-]", "", value)
if value.strip().startswith("+"):
return "+" + compact.lstrip("+")
return compact
def normalize_iban(value: str) -> str:
return re.sub(r"[^A-Z0-9]", "", value.upper())
def is_valid_email(value: str) -> bool:
if ".." in value:
return False
local, _, domain = value.rpartition("@")
if not local or not domain or "." not in domain:
return False
return not domain.startswith(".") and not domain.endswith(".")
def is_valid_url(value: str) -> bool:
candidate = value
if candidate.lower().startswith("www."):
candidate = "https://" + candidate
parsed = urlparse(candidate)
return bool(parsed.scheme and parsed.netloc and "." in parsed.netloc)
def is_valid_phone(value: str) -> bool:
digits = re.sub(r"\D", "", value)
has_separator = any(ch in value for ch in " +-().")
groups = re.findall(r"\d+", value)
if not (8 <= len(digits) <= 15 and has_separator):
return False
if any(len(group) > 4 for group in groups):
return False
if groups and len(groups[-1]) < 3:
return False
if not value.strip().startswith("+") and len(digits) > 11:
return False
if re.fullmatch(r"\d{4}-\d{1,2}-\d{1,2}(?:\s+\d{1,2})?", value):
return False
return True
def is_valid_ipv4(value: str) -> bool:
try:
return isinstance(ipaddress.ip_address(value), ipaddress.IPv4Address)
except ValueError:
return False
def is_valid_ipv6(value: str) -> bool:
try:
return isinstance(ipaddress.ip_address(value), ipaddress.IPv6Address)
except ValueError:
return False
def parse_date_text(value: str) -> bool:
candidate = normalize_space(value.replace(" of ", " "))
formats = [
"%Y-%m-%d", "%Y/%m/%d", "%d-%m-%Y", "%d/%m/%Y", "%m-%d-%Y", "%m/%d/%Y",
"%d-%m-%y", "%d/%m/%y", "%m-%d-%y", "%m/%d/%y",
"%B %d, %Y", "%B %d, %y", "%B %d %Y", "%B %d",
"%b %d, %Y", "%b %d, %y", "%b %d %Y", "%b %d",
"%d %B %Y", "%d %B", "%d %b %Y", "%d %b",
]
candidate = re.sub(r"(\d)(st|nd|rd|th)\b", r"\1", candidate, flags=re.IGNORECASE)
for fmt in formats:
try:
datetime.strptime(candidate, fmt)
return True
except ValueError:
continue
return False
def is_valid_time(value: str) -> bool:
candidate = normalize_space(value).upper().replace(".", "")
for fmt in ["%H:%M", "%I:%M %p", "%I %p"]:
try:
datetime.strptime(candidate, fmt)
return True
except ValueError:
continue
return False
def luhn_check(value: str) -> bool:
digits = re.sub(r"\D", "", value)
if not 13 <= len(digits) <= 19:
return False
total = 0
for index, digit in enumerate(digits[::-1]):
number = int(digit)
if index % 2 == 1:
number *= 2
if number > 9:
number -= 9
total += number
return total % 10 == 0
IBAN_LENGTHS = {
"AD":24,"AE":23,"AL":28,"AT":20,"AZ":28,"BA":20,"BE":16,"BG":22,"BH":22,"BR":29,
"BY":28,"CH":21,"CR":22,"CY":28,"CZ":24,"DE":22,"DK":18,"DO":28,"EE":20,"EG":29,
"ES":24,"FI":18,"FO":18,"FR":27,"GB":22,"GE":22,"GI":23,"GL":18,"GR":27,"GT":28,
"HR":21,"HU":28,"IE":22,"IL":23,"IQ":23,"IS":26,"IT":27,"JO":30,"KW":30,"KZ":20,
"LB":28,"LC":32,"LI":21,"LT":20,"LU":20,"LV":21,"MC":27,"MD":24,"ME":22,"MK":19,
"MR":27,"MT":31,"MU":30,"NL":18,"NO":15,"PK":24,"PL":28,"PS":29,"PT":25,"QA":29,
"RO":24,"RS":22,"SA":24,"SC":31,"SE":24,"SI":19,"SK":24,"SM":27,"ST":25,"SV":28,
"TL":23,"TN":24,"TR":26,"UA":29,"VA":22,"VG":24,"XK":20,
}
def iban_checksum(value: str) -> bool:
compact = normalize_iban(value)
if not 15 <= len(compact) <= 34:
return False
expected = IBAN_LENGTHS.get(compact[:2])
if expected and len(compact) != expected:
return False
rearranged = compact[4:] + compact[:4]
numeric = "".join(str(int(ch, 36)) for ch in rearranged)
return int(numeric) % 97 == 1
def is_valid_zipcode(value: str) -> bool:
c = value.strip()
return bool(
re.fullmatch(r"\d{5}(?:-\d{4})?", c)
or re.fullmatch(r"[A-Z]{1,2}\d[A-Z0-9]?\s*\d[A-Z]{2}", c, re.IGNORECASE)
or re.fullmatch(r"\d{5}", c)
or re.fullmatch(r"\d{6}", c)
or re.fullmatch(r"[A-Z]\d[A-Z]\s*\d[A-Z]\d", c, re.IGNORECASE)
or re.fullmatch(r"\d{3}-\d{4}", c)
or re.fullmatch(r"\d{4}", c)
)
def is_valid_amount(value: str) -> bool:
digits = re.sub(r"[^\d.]", "", value)
if not digits or digits.count(".") > 1:
return False
try:
return float(digits) > 0
except ValueError:
return False
def is_valid_crypto_address(value: str) -> bool:
c = value.strip()
return bool(
re.fullmatch(r"[13][A-HJ-NP-Za-km-z1-9]{25,34}", c)
or re.fullmatch(r"bc1[a-z0-9]{25,89}", c)
or re.fullmatch(r"0x[0-9a-fA-F]{40}", c)
or re.fullmatch(r"[LM3][A-HJ-NP-Za-km-z1-9]{26,33}", c)
or re.fullmatch(r"r[0-9A-HJ-NP-Za-km-z]{24,34}", c)
)
def is_valid_age(value: str) -> bool:
digits = re.sub(r"\D", "", value)
return bool(digits) and 0 <= int(digits) <= 150
def is_valid_ssn(value: str) -> bool:
digits = re.sub(r"\D", "", value)
if len(digits) != 9:
return False
area, group, serial = digits[:3], digits[3:5], digits[5:]
if area in {"000", "666"} or area.startswith("9"):
return False
return group != "00" and serial != "0000"
def is_valid_cccd(value: str) -> bool:
digits = re.sub(r"\D", "", value)
return len(digits) in {9, 12}
RULE_PATTERNS: Dict[str, RegexRule] = {
"EMAIL": RegexRule(
entity_type="EMAIL",
pattern=re.compile(r"(?<![\w.+-])[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,63}(?![\w-])", re.IGNORECASE),
validator=is_valid_email,
),
"URL": RegexRule(
entity_type="URL",
pattern=re.compile(r"(?:(?:https?://)|(?:www\.))[^\s<>()]+", re.IGNORECASE),
validator=is_valid_url,
),
"PHONE_NUMBER": RegexRule(
entity_type="PHONE_NUMBER",
pattern=re.compile(r"(?<!\w)(?:\+?\d{1,3}[\s().-]*)?(?:\(?\d{2,4}\)?[\s().-]*){1,3}\d{3,4}(?!\w)"),
validator=is_valid_phone,
normalizer=normalize_phone,
),
"IP_ADDRESS_V4": RegexRule(
entity_type="IP_ADDRESS",
pattern=re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"),
validator=is_valid_ipv4,
),
"IP_ADDRESS_V6": RegexRule(
entity_type="IP_ADDRESS",
pattern=re.compile(r"(?<![:\w])(?:[A-F0-9]{1,4}:){2,7}[A-F0-9]{1,4}(?![:\w])", re.IGNORECASE),
validator=is_valid_ipv6,
),
"DATE": RegexRule(
entity_type="DATE",
pattern=re.compile(rf"\b{DATE_TOKEN_PATTERN}\b", re.IGNORECASE),
validator=parse_date_text,
),
"TIME": RegexRule(
entity_type="TIME",
pattern=re.compile(
r"\b(?:(?:[01]?\d|2[0-3]):[0-5]\d(?::[0-5]\d)?(?:\s?(?:AM|PM))?|(?:1[0-2]|0?[1-9])\s?(?:AM|PM))\b",
re.IGNORECASE,
),
validator=is_valid_time,
),
"CREDIT_CARD_NUMBER": RegexRule(
entity_type="CREDIT_CARD_NUMBER",
pattern=re.compile(r"(?<!\d)(?:\d[ -]?){13,19}(?!\d)"),
validator=luhn_check,
normalizer=strip_number_separators,
),
"IBAN": RegexRule(
entity_type="IBAN",
pattern=re.compile(r"\b[A-Z]{2}\d{2}(?:[ -]?[A-Z0-9]){11,30}\b", re.IGNORECASE),
validator=iban_checksum,
normalizer=normalize_iban,
),
"SSN_CCCD_US": RegexRule(
entity_type="SSN_CCCD",
pattern=re.compile(r"\b(?!000|666|9\d\d)\d{3}-\d{2}-\d{4}\b"),
validator=is_valid_ssn,
),
"SSN_CCCD_VN": RegexRule(
entity_type="SSN_CCCD",
pattern=re.compile(r"\b0\d{2}[0-3]\d{8}\b"),
validator=is_valid_cccd,
),
"ZIPCODE": RegexRule(
entity_type="ZIPCODE",
pattern=re.compile(
r"(?<=[\s,])\d{5}(?:-\d{4})?(?=[\s,.\n]|$)"
r"|(?<=[\s,])[A-Z]{1,2}\d[A-Z0-9]?\s*\d[A-Z]{2}(?=[\s,.\n]|$)"
r"|(?<=[\s,])[A-Z]\d[A-Z]\s*\d[A-Z]\d(?=[\s,.\n]|$)"
r"|(?<=[\s,])\d{3}-\d{4}(?=[\s,.\n]|$)",
re.IGNORECASE,
),
validator=is_valid_zipcode,
),
"AMOUNT": RegexRule(
entity_type="AMOUNT",
pattern=re.compile(
r"(?<!\w)(?:[$€£¥₫₩₹])\s*\d[\d,. ]*\d(?:\s*(?:USD|EUR|GBP|VND|JPY|KRW|INR|AUD|CAD|CHF|CNY|SGD|THB|đồng|đ|dong))?"
r"|(?<!\w)\d[\d,. ]*\d\s*(?:USD|EUR|GBP|VND|JPY|KRW|INR|AUD|CAD|CHF|CNY|SGD|THB|đồng|đ|dong)\b",
re.IGNORECASE,
),
validator=is_valid_amount,
),
"CRYPTO_ADDRESS": RegexRule(
entity_type="CRYPTO_ADDRESS",
pattern=re.compile(
r"\b(?:0x[0-9a-fA-F]{40}|[13][A-HJ-NP-Za-km-z1-9]{25,34}|bc1[a-z0-9]{25,89}"
r"|[LM3][A-HJ-NP-Za-km-z1-9]{26,33}|r[0-9A-HJ-NP-Za-km-z]{24,34})\b"
),
validator=is_valid_crypto_address,
),
"AGE": RegexRule(
entity_type="AGE",
pattern=re.compile(
r"(?i)(?:"
r"(?:age(?:d)?|tuổi|năm tuổi|years?\s*old|ans|Jahre?\s*alt)\s*[:#]?\s*(\d{1,3})"
r"|(\d{1,3})\s*(?:tuổi|năm tuổi|years?\s*old|ans|Jahre?\s*alt)"
r")"
),
validator=is_valid_age,
),
}
_RULE_PRIORITY: List[str] = [
"CRYPTO_ADDRESS", "IBAN", "CREDIT_CARD_NUMBER",
"SSN_CCCD_US", "SSN_CCCD_VN",
"IP_ADDRESS_V6", "IP_ADDRESS_V4",
"URL", "EMAIL", "DATE", "AMOUNT", "AGE",
]
RULE_BASED_LABELS: set = {rule.entity_type for rule in RULE_PATTERNS.values()}
_TRAILING_PUNCT = re.compile(r"[.,;:)\]}]+$")
def _trim_trailing_punct(value: str, end: int) -> tuple:
trimmed = _TRAILING_PUNCT.sub("", value)
return trimmed, end - (len(value) - len(trimmed))
def detect_by_rules(text: str, allowed_labels) -> List[Dict]:
if isinstance(allowed_labels, dict):
allowed_entity_types = set(allowed_labels.keys())
else:
allowed_entity_types = set(allowed_labels)
results: List[Dict] = []
for key in _RULE_PRIORITY:
rule = RULE_PATTERNS.get(key)
if rule is None or rule.entity_type not in allowed_entity_types:
continue
for m in rule.pattern.finditer(text):
raw = m.group(0).strip()
cleaned, end = _trim_trailing_punct(raw, m.end())
if not cleaned:
continue
if not rule.validator(cleaned):
continue
results.append({
"text": cleaned, "label": rule.entity_type,
"start": m.start(), "end": end,
})
return results
# ==========================================
# 2d. Postprocessing — normalize & resolve conflicts
# ==========================================
def _find_all_occurrences(text: str, value: str) -> List[Tuple[int, int]]:
spans, start = [], 0
while True:
idx = text.find(value, start)
if idx == -1:
break
spans.append((idx, idx + len(value)))
start = idx + max(1, len(value))
return spans
def normalize_predictions(raw_preds, text: str) -> List[Dict]:
if raw_preds is None:
return []
results: List[Dict] = []
if isinstance(raw_preds, list):
for item in raw_preds:
if not isinstance(item, dict):
continue
label = item.get("label") or item.get("entity_group") or item.get("entity") or ""
start = item.get("start", -1)
end = item.get("end", -1)
value = item.get("text") or item.get("value") or item.get("word") or ""
if not label:
continue
score = float(item.get("score") or item.get("probability") or item.get("confidence") or 0.0)
if start != -1 and end != -1 and 0 <= int(start) < int(end) <= len(text):
results.append({"text": text[int(start):int(end)], "label": label,
"start": int(start), "end": int(end), "score": score})
elif value:
for s, e in _find_all_occurrences(text, str(value)):
results.append({"text": text[s:e], "label": label, "start": s, "end": e, "score": score})
elif isinstance(raw_preds, dict):
entities = raw_preds.get("entities", raw_preds)
for label, values in entities.items():
if not isinstance(values, list):
continue
for v in values:
if isinstance(v, dict):
s, e = v.get("start"), v.get("end")
surface = v.get("text") or v.get("value") or ""
sc = float(v.get("score") or v.get("probability") or v.get("confidence") or 0.0)
if s is not None and e is not None and 0 <= int(s) < int(e) <= len(text):
results.append({"text": text[int(s):int(e)], "label": label,
"start": int(s), "end": int(e), "score": sc})
elif surface:
for cs, ce in _find_all_occurrences(text, str(surface)):
results.append({"text": text[cs:ce], "label": label,
"start": cs, "end": ce, "score": sc})
else:
for cs, ce in _find_all_occurrences(text, str(v)):
results.append({"text": text[cs:ce], "label": label,
"start": cs, "end": ce, "score": 0.0})
best: Dict[tuple, Dict] = {}
for r in results:
key = (r["label"], r["start"], r["end"])
if key not in best or r.get("score", 0.0) > best[key].get("score", 0.0):
best[key] = r
return sorted(best.values(), key=lambda x: x["start"])
_LABEL_PRIORITY: Dict[str, int] = {
"CRYPTO_ADDRESS": 10, "IBAN": 10, "CREDIT_CARD_NUMBER": 10,
"SSN_CCCD_US": 10, "SSN_CCCD_VN": 10,
"IP_ADDRESS_V6": 10, "IP_ADDRESS_V4": 10,
"URL": 10, "EMAIL": 10, "DATE": 10, "AMOUNT": 10, "AGE": 10,
}
def resolve_span_conflicts(
entities: List[Dict],
strategy: str = "score_first",
rule_labels: Optional[set] = None,
) -> List[Dict]:
if not entities:
return []
_rule_labels = rule_labels if rule_labels is not None else RULE_BASED_LABELS
def _sort_key(e: Dict):
is_rule = 1 if e["label"] in _rule_labels else 0
score = 1.0 if is_rule else float(e.get("score") or 0.0)
length = e["end"] - e["start"]
prio = _LABEL_PRIORITY.get(e["label"], 0)
if strategy == "score_first":
return (is_rule, score, prio, length)
elif strategy == "longest_first":
return (is_rule, length, score, prio)
else:
return (is_rule, prio, score, length)
ranked = sorted(entities, key=_sort_key, reverse=True)
kept: List[Dict] = []
occupied: List[Tuple[int, int]] = []
def _overlaps(s: int, e: int) -> bool:
for (ks, ke) in occupied:
if s < ke and e > ks:
return True
return False
for ent in ranked:
s, e = ent["start"], ent["end"]
if not _overlaps(s, e):
kept.append(ent)
occupied.append((s, e))
return sorted(kept, key=lambda x: x["start"])
def apply_advanced_postprocessing(entities: List[Dict], original_text: str) -> List[Dict]:
if not entities:
return []
sorted_ents = sorted(entities, key=lambda x: x['start'])
final_entities = []
for i, ent in enumerate(sorted_ents):
label = ent["label"]
start = ent["start"]
end = ent["end"]
if label in ["TIME", "DATE"]:
prev_char = original_text[start-1] if start > 0 else ""
next_char = original_text[end] if end < len(original_text) else ""
if prev_char == "[" and next_char == "]":
continue
if label == "PREFIX":
has_person_after = False
if i + 1 < len(sorted_ents):
next_ent = sorted_ents[i+1]
if next_ent["label"] == "PERSONAL_NAME":
gap = next_ent["start"] - end
if 0 <= gap <= 3:
has_person_after = True
if not has_person_after:
continue
final_entities.append(ent)
return final_entities
# ==========================================
# 3. Model Wrappers
# ==========================================
class _BaseWrapper:
def cleanup(self):
if hasattr(self, "model"):
del self.model
if torch.cuda.is_available():
torch.cuda.empty_cache()
gc.collect()
def _run_model(self, text: str, label_input):
raw = None
for method in ("predict_entities", "extract_entities"):
fn = getattr(self.model, method, None)
if fn is None:
continue
try:
raw = fn(text, label_input, threshold=self.threshold)
except TypeError:
raw = fn(text, label_input)
break
return raw
class GLiNER2FinetunedWrapper(_BaseWrapper):
def __init__(
self,
finetuned_repo: str = "quynong/gliner2-rosenxt",
checkpoint: str = "best",
use_label_description: bool = False,
use_rules: bool = True,
threshold: float = 0.3,
MAX_SPAN: int = 100,
):
self.use_label_description = use_label_description
self.use_rules = use_rules
self.threshold = threshold
self.device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Downloading {checkpoint!r} from {finetuned_repo}")
local_path = snapshot_download(
repo_id=finetuned_repo,
allow_patterns=f"{checkpoint}/*"
)
checkpoint_path = os.path.join(local_path, checkpoint)
print(f"Loading model from: {checkpoint_path}")
self.model = GLiNER2.from_pretrained(checkpoint_path)
self.model.config.max_width = MAX_SPAN
self.model.max_width = MAX_SPAN
self.model.span_rep.span_rep_layer.max_width = MAX_SPAN
self.model.to(self.device).eval()
def predict(self, text: str, labels, label_descriptions: Optional[Dict[str, str]] = None) -> List[Dict]:
cleaned_text, char_map = clean_text_with_mapping(text)
raw = self._run_model(
cleaned_text,
labels if not (self.use_label_description and label_descriptions)
else {k: v for k, v in label_descriptions.items()
if k in (labels if isinstance(labels, set)
else set(labels.keys()) if isinstance(labels, dict)
else set(labels))}
)
model_preds = normalize_predictions(raw, cleaned_text)
rule_preds: List[Dict] = []
if self.use_rules:
for e in detect_by_rules(cleaned_text, labels):
rule_preds.append({**e, "score": 1.0})
combined = rule_preds + model_preds
resolved = resolve_span_conflicts(combined, strategy="rule_first")
for pred in resolved:
s, e = pred["start"], pred["end"]
pred["start"] = char_map[s] if s < len(char_map) else len(text)
pred["end"] = char_map[e] if e < len(char_map) else len(text)
pred["text"] = text[pred["start"]:pred["end"]]
return apply_advanced_postprocessing(resolved, text)
class mmBERTWrapper(_BaseWrapper):
def __init__(
self,
hf_repo: str,
subfolder: str = "",
threshold: float = 0.5,
use_rules: bool = True,
):
self.threshold = threshold
self.use_rules = use_rules
if subfolder:
print(f"Downloading subfolder '{subfolder}' from {hf_repo}")
local_repo_path = snapshot_download(
repo_id=hf_repo,
allow_patterns=f"{subfolder}/*"
)
model_path = os.path.join(local_repo_path, subfolder)
else:
model_path = hf_repo
print(f"Loading mmBERT pipeline from: {model_path}")
self.pipe = pipeline(
"ner",
model=model_path,
tokenizer=model_path,
aggregation_strategy="simple",
)
def predict(self, text: str, labels) -> List[Dict]:
cleaned_text, char_map = clean_text_with_mapping(text)
raw = self.pipe(cleaned_text)
model_preds = normalize_predictions(raw, cleaned_text)
# Filter by threshold
model_preds = [p for p in model_preds if p.get("score", 0.0) >= self.threshold]
rule_preds: List[Dict] = []
if self.use_rules:
for e in detect_by_rules(cleaned_text, labels):
rule_preds.append({**e, "score": 1.0})
combined = rule_preds + model_preds
resolved = resolve_span_conflicts(combined, strategy="rule_first")
for pred in resolved:
s, e = pred["start"], pred["end"]
pred["start"] = char_map[s] if s < len(char_map) else len(text)
pred["end"] = char_map[e] if e < len(char_map) else len(text)
pred["text"] = text[pred["start"]:pred["end"]]
return apply_advanced_postprocessing(resolved, text)
def cleanup(self):
if hasattr(self, "pipe"):
del self.pipe
if torch.cuda.is_available():
torch.cuda.empty_cache()
gc.collect()
# ==========================================
# 3b. Lazy model loaders (singleton)
# ==========================================
_MODEL_INSTANCES: Dict[str, object] = {}
_MODEL_FACTORIES = {
"multi_1": lambda: GLiNER2FinetunedWrapper(
finetuned_repo="quynong/gliner2-rosenxt-multi-v7",
use_label_description=False,
use_rules=False,
checkpoint="checkpoint-epoch-3",
threshold=0.7,
MAX_SPAN=90,
),
"XLMBERT": lambda: mmBERTWrapper(
hf_repo="quynong/multilingualv1",
subfolder="FacebookAI-xlm-roberta-base",
threshold=0.5,
use_rules=False,
),
"mmBERT-multilingual": lambda: mmBERTWrapper(
hf_repo="quynong/multilingual",
subfolder="jhu-clsp-mmBERT-base",
threshold=0.5,
use_rules=False,
),
}
MODEL_DISPLAY_NAMES = {
"multi_1": "GLiNER2 — multi-v7 (ep3)",
"XLMBERT": "XLM-RoBERTa (BERT)",
"mmBERT-multilingual": "mmBERT-multilingual",
}
def get_model(key: str):
if key not in _MODEL_INSTANCES:
print(f"🚀 Khởi tạo model: {key}")
_MODEL_INSTANCES[key] = _MODEL_FACTORIES[key]()
print(f"✅ Model '{key}' sẵn sàng.")
return _MODEL_INSTANCES[key]
# ==========================================
# 4. Hàm tạo HTML
# ==========================================
def build_html_with_tooltips(text: str, spans: List[Dict]) -> str:
spans = sorted(spans, key=lambda x: x["start"])
html_content = ""
last_idx = 0
for span in spans:
start, end, label = span["start"], span["end"], span["entity"]
if start >= end or start < last_idx:
continue
html_content += text[last_idx:start].replace('\n', '<br>')
pii_text = text[start:end].replace('\n', '<br>')
html_content += (
f'<span class="pii-entity">{pii_text}'
f'<span class="tooltip-label">{label}</span></span>'
)
last_idx = end
html_content += text[last_idx:].replace('\n', '<br>')
return f'<div class="text-container">{html_content}</div>'
# ==========================================
# 5. Core: chạy 1 model, trả về (html, json)
# ==========================================
def run_single_model(text: str, model_key: str):
try:
wrapper = get_model(model_key)
preds = wrapper.predict(text=text, labels=LABELS)
spans_for_html = [{"entity": p["label"], "start": p["start"], "end": p["end"]} for p in preds]
html = build_html_with_tooltips(text, spans_for_html)
return html, preds
except Exception as exc:
import traceback
tb = traceback.format_exc()
return f"<div style='color:red;'>Lỗi model {model_key}: {exc}</div>", {"error": str(exc), "traceback": tb}
# ==========================================
# 6. Wrapper cho Gradio (3 model song song)
# ==========================================
def process_both(text: str):
html1, json1 = run_single_model(text, "multi_1")
html2, json2 = run_single_model(text, "XLMBERT")
html3, json3 = run_single_model(text, "mmBERT-multilingual")
return html1, json1, html2, json2, html3, json3
# ==========================================
# 6b. Batch processing — test file upload
# ==========================================
def _build_entity_tags_html(preds: list) -> str:
if not preds or (isinstance(preds, dict) and "error" in preds):
return '<span style="color:#ef4444;">Lỗi</span>'
if not preds:
return '<span style="color:#9ca3af;">Không phát hiện PII</span>'
tags = []
for p in preds:
label = p.get("label", "?")
text_val = p.get("text", "")
display = text_val if len(text_val) <= 30 else text_val[:27] + "..."
display = display.replace("<", "&lt;").replace(">", "&gt;").replace('"', "&quot;")
tags.append(f'<span class="batch-entity-tag" title="{display}">{label}: {display}</span>')
return "".join(tags)
def process_test_file(file, model_choice: str):
if file is None:
return "<div class='batch-progress'>⬆️ Vui lòng upload file JSON.</div>", []
try:
file_path = file if isinstance(file, str) else file.name
with open(file_path, "r", encoding="utf-8") as f:
data = json.load(f)
except Exception as exc:
return f"<div style='color:red;padding:20px;'>Lỗi đọc file: {exc}</div>", []
if not isinstance(data, list):
return "<div style='color:red;padding:20px;'>File JSON phải là một mảng (array) các object có trường \"text\".</div>", []
model_keys = []
if model_choice in ("GLiNER2", "Tất cả"):
model_keys.append("multi_1")
if model_choice in ("XLM-RoBERTa", "Tất cả"):
model_keys.append("XLMBERT")
if model_choice in ("mmBERT-multilingual", "Tất cả"):
model_keys.append("mmBERT-multilingual")
total = len(data)
all_results = []
html_parts = []
# Summary header placeholder
total_pii_gliner = 0
total_pii_bert = 0
total_pii_mmbert = 0
for idx, item in enumerate(data):
text = item.get("text", "")
if not text:
continue
sample_num = idx + 1
short_text = text[:150].replace("<", "&lt;").replace(">", "&gt;").replace("\n", " ")
if len(text) > 150:
short_text += "..."
result_entry = {"sample": sample_num, "text_preview": text[:100]}
# Run models
html_gliner = preds_gliner = html_bert = preds_bert = html_mmbert = preds_mmbert = None
if "multi_1" in model_keys:
html_gliner, preds_gliner = run_single_model(text, "multi_1")
result_entry["gliner"] = preds_gliner
if isinstance(preds_gliner, list):
total_pii_gliner += len(preds_gliner)
if "XLMBERT" in model_keys:
html_bert, preds_bert = run_single_model(text, "XLMBERT")
result_entry["xlmbert"] = preds_bert
if isinstance(preds_bert, list):
total_pii_bert += len(preds_bert)
if "mmBERT-multilingual" in model_keys:
html_mmbert, preds_mmbert = run_single_model(text, "mmBERT-multilingual")
result_entry["mmbert_multilingual"] = preds_mmbert
if isinstance(preds_mmbert, list):
total_pii_mmbert += len(preds_mmbert)
all_results.append(result_entry)
# Build sample HTML
html_parts.append(f'<div class="batch-sample">')
html_parts.append(
f'<div class="batch-sample-header">'
f'<span>📄 Mẫu #{sample_num} / {total}</span>'
f'<span style="font-weight:normal;font-size:0.8rem;">{len(text)} ký tự</span>'
f'</div>'
)
html_parts.append(f'<div class="batch-sample-input">{short_text}</div>')
if len(model_keys) >= 2:
cols = f"repeat({len(model_keys)}, 1fr)"
html_parts.append(f'<div class="batch-models-row" style="grid-template-columns:{cols};">')
_batch_cols = [
("multi_1", "gliner", "GLiNER2", html_gliner, preds_gliner),
("XLMBERT", "bert", "XLM-RoBERTa", html_bert, preds_bert),
("mmBERT-multilingual", "mmbert", "mmBERT-multi", html_mmbert, preds_mmbert),
]
for mk, cls, name, h, p in _batch_cols:
if mk not in model_keys:
continue
html_parts.append('<div class="batch-model-col">')
html_parts.append(f'<div class="batch-model-label {cls}">{name}</div>')
if h:
html_parts.append(h)
n = len(p) if isinstance(p, list) else 0
html_parts.append(f'<div class="batch-entity-tags"><strong style="font-size:0.75rem;margin-right:6px;">{n} entities:</strong>')
html_parts.append(_build_entity_tags_html(p))
html_parts.append('</div></div>')
html_parts.append('</div>') # end batch-models-row
else:
# Single model
mk = model_keys[0]
_single_map = {
"multi_1": ("gliner", "GLiNER2", html_gliner, preds_gliner),
"XLMBERT": ("bert", "XLM-RoBERTa", html_bert, preds_bert),
"mmBERT-multilingual": ("mmbert", "mmBERT-multi", html_mmbert, preds_mmbert),
}
cls, name, h, p = _single_map[mk]
html_parts.append(f'<div style="padding:12px 16px;">')
html_parts.append(f'<div class="batch-model-label {cls}">{name}</div>')
if h:
html_parts.append(h)
n = len(p) if isinstance(p, list) else 0
html_parts.append(f'<div class="batch-entity-tags"><strong style="font-size:0.75rem;margin-right:6px;">{n} entities:</strong>')
html_parts.append(_build_entity_tags_html(p))
html_parts.append('</div></div>')
html_parts.append('</div>') # end batch-sample
# Build summary
summary_lines = [
f"<div class='batch-summary'>",
f"<strong>📊 Tổng kết:</strong> {total} mẫu đã xử lý<br>",
]
if "multi_1" in model_keys:
summary_lines.append(f"🤖 <strong>GLiNER2:</strong> {total_pii_gliner} PII entities phát hiện<br>")
if "XLMBERT" in model_keys:
summary_lines.append(f"🧠 <strong>XLM-RoBERTa:</strong> {total_pii_bert} PII entities phát hiện<br>")
if "mmBERT-multilingual" in model_keys:
summary_lines.append(f"🔶 <strong>mmBERT-multilingual:</strong> {total_pii_mmbert} PII entities phát hiện")
summary_lines.append("</div>")
final_html = "".join(summary_lines) + "".join(html_parts)
return final_html, all_results
# ==========================================
# 7. Giao diện Gradio UI — song song 2 model
# ==========================================
DEFAULT_TEXT = (
"Dear Mr. PATEL,\n\n"
"We are pleased to inform you that your application for a personal loan has been approved "
"by RiverbankFinancial. The approval was finalized on 14-05-2024 09:45 and your documents "
"will be processed within the next two business days. As a resident of Illinois, your "
"application was reviewed in accordance with all regional regulations. Please note that your "
"National ID, AID-6543217890, has been securely verified as part of our compliance process.\n\n"
"For your records, your online application was submitted from the IP address "
"2a02:4d60:1f31:4c3f:85e1:1122:abfc:0345. Your loan agreement and repayment schedule will "
"be sent to you via email by john.patel@email.com. If you have any questions, please contact "
"our support team at 555-0198.\n\n"
)
with gr.Blocks(title="PII Extractor — So sánh 3 Model", theme=gr.themes.Soft(), css=custom_css) as demo:
gr.Markdown("# 🔍 Hệ thống Trích xuất PII — So sánh 3 Model Song Song")
with gr.Tabs():
# ── Tab 1: Nhập văn bản ──────────────────────────────────────
with gr.TabItem("📝 Nhập văn bản"):
gr.Markdown(
"Nhập văn bản bên dưới và nhấn **Trích xuất PII**. "
"Kết quả của **GLiNER2 multi-v7**, **XLM-RoBERTa** và **mmBERT-multilingual** sẽ hiển thị đồng thời."
)
with gr.Row():
input_text = gr.Textbox(
lines=12,
label="📝 Văn bản đầu vào",
value=DEFAULT_TEXT,
scale=1,
)
submit_btn = gr.Button("⚡ Trích xuất PII (Cả 3 Model)", variant="primary", size="lg")
gr.Markdown("## 📊 Kết quả nhận diện")
with gr.Row(equal_height=True):
with gr.Column():
gr.HTML(
'<div class="model-header-gliner">🤖 Model 1 — GLiNER2 multi-v7 (checkpoint-epoch-3, threshold=0.7)</div>'
)
out_html_gliner = gr.HTML(label="GLiNER2 Output")
with gr.Column():
gr.HTML(
'<div class="model-header-xlmbert">🧠 Model 2 — XLM-RoBERTa (FacebookAI, threshold=0.5)</div>'
)
out_html_bert = gr.HTML(label="XLM-RoBERTa Output")
with gr.Column():
gr.HTML(
'<div class="model-header-mmbert">🔶 Model 3 — mmBERT-multilingual (jhu-clsp, threshold=0.5)</div>'
)
out_html_mmbert = gr.HTML(label="mmBERT-multilingual Output")
gr.Markdown("## 🗂️ Dữ liệu JSON thô")
with gr.Row(equal_height=True):
with gr.Column():
gr.Markdown("**GLiNER2 JSON**")
out_json_gliner = gr.JSON(label="GLiNER2 Raw JSON")
with gr.Column():
gr.Markdown("**XLM-RoBERTa JSON**")
out_json_bert = gr.JSON(label="XLM-RoBERTa Raw JSON")
with gr.Column():
gr.Markdown("**mmBERT-multilingual JSON**")
out_json_mmbert = gr.JSON(label="mmBERT-multilingual Raw JSON")
submit_btn.click(
fn=process_both,
inputs=[input_text],
outputs=[out_html_gliner, out_json_gliner, out_html_bert, out_json_bert, out_html_mmbert, out_json_mmbert],
)
# ── Tab 2: Test File ─────────────────────────────────────────
with gr.TabItem("📂 Test File"):
gr.Markdown(
"Upload file JSON chứa các mẫu test (format: `[{\"text\": \"...\"}, ...]`). "
"Hệ thống sẽ chạy model trên từng mẫu và hiển thị kết quả trực quan."
)
with gr.Row():
with gr.Column(scale=2):
test_file_input = gr.File(
label="📁 Upload file JSON test",
file_types=[".json"],
type="filepath",
)
with gr.Column(scale=1):
model_choice = gr.Radio(
choices=["Tất cả", "GLiNER2", "XLM-RoBERTa", "mmBERT-multilingual"],
value="Tất cả",
label="🔧 Chọn model",
)
test_btn = gr.Button("🚀 Chạy Test Batch", variant="primary", size="lg")
gr.Markdown("## 📊 Kết quả Test")
test_output_html = gr.HTML(label="Kết quả trực quan")
with gr.Accordion("🗂️ JSON chi tiết", open=False):
test_output_json = gr.JSON(label="Kết quả JSON")
test_btn.click(
fn=process_test_file,
inputs=[test_file_input, model_choice],
outputs=[test_output_html, test_output_json],
)
if __name__ == "__main__":
demo.launch(server_name="0.0.0.0", server_port=7860)