Improve GLiNER detector caching (#119)

diff --git a/detector/detector/gliner_layer.py b/detector/detector/gliner_layer.py
index e9290c521a3e7ea30d2e421ce97ab76353d78b3a..be58f10107833d29f811b760c736692e04ce48a0 100644 (file)
--- 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 1061158a9c1f00468213a743064721d5de984b3b..55a255a5df3c9ffda5797241f0c5ce7894638ea1 100644 (file)
--- 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