Spaces:
Sleeping
Sleeping
| 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 | |
| # ========================================== | |
| 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("<", "<").replace(">", ">").replace('"', """) | |
| 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("<", "<").replace(">", ">").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) | |