backend
backend
¶
Native GLiNER2 loading and inference for the notebook server.
Functions:
| Name | Description |
|---|---|
resolve_device |
Resolve and validate a notebook inference device. |
load_model |
Download the pinned checkpoint and load it on |
overlap_policy |
Translate the existing server flag to GLiNER2 overlap semantics. |
flatten_result |
Convert GLiNER2's label-keyed result into Anonymizer entity dictionaries. |
detect_entities_for_texts |
Detect entities while retaining the existing character-chunk contract. |
resolve_device(requested)
¶
Resolve and validate a notebook inference device.
Source code in src/anonymizer/notebooks/local_inference/gliner2/backend.py
def resolve_device(requested: str) -> str:
"""Resolve and validate a notebook inference device."""
import torch # ty: ignore[unresolved-import] -- installed in the isolated server environment
normalized = requested.strip().lower()
if normalized not in SUPPORTED_DEVICES:
raise ValueError(f"Unsupported GLiNER2 device {requested!r}; expected one of {SUPPORTED_DEVICES!r}.")
if normalized == "auto":
if torch.cuda.is_available():
return "cuda"
if torch.backends.mps.is_available():
return "mps"
return "cpu"
if normalized == "cuda" and not torch.cuda.is_available():
raise RuntimeError(
"CUDA was requested for local GLiNER2, but it is unavailable. "
"In Colab, select Runtime > Change runtime type > GPU and retry."
)
if normalized == "mps" and not torch.backends.mps.is_available():
raise RuntimeError("MPS was requested for local GLiNER2, but it is unavailable on this machine.")
return normalized
load_model(device)
¶
Download the pinned checkpoint and load it on device.
Source code in src/anonymizer/notebooks/local_inference/gliner2/backend.py
def load_model(device: str) -> Any:
"""Download the pinned checkpoint and load it on ``device``."""
from gliner2 import GLiNER2 # ty: ignore[unresolved-import] -- isolated server dependency
from huggingface_hub import snapshot_download
snapshot_path = snapshot_download(repo_id=MODEL_ID, revision=MODEL_REVISION)
compatible_snapshot = _prepare_transformers_v4_snapshot(Path(snapshot_path))
return GLiNER2.from_pretrained(compatible_snapshot, map_location=device)
overlap_policy(flat_ner)
¶
Translate the existing server flag to GLiNER2 overlap semantics.
Source code in src/anonymizer/notebooks/local_inference/gliner2/backend.py
def overlap_policy(flat_ner: bool) -> str:
"""Translate the existing server flag to GLiNER2 overlap semantics."""
return "disallow" if flat_ner else "longest"
flatten_result(result)
¶
Convert GLiNER2's label-keyed result into Anonymizer entity dictionaries.
Source code in src/anonymizer/notebooks/local_inference/gliner2/backend.py
def flatten_result(result: object) -> list[dict[str, Any]]:
"""Convert GLiNER2's label-keyed result into Anonymizer entity dictionaries."""
if not isinstance(result, dict):
raise ValueError(f"Unexpected GLiNER2 result type: {type(result)!r}")
entities_by_label = result.get("entities", {})
if not isinstance(entities_by_label, dict):
raise ValueError("Unexpected GLiNER2 result: 'entities' must be a mapping.")
flattened: list[dict[str, Any]] = []
for label, values in entities_by_label.items():
if not isinstance(values, list):
raise ValueError(f"Unexpected GLiNER2 values for label {label!r}: expected a list.")
for value in values:
if not isinstance(value, dict):
raise ValueError(f"Unexpected GLiNER2 entity for label {label!r}: {value!r}")
entity = {
"text": str(value["text"]),
"label": str(label),
"start": int(value["start"]),
"end": int(value["end"]),
"score": float(value["confidence"]),
}
flattened.append(entity)
return flattened
detect_entities_for_texts(model, texts, labels, *, threshold, chunk_length, overlap, flat_ner, inference_batch_size)
¶
Detect entities while retaining the existing character-chunk contract.
Source code in src/anonymizer/notebooks/local_inference/gliner2/backend.py
def detect_entities_for_texts(
model: Any,
texts: list[str],
labels: list[str],
*,
threshold: float,
chunk_length: int,
overlap: int,
flat_ner: bool,
inference_batch_size: int,
) -> list[list[dict[str, Any]]]:
"""Detect entities while retaining the existing character-chunk contract."""
if not labels:
return [[] for _ in texts]
_validate_chunk_params(chunk_length, overlap)
if inference_batch_size < 1:
raise ValueError("inference_batch_size must be >= 1")
chunk_records: list[tuple[int, int, str]] = []
for text_index, source_text in enumerate(texts):
for chunk, offset in _create_text_chunks(source_text, chunk_length, overlap):
chunk_records.append((text_index, offset, chunk))
results: list[list[dict[str, Any]]] = [[] for _ in texts]
if not chunk_records:
return results
raw_results = model.batch_extract_entities(
[chunk for _, _, chunk in chunk_records],
labels,
batch_size=inference_batch_size,
threshold=threshold,
include_confidence=True,
include_spans=True,
overlap_policy=overlap_policy(flat_ner),
)
if not isinstance(raw_results, list) or len(raw_results) != len(chunk_records):
raise ValueError("Unexpected GLiNER2 batch result shape.")
for (text_index, offset, _), raw_result in zip(chunk_records, raw_results, strict=True):
for entity in flatten_result(raw_result):
entity["start"] += offset
entity["end"] += offset
results[text_index].append(entity)
policy = overlap_policy(flat_ner)
return [_resolve_document_overlaps(entities, policy) for entities in results]