ner.py

import sys
from typing import Any, Dict, List, Union
from exception import CustomExceptionHandling
from logger import logging
from gliner import GLiNER
import gradio as gr

# Load the model once per process
model = GLiNER.from_pretrained("nvidia/gliner-PII")


def pii_ner(
    text: str,
    labels: str,
    threshold: float,
    nested_ner: bool,
    multi_label: bool,
) -> Dict[str, Any]:
    """Perform PII detection on the given text."""
    try:
        if not text:
            raise gr.Error("No text provided")
        if not labels:
            raise gr.Error("No labels provided")
            
        # Parse and clean labels
        label_list: List[str] = [
            l.strip() for l in labels.split(",") if l.strip()
        ]

        entities = model.predict_entities(
            text,
            label_list,
            flat_ner=not nested_ner,
            threshold=threshold,
            multi_label=multi_label,
        )

        result = {
            "text": text,
            "entities": [
                {
                    "entity": entity["label"],
                    "word": entity["text"],
                    "start": entity["start"],
                    "end": entity["end"],
                    "score": float(entity.get("score", 0.0)),
                }
                for entity in entities
            ],
        }

        logging.info("PII detection completed successfully")
        return result

    except Exception as e:
        # Custom exception handling
        raise CustomExceptionHandling(e, sys) from e