| import asyncio
|
| import atexit
|
| import json
|
| import os
|
| import re
|
| import tempfile
|
| import unicodedata
|
| from functools import lru_cache
|
| from pathlib import Path
|
| from typing import Dict, List, Optional, Tuple
|
|
|
|
|
| def _close_event_loop() -> None:
|
| try:
|
| loop = asyncio.get_event_loop()
|
| if not loop.is_closed():
|
| loop.close()
|
| except Exception:
|
| pass
|
|
|
|
|
| atexit.register(_close_event_loop)
|
|
|
| import gradio as gr
|
|
|
| try:
|
| from huggingface_hub import login as hf_login
|
| except Exception:
|
| hf_login = None
|
|
|
| try:
|
| import torch
|
| except Exception:
|
| torch = None
|
|
|
| from gliner2 import GLiNER2
|
| from rule_detector import DEFAULT_RULEBASE_ALLOWED_ENTITY_TYPES, detect_by_rules
|
|
|
|
|
| DEFAULT_MODEL_ID = os.getenv("MODEL_ID", "AITeamUIT/gliner2-multi-v1-e3-25-6")
|
| DEFAULT_THRESHOLD = float(os.getenv("THRESHOLD", "0.5"))
|
| DEFAULT_CHUNK_CHARS = int(os.getenv("CHUNK_CHARS", "1000"))
|
| DEFAULT_CHUNK_TOKENS = int(os.getenv("CHUNK_TOKENS", "768"))
|
| CHUNK_SAFETY_MARGIN = int(os.getenv("CHUNK_SAFETY_MARGIN", "20"))
|
| MAX_CHARS_PER_CHUNK = int(os.getenv("MAX_CHARS_PER_CHUNK", "8000"))
|
| MAX_WIDTH = int(os.getenv("MAX_WIDTH", "30"))
|
|
|
| PII_LABEL_DESCRIPTIONS: Dict[str, str] = {
|
| "PERSON": "Person names",
|
| "DATE": "Calendar dates",
|
| "JOB_TITLE": "Job titles",
|
| "TIME": "Time values",
|
| "LOCATION": "General locations",
|
| "PREFIX": "Titles before names",
|
| "ORGANIZATION": "Organization names",
|
| "EMAIL": "Email addresses",
|
| "PHONE": "Phone numbers",
|
| "ADDRESS": "Street addresses",
|
| "GENDER": "Sex or gender",
|
| "COORDINATE": "Geographic coordinates",
|
| "MONEY": "Monetary amounts",
|
| "IP": "IP addresses",
|
| "ZIP_CODE": "Postal or ZIP codes",
|
| "NATIONALITY": "Legal nationality",
|
| "AGE": "Age values",
|
| "USERNAME": "Usernames or handles",
|
| "URL": "Web links or URLs",
|
| "MARITAL": "Marital status",
|
| "TRADE_UNION": "Trade union membership",
|
| "BIRTHDATE": "Date of birth",
|
| "CARD_ISSUER": "Card brands or issuers",
|
| "PIN": "Authentication PINs",
|
| "PASSPORT": "Passport numbers",
|
| "MEDICAL_INFO": "Health or medical information",
|
| "IBAN": "International bank account numbers",
|
| "PASSWORD": "Account passwords",
|
| "ETHNICITY": "Ethnic or racial group",
|
| "RELIGION": "Religious affiliation",
|
| "EMPLOYEE_ID": "Employee identifiers",
|
| "INSURANCE_ID": "Insurance identifiers",
|
| "TIN": "Tax identification numbers",
|
| "SWIFT": "Bank SWIFT or BIC codes",
|
| "CARD_NUMBER": "Payment card numbers",
|
| "ACCOUNT_ID": "Account identifiers",
|
| "WALLET": "Crypto wallet addresses",
|
| "NATIONAL_ID": "National ID numbers",
|
| "CVV": "Card security codes",
|
| "BANK_ACCOUNT": "Bank account numbers",
|
| "LICENSE": "Driver license numbers",
|
| "TICKET_ID": "Support or case IDs",
|
| "PLATE": "Vehicle plate numbers",
|
| "API_KEY": "API keys or tokens",
|
| }
|
|
|
| LABELS = list(PII_LABEL_DESCRIPTIONS.keys())
|
| RULE_LABELS = set(DEFAULT_RULEBASE_ALLOWED_ENTITY_TYPES)
|
|
|
| LABEL_PRIORITY: Dict[str, int] = {
|
| "WALLET": 10,
|
| "IBAN": 10,
|
| "CARD_NUMBER": 10,
|
| "NATIONAL_ID": 10,
|
| "IP": 10,
|
| "URL": 10,
|
| "EMAIL": 10,
|
| "MONEY": 10,
|
| "AGE": 10,
|
| }
|
|
|
| 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]+")
|
| NUMERIC_MASK_LABELS = {
|
| "CARD_NUMBER",
|
| "PHONE",
|
| "PIN",
|
| "CVV",
|
| "TIN",
|
| "BANK_ACCOUNT",
|
| "IBAN",
|
| "SWIFT",
|
| "WALLET",
|
| "IP",
|
| "ZIP_CODE",
|
| }
|
|
|
|
|
| try:
|
| from gliner2.processor import WhitespaceTokenSplitter
|
|
|
| IMPROVED_PATTERN = re.compile(
|
| r"""
|
| # Scheme / www URL
|
| (?:(?:https?|ftp)://|www\.)[^\s<>"'`{}|\\^]*[^\s<>"'`{}|\\^.,;:!?()\[\]]
|
|
|
| # Email address
|
| |[a-z0-9.!#$%&'*+/=?^_`{|}~-]+@(?:[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?\.)+[a-z]{2,63}
|
|
|
| # Social/user handle
|
| |@[a-z0-9_](?:[a-z0-9_.-]{0,28}[a-z0-9_])?
|
|
|
| # IPv4 address. Allows sentence punctuation after the IP, but avoids
|
| # matching inside a longer dotted numeric sequence.
|
| |(?:25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(?:\.(?:25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3}(?=$|[^\w])(?!\.\d)
|
|
|
| # IPv6 address, including compressed forms and optional zone IDs.
|
| |(?:
|
| (?:[0-9a-f]{1,4}:){7}[0-9a-f]{1,4}
|
| |[0-9a-f]{1,4}:(?:(?::[0-9a-f]{1,4}){1,6})
|
| |(?:[0-9a-f]{1,4}:){1,2}(?::[0-9a-f]{1,4}){1,5}
|
| |(?:[0-9a-f]{1,4}:){1,3}(?::[0-9a-f]{1,4}){1,4}
|
| |(?:[0-9a-f]{1,4}:){1,4}(?::[0-9a-f]{1,4}){1,3}
|
| |(?:[0-9a-f]{1,4}:){1,5}(?::[0-9a-f]{1,4}){1,2}
|
| |(?:[0-9a-f]{1,4}:){1,6}:[0-9a-f]{1,4}
|
| |:(?:(?::[0-9a-f]{1,4}){1,7}|:)
|
| |(?:[0-9a-f]{1,4}:){1,7}:
|
| )(?:%[0-9a-z_.-]+)?(?![0-9a-f:])
|
|
|
| # Strict JWT-like token
|
| |[a-z0-9_-]{10,}\.[a-z0-9_-]{10,}\.[a-z0-9_-]{10,}
|
|
|
| # Shorter dot-separated token-like string, e.g. sk-abc.def.ghi
|
| |(?=[a-z0-9_.-]*[-_\d])[a-z0-9_-]{2,}(?:\.[a-z0-9_-]{2,}){2,}
|
|
|
| # Bare domain / URL-like candidate. Validation belongs downstream; this
|
| # tokenizer only keeps plausible URL spans from being split into pieces.
|
| |(?<![\w@])
|
| (?:[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?\.)+
|
| (?:[a-z]{2,63}|xn--[a-z0-9-]{2,59})
|
| (?=$|/|:|[^\w-])
|
| (?::\d{2,5})?
|
| (?:/(?:[^\s<>"'`{}|\\^]*[^\s<>"'`{}|\\^.,;:!?()\[\]])?)?
|
|
|
| # Letter-containing word token: must have at least one letter or
|
| # underscore (the [^\W\d] anchor). Only underscores join tokens;
|
| # hyphens always split (except digit-digit handled below).
|
| |(?:\d*[^\W\d]\w*)(?:_\w+)*
|
|
|
| # Pure digit sequence: hyphens/underscores only join other digit groups.
|
| # "555-1234" stays atomic; "13-tuổi" → "13" / "-" / "tuổi" so the
|
| # numeric entity (AGE, DATE, …) is not fused with the surrounding word.
|
| |\d+(?:[-_]\d+)*
|
|
|
| # Fallback: any single non-space character.
|
| |\S
|
| """,
|
| re.VERBOSE | re.IGNORECASE,
|
| )
|
| WhitespaceTokenSplitter._PATTERN = IMPROVED_PATTERN
|
| except Exception:
|
| pass
|
|
|
|
|
| @lru_cache(maxsize=2)
|
| def load_model(model_id: str):
|
| hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
|
| if hf_token and hf_login is not None:
|
| hf_login(token=hf_token)
|
| model = GLiNER2.from_pretrained(model_id)
|
| model.config.max_width = MAX_WIDTH
|
| model.max_width = MAX_WIDTH
|
| if hasattr(model, "span_rep") and hasattr(model.span_rep, "span_rep_layer"):
|
| model.span_rep.span_rep_layer.max_width = MAX_WIDTH
|
| if torch is not None:
|
| device = "cuda" if torch.cuda.is_available() else "cpu"
|
| model = model.to(device)
|
| model.eval()
|
| return model
|
|
|
|
|
| def build_schema(model, use_descriptions: bool):
|
| schema = model.create_schema()
|
| if use_descriptions:
|
| schema.entities(PII_LABEL_DESCRIPTIONS)
|
| else:
|
| schema.entities(LABELS)
|
| return schema
|
|
|
|
|
| def manual_chunks(text: str, max_chars: int) -> List[Tuple[str, int]]:
|
| if len(text) <= max_chars:
|
| return [(text, 0)]
|
| tokens = list(re.finditer(r"\S+", text))
|
| if not tokens:
|
| return [(text, 0)]
|
| chunks: List[Tuple[str, int]] = []
|
| chunk_start = 0
|
| for index in range(1, len(tokens)):
|
| seg_start = tokens[chunk_start].start()
|
| seg_end = tokens[index].end()
|
| if seg_end - seg_start > max_chars:
|
| prev_end = tokens[index - 1].end()
|
| chunks.append((text[seg_start:prev_end], seg_start))
|
| chunk_start = index
|
| seg_start = tokens[chunk_start].start()
|
| chunks.append((text[seg_start:tokens[-1].end()], seg_start))
|
| return chunks
|
|
|
|
|
| def word_token_count(text: str) -> int:
|
| """Count word tokens the way the (patched) model tokenizer splits them."""
|
| try:
|
| from gliner2.processor import WhitespaceTokenSplitter
|
|
|
| splitter = WhitespaceTokenSplitter()
|
| return sum(1 for _ in splitter(text, lower=False))
|
| except Exception:
|
| return len(text.split())
|
|
|
|
|
| class RecursiveWordChunker:
|
| """Label-aware recursive chunker built on chonkie.RecursiveChunker.
|
|
|
| The chunk-size budget is measured in the model's own word-token space and
|
| already subtracts the entity-schema label tokens, so each emitted chunk plus
|
| the labels fits the encoder. Recursion falls back paragraph -> sentence ->
|
| whitespace. ``chunk()`` returns ``(chunk_text, start_index)`` tuples.
|
| """
|
|
|
| def __init__(self, budget_words: int, max_chars: int = MAX_CHARS_PER_CHUNK, min_chars_per_chunk: int = 24):
|
|
|
| from chonkie.chunker.recursive import RecursiveChunker
|
| from chonkie.types import RecursiveLevel, RecursiveRules
|
|
|
| self.budget = budget_words
|
| self.max_chars = max_chars
|
| self._chunker = RecursiveChunker(
|
| tokenizer=word_token_count,
|
| chunk_size=budget_words,
|
| rules=RecursiveRules(
|
| levels=[
|
| RecursiveLevel(delimiters=["\n\n", "\r\n", "\n", "\r"]),
|
| RecursiveLevel(delimiters=[". ", "! ", "? ", ".\n", "!\n", "?\n"]),
|
| RecursiveLevel(whitespace=True),
|
| ]
|
| ),
|
| min_characters_per_chunk=min_chars_per_chunk,
|
| )
|
|
|
| def chunk(self, text: str) -> List[Tuple[str, int]]:
|
| return [(c.text, int(c.start_index)) for c in self._chunker.chunk(text)]
|
|
|
|
|
| @lru_cache(maxsize=4)
|
| def build_word_chunker(
|
| chunk_tokens: int = DEFAULT_CHUNK_TOKENS,
|
| safety_margin: int = CHUNK_SAFETY_MARGIN,
|
| max_chars: int = MAX_CHARS_PER_CHUNK,
|
| ) -> RecursiveWordChunker:
|
| """Build a :class:`RecursiveWordChunker` with a label-aware token budget.
|
|
|
| The label tokens (``"<LABEL> <description>"`` plus a +2 separator budget each)
|
| are subtracted from ``chunk_tokens`` so the encoder never overflows.
|
| """
|
| labels_words = [f"{k} {v}" for k, v in PII_LABEL_DESCRIPTIONS.items()]
|
| label_token_count = sum(word_token_count(lbl) + 2 for lbl in labels_words)
|
| budget = max(64, chunk_tokens - label_token_count - safety_margin)
|
| return RecursiveWordChunker(budget_words=budget, max_chars=max_chars)
|
|
|
|
|
| def chunk_text(text: str) -> List[Tuple[str, int]]:
|
| """Split text into (chunk_text, start_index) using the recursive word chunker.
|
|
|
| Falls back to a plain whitespace chunker only if chonkie is unavailable.
|
| """
|
| try:
|
| return build_word_chunker(DEFAULT_CHUNK_TOKENS).chunk(text)
|
| except Exception:
|
| return manual_chunks(text, DEFAULT_CHUNK_CHARS)
|
| spans = []
|
| if not value:
|
| return spans
|
| start = 0
|
| while True:
|
| index = text.find(value, start)
|
| if index == -1:
|
| break
|
| spans.append((index, index + len(value)))
|
| start = index + max(1, len(value))
|
| return spans
|
|
|
|
|
| def normalize_predictions(raw_preds, text: str, emit_duplicates: bool = False) -> List[Dict]:
|
| spans = []
|
| if isinstance(raw_preds, dict):
|
| entities = raw_preds.get("entities", raw_preds)
|
| if isinstance(entities, list):
|
| entities = entities[0] if entities and isinstance(entities[0], dict) else {}
|
| if not isinstance(entities, dict):
|
| entities = {}
|
| for label, items in entities.items():
|
| if not isinstance(items, list):
|
| continue
|
| label = str(label).upper().strip()
|
| for item in items:
|
| if isinstance(item, dict):
|
| start = item.get("start")
|
| end = item.get("end")
|
| value = item.get("text") or item.get("value") or ""
|
| score = float(item.get("score", item.get("confidence", 1.0)) or 0.0)
|
| has_span = start is not None and end is not None and int(end) > int(start)
|
| if emit_duplicates and value:
|
| for match_start, match_end in find_all_occurrences(text, str(value)):
|
| spans.append({"label": label, "start": match_start, "end": match_end, "text": str(value), "score": score})
|
| elif has_span:
|
| start_int = int(start)
|
| end_int = int(end)
|
| spans.append({"label": label, "start": start_int, "end": end_int, "text": item.get("text", text[start_int:end_int]), "score": score})
|
| elif value:
|
| for match_start, match_end in find_all_occurrences(text, str(value)):
|
| spans.append({"label": label, "start": match_start, "end": match_end, "text": str(value), "score": score})
|
| elif isinstance(item, str):
|
| for match_start, match_end in find_all_occurrences(text, item):
|
| spans.append({"label": label, "start": match_start, "end": match_end, "text": item, "score": 0.0})
|
| elif isinstance(raw_preds, list):
|
| spans = list(raw_preds)
|
|
|
| best: Dict[tuple, Dict] = {}
|
| for span in spans:
|
| try:
|
| key = (span.get("label"), int(span.get("start")), int(span.get("end")))
|
| except Exception:
|
| continue
|
| if key not in best or float(span.get("score", 0) or 0) > float(best[key].get("score", 0) or 0):
|
| span["label"] = str(span.get("label", "")).upper().strip()
|
| span["start"] = key[1]
|
| span["end"] = key[2]
|
| span["text"] = str(span.get("text") or text[key[1]:key[2]])
|
| span["score"] = float(span.get("score", 0) or 0)
|
| best[key] = span
|
| return list(best.values())
|
|
|
|
|
| def clean_text_with_mapping(text: str) -> Tuple[str, List[int]]:
|
| if not text:
|
| return text, []
|
| original = unicodedata.normalize("NFC", text)
|
| cleaned_chars: List[str] = []
|
| char_map: List[int] = []
|
| for original_pos, char in enumerate(original):
|
| if char in STRIP_CHARS or C0_RE.match(char):
|
| continue
|
| if C1_RE.match(char):
|
| cleaned_chars.append(" ")
|
| char_map.append(original_pos)
|
| continue
|
| cleaned_chars.append(char)
|
| char_map.append(original_pos)
|
|
|
| cleaned = MULTI_SPACE_RE.sub(" ", "".join(cleaned_chars)).strip()
|
| final_map: List[int] = []
|
| map_index = 0
|
| for char in cleaned:
|
| while map_index < len(char_map) and (
|
| original[char_map[map_index]] != char
|
| and not (char == " " and original[char_map[map_index]].isspace())
|
| ):
|
| map_index += 1
|
| final_map.append(char_map[map_index] if map_index < len(char_map) else len(text))
|
| map_index += 1
|
| final_map.append(len(text))
|
| return cleaned, final_map
|
|
|
|
|
| def spans_overlap(left: Dict, right: Dict) -> bool:
|
| return max(int(left["start"]), int(right["start"])) < min(int(left["end"]), int(right["end"]))
|
|
|
|
|
| def resolve_span_conflicts(entities: List[Dict], strategy: str = "score_first") -> List[Dict]:
|
| if not entities:
|
| return []
|
|
|
| def sort_key(entity: Dict):
|
| is_rule = 1 if entity.get("source") == "rule" else 0
|
| score = 1.0 if is_rule else float(entity.get("score") or 0.0)
|
| length = int(entity["end"]) - int(entity["start"])
|
| priority = LABEL_PRIORITY.get(str(entity.get("label", "")), 0)
|
| if strategy == "rule_first":
|
| return (is_rule, priority, score, length)
|
| if strategy == "longest_first":
|
| return (is_rule, length, score, priority)
|
| return (is_rule, score, priority, length)
|
|
|
| ranked = sorted(entities, key=sort_key, reverse=True)
|
| kept: List[Dict] = []
|
| for entity in ranked:
|
| if not any(spans_overlap(entity, existing) for existing in kept):
|
| kept.append(entity)
|
| return sorted(kept, key=lambda item: int(item["start"]))
|
|
|
|
|
| def apply_span_validators(entities: List[Dict]) -> List[Dict]:
|
| kept = []
|
| for entity in entities:
|
| value = str(entity.get("text") or "")
|
| label = str(entity.get("label") or "").upper()
|
| if re.search(r"\*{2,}", value):
|
| continue
|
| if label in NUMERIC_MASK_LABELS and re.search(r"X{2,}", value):
|
| continue
|
| kept.append(entity)
|
| return kept
|
|
|
|
|
| def apply_postprocessing(entities: List[Dict], original_text: str) -> List[Dict]:
|
| entities = apply_span_validators(entities)
|
| sorted_entities = sorted(entities, key=lambda item: int(item["start"]))
|
| final_entities: List[Dict] = []
|
| for index, entity in enumerate(sorted_entities):
|
| label = entity.get("label")
|
| start = int(entity["start"])
|
| end = int(entity["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 index + 1 < len(sorted_entities):
|
| next_entity = sorted_entities[index + 1]
|
| gap = int(next_entity["start"]) - end
|
| if next_entity.get("label") == "PERSON" and 0 <= gap <= 3:
|
| has_person_after = True
|
| if not has_person_after:
|
| continue
|
| final_entities.append(entity)
|
| return final_entities
|
|
|
|
|
| def tokenize_chunk(text: str) -> List[Tuple[str, int, int]]:
|
| """Run WhitespaceTokenSplitter on text, return list of (token, start, end)."""
|
| try:
|
| from gliner2.processor import WhitespaceTokenSplitter
|
| splitter = WhitespaceTokenSplitter()
|
| return list(splitter(text, lower=False))
|
| except Exception:
|
| return []
|
|
|
|
|
| def run_inference(text: str, model_id: str, threshold: float, use_rules: bool, use_descriptions: bool, emit_duplicates: bool) -> Tuple[List[Dict], List[Dict], Dict]:
|
| """Returns (resolved_entities, raw_model_entities, debug_info)."""
|
| text = text.rstrip()
|
| debug_info = {"chunks": [], "tokens_per_chunk": [], "raw_preds_per_chunk": [], "token_merge_info": []}
|
| if not text or not text.strip():
|
| return [], [], debug_info
|
| model = load_model(model_id)
|
| schema = build_schema(model, use_descriptions)
|
| all_predictions: List[Dict] = []
|
|
|
| chunks = chunk_text(text)
|
| for chunk_idx, (chunk_text_value, chunk_start) in enumerate(chunks):
|
|
|
| debug_info["chunks"].append({"index": chunk_idx, "start": chunk_start, "end": chunk_start + len(chunk_text_value), "text": chunk_text_value})
|
|
|
|
|
| tokens = tokenize_chunk(chunk_text_value)
|
| debug_info["tokens_per_chunk"].append({"chunk_index": chunk_idx, "tokens": [(t, s, e) for t, s, e in tokens]})
|
|
|
| raw = model.extract(
|
| chunk_text_value,
|
| schema,
|
| threshold=threshold,
|
| include_spans=not emit_duplicates,
|
| include_confidence=not emit_duplicates,
|
| format_results=emit_duplicates,
|
| )
|
| predictions = normalize_predictions(raw, chunk_text_value, emit_duplicates=emit_duplicates)
|
|
|
|
|
| debug_info["raw_preds_per_chunk"].append({"chunk_index": chunk_idx, "predictions": [dict(p) for p in predictions]})
|
|
|
|
|
| for pred in predictions:
|
| covered_tokens = [(t, s, e) for t, s, e in tokens if s < pred["end"] and e > pred["start"]]
|
| debug_info["token_merge_info"].append({
|
| "chunk_index": chunk_idx,
|
| "label": pred.get("label", ""),
|
| "text": pred.get("text", ""),
|
| "start": pred["start"],
|
| "end": pred["end"],
|
| "score": pred.get("score", 0.0),
|
| "tokens_merged": covered_tokens,
|
| })
|
|
|
| for prediction in predictions:
|
| prediction["start"] += chunk_start
|
| prediction["end"] += chunk_start
|
| all_predictions.extend(predictions)
|
|
|
|
|
| raw_model_entities = [dict(p, source="model") for p in all_predictions]
|
|
|
| if use_rules:
|
| cleaned, char_map = clean_text_with_mapping(text)
|
| original_to_clean: Dict[int, int] = {}
|
| for clean_index, original_pos in enumerate(char_map):
|
| original_to_clean.setdefault(original_pos, clean_index)
|
|
|
| model_predictions_clean = []
|
| for prediction in all_predictions:
|
| clean_start = original_to_clean.get(prediction["start"])
|
| clean_end = original_to_clean.get(prediction["end"])
|
| if clean_end is None:
|
|
|
|
|
| for k in range(p["end"] + 1, len(text) + 2):
|
| clean_end = orig_to_clean.get(k)
|
| if clean_end is not None:
|
| break
|
| if clean_start is None or clean_end is None or clean_end <= clean_start:
|
| continue
|
| model_predictions_clean.append({
|
| "label": prediction["label"],
|
| "text": cleaned[clean_start:clean_end],
|
| "start": clean_start,
|
| "end": clean_end,
|
| "score": prediction.get("score", 0.0),
|
| "source": "model",
|
| })
|
|
|
| rule_predictions = [
|
| {**entity, "score": 1.0, "source": "rule"}
|
| for entity in detect_by_rules(cleaned, LABELS)
|
| ]
|
| combined = resolve_span_conflicts(rule_predictions + model_predictions_clean, strategy="rule_first")
|
| restored = []
|
| for entity in combined:
|
| start = int(entity["start"])
|
| end = int(entity["end"])
|
| original_start = char_map[start] if start < len(char_map) else None
|
| original_end = char_map[end - 1] + 1 if (end - 1) < len(char_map) else None
|
| if original_start is None or original_end is None or original_end <= original_start:
|
| continue
|
| restored.append({
|
| "label": entity["label"],
|
| "start": original_start,
|
| "end": original_end,
|
| "text": text[original_start:original_end],
|
| "score": float(entity.get("score", 1.0)),
|
| "source": entity.get("source", "model"),
|
| })
|
| all_predictions = restored
|
| else:
|
| for prediction in all_predictions:
|
| prediction.setdefault("source", "model")
|
| all_predictions = resolve_span_conflicts(all_predictions)
|
|
|
| return apply_postprocessing(all_predictions, text), raw_model_entities, debug_info
|
|
|
|
|
| def extract_with_llamaparse(path: str) -> List[str]:
|
| """Use LlamaParse (cloud) to parse a PDF or DOCX file — returns one string per document/page."""
|
| api_key = os.getenv("LLAMA_CLOUD_API_KEY") or os.getenv("LLAMA_API_KEY")
|
| if not api_key:
|
| raise RuntimeError("Set LLAMA_CLOUD_API_KEY in environment to use LlamaParse")
|
|
|
| try:
|
| from llama_parse import LlamaParse
|
| except Exception:
|
| try:
|
| from llama_cloud_services import LlamaParse
|
| except Exception as exc:
|
| raise RuntimeError("Install llama-parse to use LlamaParse parsing") from exc
|
|
|
| parser = LlamaParse(api_key=api_key, result_type="markdown", verbose=False)
|
| documents = parser.load_data(path)
|
| if not documents:
|
| return [""]
|
| return [getattr(document, "text", "") or "" for document in documents]
|
|
|
|
|
| def extract_with_llamaindex_local(path: str, suffix: str) -> List[str]:
|
| """Use LlamaIndex local readers to parse a PDF or DOCX file — returns one string per page/document."""
|
| if suffix == ".docx":
|
| from llama_index.readers.file import DocxReader
|
| docs = DocxReader().load_data(Path(path))
|
| return ["\n".join(d.text for d in docs)] if docs else [""]
|
|
|
| if suffix == ".pdf":
|
| for cls_name in ("PyMuPDFReader", "PDFReader"):
|
| try:
|
| module = __import__("llama_index.readers.file", fromlist=[cls_name])
|
| reader_cls = getattr(module, cls_name)
|
| docs = reader_cls().load_data(Path(path))
|
| if any(d.text.strip() for d in docs):
|
| return [d.text for d in docs]
|
| except Exception:
|
| continue
|
| return [""]
|
|
|
| raise ValueError(f"Unsupported file type: {suffix}")
|
|
|
|
|
| def extract_input_text(text_input: str, file_obj, pdf_backend: str) -> Tuple[str, str, List[int]]:
|
| if file_obj is None:
|
| return text_input or "", "text", []
|
|
|
| file_path = Path(file_obj.name if hasattr(file_obj, "name") else str(file_obj))
|
| suffix = file_path.suffix.lower()
|
| if suffix not in (".docx", ".pdf"):
|
| raise ValueError("Only .docx and .pdf files are supported")
|
|
|
| if pdf_backend == "llamaparse":
|
| pages = extract_with_llamaparse(str(file_path))
|
| else:
|
| pages = extract_with_llamaindex_local(str(file_path), suffix)
|
|
|
| if suffix == ".docx":
|
| return "\n".join(pages), f"docx: {file_path.name}", []
|
|
|
|
|
| offsets: List[int] = []
|
| parts: List[str] = []
|
| cursor = 0
|
| for page_index, page_text in enumerate(pages, start=1):
|
| if page_index > 1:
|
| separator = "\n\n"
|
| parts.append(separator)
|
| cursor += len(separator)
|
| offsets.append(cursor)
|
| parts.append(page_text)
|
| cursor += len(page_text)
|
| return "".join(parts), f"pdf: {file_path.name}", offsets
|
|
|
|
|
| def add_page_numbers(entities: List[Dict], page_offsets: List[int]) -> List[Dict]:
|
| if not page_offsets:
|
| return entities
|
| enriched = []
|
| for entity in entities:
|
| page = 1
|
| for index, offset in enumerate(page_offsets, start=1):
|
| if entity["start"] >= offset:
|
| page = index
|
| else:
|
| break
|
| enriched.append({**entity, "page": page})
|
| return enriched
|
|
|
|
|
| def entities_to_table(entities: List[Dict]) -> List[List]:
|
| rows = []
|
| for index, entity in enumerate(entities, start=1):
|
| rows.append([
|
| index,
|
| entity.get("label", ""),
|
| entity.get("text", ""),
|
| int(entity.get("start", 0)),
|
| int(entity.get("end", 0)),
|
| round(float(entity.get("score", 0.0)), 4),
|
| entity.get("source", "model"),
|
| entity.get("page", ""),
|
| ])
|
| return rows
|
|
|
|
|
| def make_highlighted_text(text: str, entities: List[Dict]):
|
| if not text:
|
| return []
|
| valid_entities = sorted(
|
| [entity for entity in entities if 0 <= int(entity.get("start", -1)) < int(entity.get("end", -1)) <= len(text)],
|
| key=lambda item: int(item["start"]),
|
| )
|
| highlights = []
|
| cursor = 0
|
| for entity in valid_entities:
|
| start = int(entity["start"])
|
| end = int(entity["end"])
|
| if start > cursor:
|
| highlights.append((text[cursor:start], None))
|
| highlights.append((text[start:end], entity.get("label", "PII")))
|
| cursor = end
|
| if cursor < len(text):
|
| highlights.append((text[cursor:], None))
|
| return highlights
|
|
|
|
|
| def write_json_result(payload: Dict) -> str:
|
| temp_dir = tempfile.mkdtemp(prefix="guardpii_demo_")
|
| output_path = os.path.join(temp_dir, "pii_predictions.json")
|
| with open(output_path, "w", encoding="utf-8") as output_file:
|
| json.dump(payload, output_file, ensure_ascii=False, indent=2)
|
| return output_path
|
|
|
|
|
| def detect_pii(text_input, file_obj, model_id, threshold, use_rules, use_descriptions, emit_duplicates, pdf_backend):
|
| text, source, page_offsets = extract_input_text(text_input, file_obj, pdf_backend)
|
| if not text.strip():
|
| empty_payload = {"source": source, "num_entities": 0, "entities": [], "text": text}
|
| return [], [], empty_payload, write_json_result(empty_payload), "No text was extracted.", "", "", "", ""
|
|
|
| entities, raw_model_entities, debug_info = run_inference(
|
| text=text,
|
| model_id=(model_id or DEFAULT_MODEL_ID).strip(),
|
| threshold=float(threshold),
|
| use_rules=bool(use_rules),
|
| use_descriptions=bool(use_descriptions),
|
| emit_duplicates=bool(emit_duplicates),
|
| )
|
| entities = add_page_numbers(entities, page_offsets)
|
| raw_model_entities = add_page_numbers(raw_model_entities, page_offsets)
|
| payload = {
|
| "source": source,
|
| "num_chars": len(text),
|
| "num_entities": len(entities),
|
| "model": (model_id or DEFAULT_MODEL_ID).strip(),
|
| "threshold": float(threshold),
|
| "use_rules": bool(use_rules),
|
| "entities": entities,
|
| "text": text,
|
| }
|
| status = f"Found {len(entities)} entities in {source}."
|
|
|
|
|
| chunks_md = "## Chunks Created\n\n"
|
| for chunk in debug_info["chunks"]:
|
| chunks_md += f"### Chunk {chunk['index']} (chars {chunk['start']}–{chunk['end']}, len={chunk['end'] - chunk['start']})\n"
|
| preview = chunk["text"][:500] + ("..." if len(chunk["text"]) > 500 else "")
|
| chunks_md += f"```\n{preview}\n```\n\n"
|
|
|
|
|
| tokens_md = "## WhitespaceTokenSplitter Output\n\n"
|
| for item in debug_info["tokens_per_chunk"]:
|
| tokens_md += f"### Chunk {item['chunk_index']} — {len(item['tokens'])} tokens\n\n"
|
| tokens_md += "| # | Token | Start | End |\n|---|-------|-------|-----|\n"
|
| for idx, (tok, s, e) in enumerate(item["tokens"][:200]):
|
| tok_display = tok.replace("|", "\\|")
|
| tokens_md += f"| {idx} | `{tok_display}` | {s} | {e} |\n"
|
| if len(item["tokens"]) > 200:
|
| tokens_md += f"\n... và {len(item['tokens']) - 200} tokens nữa\n"
|
| tokens_md += "\n"
|
|
|
|
|
| merge_md = "## Token → Prediction Merge\n\n"
|
| merge_md += "Hiển thị cách model predict trên token spans và ghép lại thành entity:\n\n"
|
| for info in debug_info["token_merge_info"]:
|
| tokens_str = " + ".join([f"`{t}`" for t, s, e in info["tokens_merged"]])
|
| merge_md += f"- **[{info['label']}]** \"{info['text']}\" (chunk {info['chunk_index']}, chars {info['start']}–{info['end']}, score={info['score']:.4f})\n"
|
| merge_md += f" - Tokens merged: {tokens_str}\n\n"
|
|
|
|
|
| per_chunk_md = "## Per-Chunk Model Predictions\n\n"
|
| for item in debug_info["raw_preds_per_chunk"]:
|
| per_chunk_md += f"### Chunk {item['chunk_index']} — {len(item['predictions'])} predictions\n\n"
|
| if item["predictions"]:
|
| per_chunk_md += "| # | Label | Text | Start | End | Score |\n|---|-------|------|-------|-----|-------|\n"
|
| for idx, pred in enumerate(item["predictions"], 1):
|
| text_display = pred.get("text", "").replace("|", "\\|")[:80]
|
| per_chunk_md += f"| {idx} | {pred.get('label', '')} | {text_display} | {pred.get('start', '')} | {pred.get('end', '')} | {pred.get('score', 0.0):.4f} |\n"
|
| per_chunk_md += "\n"
|
| else:
|
| per_chunk_md += "Không có prediction nào.\n\n"
|
|
|
| return make_highlighted_text(text, entities), entities_to_table(raw_model_entities), payload, write_json_result(payload), status, chunks_md, tokens_md, merge_md, per_chunk_md
|
|
|
|
|
| with gr.Blocks(title="GuardPII GLiNER2 Demo") as demo:
|
| gr.Markdown("# GuardPII GLiNER2 PII Demo")
|
| gr.Markdown("Detect PII from text, DOCX, or PDF files with GLiNER2 plus the high-precision rulebase.")
|
|
|
| with gr.Row():
|
| with gr.Column(scale=2):
|
| text_input = gr.Textbox(
|
| label="Text input",
|
| lines=10,
|
| placeholder="Paste text here, or upload a DOCX/PDF below.",
|
| )
|
| file_input = gr.File(label="DOCX or PDF input", file_types=[".docx", ".pdf"])
|
| run_button = gr.Button("Detect PII", variant="primary")
|
| with gr.Column(scale=1):
|
| model_id = gr.Textbox(label="Model ID", value=DEFAULT_MODEL_ID)
|
| threshold = gr.Slider(label="Threshold", minimum=0.05, maximum=0.95, step=0.05, value=DEFAULT_THRESHOLD)
|
| use_rules = gr.Checkbox(label="Use rulebase gap-filler", value=True)
|
| use_descriptions = gr.Checkbox(label="Use label descriptions", value=False)
|
| emit_duplicates = gr.Checkbox(label="Value broadcast mode", value=False)
|
| pdf_backend = gr.Dropdown(
|
| label="Document parser",
|
| choices=["llamaparse", "llamaindex"],
|
| value=os.getenv("PDF_BACKEND", "llamaparse"),
|
| )
|
|
|
| status = gr.Textbox(label="Status", interactive=False)
|
| highlighted = gr.HighlightedText(label="Highlighted text", combine_adjacent=True, show_legend=True)
|
| table = gr.Dataframe(
|
| label="Entities",
|
| headers=["#", "label", "text", "start", "end", "score", "source", "page"],
|
| datatype=["number", "str", "str", "number", "number", "number", "str", "str"],
|
| wrap=True,
|
| )
|
| json_output = gr.JSON(label="JSON result")
|
| json_file = gr.File(label="Download JSON")
|
|
|
| gr.Markdown("---")
|
| gr.Markdown("# Debug / Pipeline Visualization")
|
|
|
| with gr.Accordion("1. Chunks (text splitting)", open=False):
|
| debug_chunks = gr.Markdown(value="")
|
|
|
| with gr.Accordion("2. WhitespaceTokenSplitter Output", open=False):
|
| debug_tokens = gr.Markdown(value="")
|
|
|
| with gr.Accordion("3. Token → Prediction Merge (how tokens are joined into entities)", open=False):
|
| debug_merge = gr.Markdown(value="")
|
|
|
| with gr.Accordion("4. Per-Chunk Model Predictions", open=False):
|
| debug_per_chunk = gr.Markdown(value="")
|
|
|
| run_button.click(
|
| detect_pii,
|
| inputs=[text_input, file_input, model_id, threshold, use_rules, use_descriptions, emit_duplicates, pdf_backend],
|
| outputs=[highlighted, table, json_output, json_file, status, debug_chunks, debug_tokens, debug_merge, debug_per_chunk],
|
| )
|
|
|
|
|
| if __name__ == "__main__":
|
| demo.launch() |
