From: Elad Kapusta <redacted>
Date: Sat, 27 Jun 2026 07:23:34 +0000 (+0300)
Subject: Improve GLiNER detector caching (#119)
X-Git-Tag: v0.7.1~1
X-Git-Url: http://git.99rst.org/?a=commitdiff_plain;h=9de255dc805616f4b17705b4a6dadde2da148a7a;p=sgasser-llm-shield.git

Improve GLiNER detector caching (#119)
---

diff --git a/detector/detector/gliner_layer.py b/detector/detector/gliner_layer.py
index e9290c5..be58f10 100644
--- a/detector/detector/gliner_layer.py
+++ b/detector/detector/gliner_layer.py
@@ -11,10 +11,14 @@ from __future__ import annotations
 import os
 import re
 import threading
+from functools import lru_cache
 from typing import Any
 
 from .entities import LOCATION, PERSON, Span
 
+# Size of the LRU cache, in units of WINDOW_SIZE input to the model
+MODEL_INFERENCE_CACHE_SIZE = 4096
+
 DEFAULT_MODEL = "urchade/gliner_multi_pii-v1"
 
 
@@ -104,6 +108,11 @@ def load_model() -> None:
         _model = GLiNER.from_pretrained(_model_name())
 
 
+@lru_cache(maxsize=MODEL_INFERENCE_CACHE_SIZE)
+def _predict_window(text: str):
+    return _model.predict_entities(text, _PREDICT_LABELS, threshold=max(0.0, _PREDICT_FLOOR))
+
+
 def detect_gliner(text: str, score_threshold: float = 0.0) -> list[Span]:
     if not text:
         return []
@@ -114,9 +123,7 @@ def detect_gliner(text: str, score_threshold: float = 0.0) -> list[Span]:
     best: dict[tuple[int, int, str], float] = {}
     with _infer_lock:
         for offset, sub in _windows(text):
-            for ent in _model.predict_entities(
-                sub, _PREDICT_LABELS, threshold=max(0.0, _PREDICT_FLOOR)
-            ):
+            for ent in _predict_window(sub):
                 key = (offset + int(ent["start"]), offset + int(ent["end"]), ent["label"])
                 score = float(ent["score"])
                 if score > best.get(key, -1.0):
diff --git a/detector/tests/test_gliner_layer.py b/detector/tests/test_gliner_layer.py
index 1061158..55a255a 100644
--- a/detector/tests/test_gliner_layer.py
+++ b/detector/tests/test_gliner_layer.py
@@ -4,6 +4,8 @@ The role-noun suppressor labels and per-label floors are the precision layer; th
 full model integration is covered by benchmarks/pii-accuracy.
 """
 
+from unittest.mock import Mock
+
 import pytest
 
 from detector.gliner_layer import (
@@ -11,7 +13,9 @@ from detector.gliner_layer import (
     _SUPPRESS_LABELS,
     _TOKEN_RE,
     PER_LABEL_FLOOR,
+    Span,
     _windows,
+    detect_gliner,
 )
 
 
@@ -55,3 +59,34 @@ def test_per_label_floors_present_and_ordered():
     # Person carries a stricter floor than location: higher volume and no
     # structural validator, so a higher floor curbs false positives.
     assert PER_LABEL_FLOOR["location"] <= PER_LABEL_FLOOR["person"]
+
+
+def test_window_inference_cache(monkeypatch):
+    repeated_text = "This is going to repeat, don't waste resources more than once!"
+    cache_missed_text = "This appears for the first time, it will increase the model call count."
+
+    mock_predictions = {
+        repeated_text: [{"start": 0, "end": 1, "label": "person", "score": 1.0}],
+        cache_missed_text: [{"start": 0, "end": 1, "label": "person", "score": 1.0}],
+    }
+
+    mock_gliner = Mock()
+    mock_gliner.predict_entities.side_effect = lambda text, *_, **__: mock_predictions[text]
+    monkeypatch.setattr("detector.gliner_layer._model", mock_gliner)
+
+    expected = [
+        Span(
+            entity_type="PERSON",
+            start=0,
+            end=1,
+            score=1.0,
+        )
+    ]
+
+    assert detect_gliner(cache_missed_text, 0.0) == expected
+
+    for _ in range(5):
+        assert detect_gliner(repeated_text, 0.0) == expected
+
+    # Hits once for the original cache miss, and only once more for the repeated content
+    assert mock_gliner.predict_entities.call_count == 2