app/services/gliner_service.py

"""Local GLiNER2 extraction service for the ``/json/no-ai`` route.

GLiNER2 performs zero-shot entity / structured-JSON extraction entirely
on-device -- no external AI/LLM API is called, hence the "no-ai" route name.

The model is loaded ONCE and cached for the life of the process (warmed at
application startup via the lifespan hook in ``app/api/server.py`` and lazily
loaded on first use if startup warm-up failed). Inference is CPU-bound and is
dispatched on the shared thread pool by the route layer; a bounded semaphore
serializes concurrent forwards on the shared model instance so a burst of
requests cannot oversubscribe the CPU or the model's memory buffers.

Inference backend: the ONNX export ``lion-ai/gliner2-base-v1-onnx`` running on
ONNX Runtime via ``GLiNER2ONNXEngine`` (see ``gliner_onnx.py``). This removes
the PyTorch dependency from the extraction hot path.
"""

from __future__ import annotations

import threading
import time
from typing import Any, Dict, List, Optional

from app.config import get_settings
from app.core.logger import get_logger
from app.services.gliner_onnx import GLiNER2ONNXEngine

logger = get_logger(__name__)

DEFAULT_MODEL_ID = "lion-ai/gliner2-base-v1-onnx"
DEFAULT_DEVICE = "cpu"
DEFAULT_MAX_CONCURRENT = 2
DEFAULT_MAX_CONTENT_LENGTH = 100_000


class GLiNERService:
    """Cached, concurrency-bounded wrapper around the GLiNER2 ONNX engine."""

    def __init__(
        self,
        model_id: str = DEFAULT_MODEL_ID,
        device: str = DEFAULT_DEVICE,
        max_concurrent: int = DEFAULT_MAX_CONCURRENT,
        max_content_length: int = DEFAULT_MAX_CONTENT_LENGTH,
    ) -> None:
        self._model_id = model_id
        self._device = device
        self._max_content_length = max_content_length
        self._max_concurrent = max(1, max_concurrent)
        self._engine: Optional[GLiNER2ONNXEngine] = None
        self._lock = threading.Lock()
        self._semaphore = threading.BoundedSemaphore(self._max_concurrent)
        self._load_error: Optional[str] = None

    # ------------------------------------------------------------------ #
    # Lifecycle
    # ------------------------------------------------------------------ #

    def load_model(self) -> None:
        """Load and cache the GLiNER2 ONNX model. Idempotent and thread-safe."""
        if self.is_loaded():
            return
        with self._lock:
            if self.is_loaded():
                return
            t0 = time.perf_counter()
            try:
                engine = GLiNER2ONNXEngine(model_id=self._model_id, device=self._device)
                engine.load()
                self._engine = engine
                self._load_error = None
                logger.info(
                    "GLiNER2 ONNX model loaded and cached (%s, device=%s) in %.2fs",
                    self._model_id,
                    self._device,
                    time.perf_counter() - t0,
                )
            except Exception as exc:  # noqa: BLE001
                self._load_error = str(exc)
                logger.exception("GLiNER2 ONNX model load failed")
                raise

    def unload(self) -> None:
        """Release the cached model. Used at shutdown."""
        with self._lock:
            self._engine = None
            self._load_error = None

    def is_loaded(self) -> bool:
        return self._engine is not None

    def load_error(self) -> Optional[str]:
        return self._load_error

    @property
    def model_id(self) -> str:
        return self._model_id

    @property
    def device(self) -> str:
        return self._device

    @property
    def max_content_length(self) -> int:
        return self._max_content_length

    # ------------------------------------------------------------------ #
    # Inference (called on the shared thread pool)
    # ------------------------------------------------------------------ #

    def extract_json(self, text: str, structure: Dict[str, Any], threshold: float = 0.5) -> Dict[str, Any]:
        """Structured JSON extraction with the cached local ONNX model."""
        self._ensure_loaded()
        with self._semaphore:
            return self._engine.extract_json(text, structure, threshold=threshold)

    def extract_entities(self, text: str, labels: List[str], threshold: float = 0.5) -> Dict[str, Any]:
        """Zero-shot entity extraction with the cached local ONNX model."""
        self._ensure_loaded()
        with self._semaphore:
            return self._engine.extract_entities(text, labels, threshold=threshold)

    def _ensure_loaded(self) -> None:
        if not self.is_loaded():
            self.load_model()


_settings = get_settings()

gliner_service = GLiNERService(
    model_id=_settings.gliner_model,
    device=_settings.gliner_device,
    max_concurrent=_settings.gliner_max_concurrent,
    max_content_length=_settings.gliner_max_content_length,
)


# =============================================================================
# DEPRECATED -- PyTorch (``gliner2`` / torch) backend.
# -----------------------------------------------------------------------------
# Kept for reference and easy rollback / later fallback implementation. The old
# implementation loaded ``fastino/gliner2-base-v1`` via the ``gliner2`` PyTorch
# library, which pulled torch into the extraction hot path and required
# redirecting stdout to swallow the emoji config banner.
#
# To restore the PyTorch backend:
#   1. Set GLINER_MODEL back to a PyTorch repo (e.g. fastino/gliner2-base-v1)
#      in config/.env/.env.example and the Dockerfile.
#   2. Un-comment the whole block below and delete/rename the ONNX
#      ``GLiNERService`` above (or point the module-level ``gliner_service``
#      instantiation at the deprecated class).
# =============================================================================
# from __future__ import annotations
#
# import contextlib
# import io
# import threading
# import time
# from typing import Any, Dict, List, Optional
#
# from app.config import get_settings
# from app.core.logger import get_logger
#
# logger = get_logger(__name__)
#
# DEFAULT_MODEL_ID = "fastino/gliner2-base-v1"
# DEFAULT_DEVICE = "cpu"
# DEFAULT_MAX_CONCURRENT = 2
# DEFAULT_MAX_CONTENT_LENGTH = 100_000
#
#
# class GLiNERService:
#     """Cached, concurrency-bounded wrapper around a GLiNER2 model instance."""
#
#     def __init__(
#         self,
#         model_id: str = DEFAULT_MODEL_ID,
#         device: str = DEFAULT_DEVICE,
#         max_concurrent: int = DEFAULT_MAX_CONCURRENT,
#         max_content_length: int = DEFAULT_MAX_CONTENT_LENGTH,
#     ) -> None:
#         self._model_id = model_id
#         self._device = device
#         self._max_content_length = max_content_length
#         self._max_concurrent = max(1, max_concurrent)
#         self._model: Any = None
#         self._lock = threading.Lock()
#         self._semaphore = threading.BoundedSemaphore(self._max_concurrent)
#         self._load_error: Optional[str] = None
#
#     # ------------------------------------------------------------------ #
#     # Lifecycle
#     # ------------------------------------------------------------------ #
#
#     def load_model(self) -> None:
#         """Load and cache the GLiNER2 model. Idempotent and thread-safe."""
#         if self.is_loaded():
#             return
#         with self._lock:
#             if self.is_loaded():
#                 return
#             t0 = time.perf_counter()
#             try:
#                 from gliner2 import GLiNER2
#
#                 # GLiNER2 prints an emoji config banner to stdout on load, which
#                 # crashes consoles with a non-UTF-8 encoding (e.g. Windows cp1252).
#                 # Swallow it so loading works everywhere.
#                 with contextlib.redirect_stdout(io.StringIO()):
#                     self._model = GLiNER2.from_pretrained(self._model_id, map_location=self._device)
#                 self._load_error = None
#                 logger.info(
#                     "GLiNER2 model loaded and cached (%s, device=%s) in %.2fs",
#                     self._model_id,
#                     self._device,
#                     time.perf_counter() - t0,
#                 )
#             except Exception as exc:  # noqa: BLE001
#                 self._load_error = str(exc)
#                 logger.exception("GLiNER2 model load failed")
#                 raise
#
#     def unload(self) -> None:
#         """Release the cached model (frees ~1.3 GB). Used at shutdown."""
#         with self._lock:
#             self._model = None
#             self._load_error = None
#
#     def is_loaded(self) -> bool:
#         return self._model is not None
#
#     def load_error(self) -> Optional[str]:
#         return self._load_error
#
#     @property
#     def model_id(self) -> str:
#         return self._model_id
#
#     @property
#     def device(self) -> str:
#         return self._device
#
#     @property
#     def max_content_length(self) -> int:
#         return self._max_content_length
#
#     # ------------------------------------------------------------------ #
#     # Inference (called on the shared thread pool)
#     # ------------------------------------------------------------------ #
#
#     def extract_json(self, text: str, structure: Dict[str, Any], threshold: float = 0.5) -> Dict[str, Any]:
#         """Structured JSON extraction with the cached local model."""
#         self._ensure_loaded()
#         with self._semaphore:
#             return self._model.extract_json(text, structure, threshold=threshold)
#
#     def extract_entities(self, text: str, labels: List[str], threshold: float = 0.5) -> Dict[str, Any]:
#         """Zero-shot entity extraction with the cached local model."""
#         self._ensure_loaded()
#         with self._semaphore:
#             return self._model.extract_entities(text, labels, threshold=threshold)
#
#     def _ensure_loaded(self) -> None:
#         if not self.is_loaded():
#             self.load_model()
#
#
# _settings = get_settings()
#
# gliner_service = GLiNERService(
#     model_id=_settings.gliner_model,
#     device=_settings.gliner_device,
#     max_concurrent=_settings.gliner_max_concurrent,
#     max_content_length=_settings.gliner_max_content_length,
# )
# =============================================================================