Skip to content

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 device.

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]