unstructured2graph: GLiNER2Backend uses extract_long(), not extract() (#336)

unstructured2graph/src/unstructured2graph/gliner2_backend.py

@@ -111,6 +111,8 @@ def __init__(
         model: Any | None = None,
         entity_confidence_threshold: float | None = None,
         relation_confidence_threshold: float | None = None,
+        chunk_size: int = 384,
+        chunk_overlap: int = 64,
     ) -> None:
         """
         Args:
@@ -130,7 +132,7 @@ def __init__(
                 (letters, digits, underscore, not starting with a digit).
             model: Pre-loaded GLiNER2 extractor (e.g. an AutoExtractor
                 instance), or a test fake exposing the same
-                create_schema()/extract() methods. Bypasses loading
+                create_schema()/extract_long() methods. Bypasses loading
                 `model_name` and importing `gliner2` entirely -- the only
                 supported way to use this class without the `gliner2`
                 package installed.
@@ -140,6 +142,13 @@ def __init__(
                 head/tail confidence falls below this (relations carry no
                 confidence of their own -- see _extract_sync). None (default)
                 keeps everything the model returns.
+            chunk_size: Word-window size `model.extract_long()` scans `text`
+                with (see _extract_sync for why extract() alone isn't used).
+                GLiNER2's own default.
+            chunk_overlap: Overlap, in words, between consecutive windows --
+                gives an entity/relation spanning a window boundary a chance
+                to fall fully inside at least one window. GLiNER2's own
+                default.
 
         Raises:
             ValueError: if `workspace` isn't a valid Cypher identifier.
@@ -177,6 +186,8 @@ def __init__(
         self._workspace = workspace
         self.entity_confidence_threshold = entity_confidence_threshold
         self.relation_confidence_threshold = relation_confidence_threshold
+        self._chunk_size = chunk_size
+        self._chunk_overlap = chunk_overlap
 
     @property
     def workspace_label(self) -> str:
@@ -188,16 +199,30 @@ def _extract_sync(self, text: str) -> tuple[list[ExtractedEntity], list[Extracte
         unlike LightRAG's network-bound LLM call) -- run via asyncio.to_thread
         from aingest_chunk() so it doesn't block the event loop.
 
-        One combined `model.extract()` call covers both entities and
-        relations in the same forward pass: `model.create_schema().entities(
+        One combined `model.extract_long()` call covers both entities and
+        relations in the same pass: `model.create_schema().entities(
         self._entity_schema)`, `.relations(self._relation_schema)` when
-        relations are configured, then `model.extract(text, schema,
+        relations are configured, then `model.extract_long(text, schema,
+        chunk_size=self._chunk_size, chunk_overlap=self._chunk_overlap,
         include_spans=True, include_confidence=True)`. Confirmed against
         gliner2==2.0.0 that this returns `{"entities": {...},
         "relation_extraction": {...}}` -- the same per-key shape
         extract_entities()/extract_relations() each return individually,
         just merged into one result.
 
+        `extract_long()`, not the plain single-pass `extract()`: GLiNER2 is
+        an encoder-only span/boundary classifier with a fixed effective
+        context, and a session's combined text (one document per session,
+        not per turn) routinely exceeds it. Past that point `extract()`
+        both slows down and silently drops most entities, with no error --
+        confirmed on a real 16k-char session: `extract()` took 16.0s and
+        found 18 entities, `extract_long()` (this method, chunk_size=384,
+        chunk_overlap=64 -- GLiNER2's own defaults) took 4.3s and found 249.
+        `extract_long()` windows `text` into overlapping word chunks and
+        merges results back into text-global coordinates itself, so this
+        method's own span-based entity/relation matching below is unaffected
+        either way (see #336).
+
         This used to be two independent calls (extract_entities() then
         extract_relations()). A relation's head/tail span only ever carries
         a text span, not a reference to a specific already-extracted entity,
@@ -228,7 +253,14 @@ def _extract_sync(self, text: str) -> tuple[list[ExtractedEntity], list[Extracte
         schema = self.model.create_schema().entities(self._entity_schema)
         if self._relation_schema:
             schema = schema.relations(self._relation_schema)
-        raw = self.model.extract(text, schema, include_spans=True, include_confidence=True)
+        raw = self.model.extract_long(
+            text,
+            schema,
+            chunk_size=self._chunk_size,
+            chunk_overlap=self._chunk_overlap,
+            include_spans=True,
+            include_confidence=True,
+        )
 
         entities: list[ExtractedEntity] = []
         for entity_type, spans in raw.get("entities", {}).items():

unstructured2graph/tests/test_gliner2_backend.py

@@ -41,9 +41,11 @@ class _FakeModel:
     """Records every call so tests can assert on what schema was passed, and
     returns whatever canned result the test configured -- merged into the
     single {"entities": ..., "relation_extraction": ...} shape the real
-    model.extract() returns from one combined call (gliner2_backend.py's
+    model.extract_long() returns from one combined call (gliner2_backend.py's
     _extract_sync no longer calls extract_entities()/extract_relations()
-    separately)."""
+    separately, and uses extract_long() rather than plain extract() -- see
+    #336: extract() silently undercounts entities on text longer than
+    GLiNER2's effective context)."""
 
     def __init__(self, entities_result=None, relations_result=None):
         self.result: dict = dict(entities_result or {"entities": {}})
@@ -53,7 +55,7 @@ def __init__(self, entities_result=None, relations_result=None):
     def create_schema(self):
         return _FakeSchema()
 
-    def extract(self, text, schema, **kwargs):
+    def extract_long(self, text, schema, **kwargs):
         self.extract_calls.append((text, schema))
         return self.result
 
@@ -66,6 +68,40 @@ def test_entity_schema_built_from_ontology():
     assert backend._relation_schema == {}
 
 
+def test_chunk_size_and_overlap_default_and_are_configurable():
+    default_backend = GLiNER2Backend(ontology=ENTITY_ONLY_ONTOLOGY, model=_FakeModel())
+    assert default_backend._chunk_size == 384
+    assert default_backend._chunk_overlap == 64
+
+    custom_backend = GLiNER2Backend(ontology=ENTITY_ONLY_ONTOLOGY, model=_FakeModel(), chunk_size=128, chunk_overlap=32)
+    assert custom_backend._chunk_size == 128
+    assert custom_backend._chunk_overlap == 32
+
+
+@pytest.mark.asyncio
+async def test_extract_long_is_called_with_the_configured_chunk_size_and_overlap():
+    """Regression test for #336: _extract_sync must call extract_long(),
+    not extract() -- extract() silently undercounts entities on text longer
+    than GLiNER2's effective context, with no error to catch the regression
+    otherwise."""
+    model = MagicMock()
+    model.create_schema.return_value = _FakeSchema()
+    model.extract_long.return_value = {"entities": {}}
+    backend = GLiNER2Backend(ontology=ENTITY_ONLY_ONTOLOGY, model=model, chunk_size=128, chunk_overlap=32)
+
+    with patch("unstructured2graph.gliner2_backend.create_nodes_from_list"):
+        await backend.aingest_chunk(MagicMock(), Chunk(text="Alice works here.", hash="h1"))
+
+    model.extract_long.assert_called_once()
+    model.extract.assert_not_called()
+    _text, _schema = model.extract_long.call_args.args
+    kwargs = model.extract_long.call_args.kwargs
+    assert kwargs["chunk_size"] == 128
+    assert kwargs["chunk_overlap"] == 32
+    assert kwargs["include_spans"] is True
+    assert kwargs["include_confidence"] is True
+
+
 def test_workspace_label_defaults_and_is_configurable():
     model = _FakeModel()
     assert GLiNER2Backend(ontology=ENTITY_ONLY_ONTOLOGY, model=model).workspace_label == "gliner2"
@@ -110,9 +146,9 @@ async def test_entity_only_ontology_never_requests_relations_schema():
 
 @pytest.mark.asyncio
 async def test_entity_and_relation_ontology_requests_combined_schema_in_one_call():
-    """The whole point of the joint pass: one model.extract() call carries
-    both schemas, not two separate extract_entities()/extract_relations()
-    calls."""
+    """The whole point of the joint pass: one model.extract_long() call
+    carries both schemas, not two separate extract_entities()/
+    extract_relations() calls."""
     model = _FakeModel()
     backend = GLiNER2Backend(ontology=ENTITY_AND_RELATION_ONTOLOGY, model=model)