| """Local native GLiNER runtime. Character offsets always refer to untouched text.""" |
|
|
| import math |
| import re |
| import threading |
|
|
| from ocr.schemas import SafeError |
|
|
| from .model import MODEL_REVISION, verify_cache |
| from .normalization import dehyphenate |
| from .replacement import validate_entities |
| from .schemas import LABELS, Entity |
|
|
| DEFAULT_PROMPTS = { |
| "PERSON": "person", |
| |
| |
| "CITY": "location", |
| "RELIGION": "religion", |
| "EMAIL": "email address", |
| "PHONE": "phone number", |
| "ORGANIZATION": "organization", |
| } |
|
|
|
|
| class GLiNERRecognizer: |
| local = True |
| revision = MODEL_REVISION |
| categories = LABELS |
|
|
| def __init__(self, cache, *, device="cpu", threshold=0.5, prompts=None, threads=2): |
| self._lock = threading.Lock() |
| try: |
| self.prompts = dict(DEFAULT_PROMPTS if prompts is None else prompts) |
| except Exception: |
| raise SafeError("unsupported_required_categories") from None |
| if set(self.prompts) != set(LABELS) or ( |
| any( |
| not isinstance(v, str) or not v.strip() or len(v) > 80 |
| for v in self.prompts.values() |
| ) |
| or len(set(self.prompts.values())) != len(LABELS) |
| ): |
| raise SafeError("unsupported_required_categories") |
| if ( |
| type(threshold) not in (int, float) |
| or not math.isfinite(threshold) |
| or not 0 < threshold < 1 |
| ): |
| raise SafeError("invalid_recognizer_configuration") |
| if device not in ("cpu", "cuda") or type(threads) is not int or not 1 <= threads <= 16: |
| raise SafeError("invalid_recognizer_configuration") |
| self.threshold = threshold |
| try: |
| path = verify_cache(cache) |
| import torch |
| from gliner import GLiNER |
|
|
| torch.set_num_threads(threads) |
| self.model = GLiNER.from_pretrained( |
| str(path), |
| local_files_only=True, |
| map_location=device, |
| _attn_implementation="eager", |
| strict=True, |
| ) |
| self.model.eval() |
| except Exception: |
| raise SafeError("anonymization_model_unavailable") from None |
|
|
| def recognize(self, text): |
| if not isinstance(text, str) or len(text) > 2_000_000: |
| raise SafeError("anonymization_text_limit") |
| normalized, positions = dehyphenate(text) |
| predictions = self._recognize_view(normalized) |
| if normalized == text: |
| return predictions |
| mapped = [] |
| for entity in predictions: |
| start = positions[entity.start] |
| end = positions[entity.end - 1] + 1 |
| mapped.append( |
| entity.model_copy( |
| update={ |
| "start": start, |
| "end": end, |
| "text": text[start:end], |
| "normalization": "dehyphenation-v1", |
| } |
| ) |
| ) |
| return validate_entities(text, mapped) |
|
|
| def _recognize_view(self, text): |
| if not isinstance(text, str) or len(text) > 2_000_000: |
| raise SafeError("anonymization_text_limit") |
| |
| |
| words = list(re.finditer(r"\S+", text)) |
| if any(len(w[0]) > 256 for w in words): |
| raise SafeError("anonymization_token_limit") |
| entities = [] |
| reverse = {v: k for k, v in self.prompts.items()} |
| try: |
| import torch |
|
|
| with self._lock, torch.inference_mode(): |
| for start in range(0, len(words), 96): |
| window = words[start : start + 128] |
| left, right = window[0].start(), window[-1].end() |
| chunk = text[left:right] |
| encoded = self.model.data_processor.transformer_tokenizer(chunk)["input_ids"] |
| if len(encoded) > 1536: |
| raise SafeError("anonymization_token_limit") |
| raw = self.model.predict_entities( |
| chunk, |
| list(self.prompts.values()), |
| threshold=self.threshold, |
| flat_ner=True, |
| ) |
| if not isinstance(raw, list): |
| raise ValueError() |
| for item in raw: |
| if not isinstance(item, dict) or set(item) != { |
| "start", |
| "end", |
| "text", |
| "label", |
| "score", |
| }: |
| raise ValueError() |
| |
| local = Entity( |
| start=item["start"], |
| end=item["end"], |
| text=item["text"], |
| label=reverse[item["label"]], |
| score=item["score"], |
| source="multilingual_ner", |
| model_revision=self.revision, |
| ) |
| validate_entities(chunk, [local]) |
| entities.append( |
| local.model_copy( |
| update={"start": local.start + left, "end": local.end + left} |
| ) |
| ) |
| return validate_entities(text, entities) |
| except SafeError: |
| raise |
| except Exception: |
| raise SafeError("anonymization_inference_failed") from None |
|
|
