Extraction: GLiNER windows over long chunks, per-model thresholds, 'Unknown' band - #320
Open
Naseem77 wants to merge 5 commits into
Open
Conversation
GLiNER confidence scores are on a different scale per model: the bi-encoder models return almost nothing at the 0.75 that suits gliner_medium-v2.1 (measured at 2 entities for an entire corpus, with no error raised). A single hard-coded 0.75 therefore silently broke any model swap. DEFAULT_THRESHOLDS records the measured operating point for each known model (gliner_medium-v2.1 / large-v2.1 / medium-v2.5: 0.75; gliner-bi-small / bi-base-v2.0: 0.50); threshold=None (the default) looks the model up and falls back to 0.5 with a warning for unknown models. The default model stays urchade/gliner_medium-v2.1: gliner-bi-small-v2.0 measured better in August (ceiling recall 0.805 vs 0.709, half the memory, half the time) but its score calibration collapses to <= 0.03 under gliner 0.2.27 / transformers 5.6 and it returns zero entities; the table keeps its tuned threshold for when the bi-encoder loading path is stable again.
… silently GLiNER has a hard input limit (config.max_len) and discards anything beyond it silently - no exception, no warning, the tail simply never reaches the model. Measured on gliner_medium-v2.1 (max_len 384): a probe entity placed at word-token 388 is always returned, at 389 never. The new default model raises the limit to 2048, but real documents still exceed it: disabling windowing on the benchmark corpus dropped entity recall from 0.805 to 0.621. Windowing matters regardless of the model. Text longer than window_tokens is now processed as a series of overlapping windows and the predictions merged. Short text takes a fast path and behaves exactly as before, so nothing changes for chunk-sized input. Two details that matter for correctness: - Predictions are re-offset into the original text, so character spans stay valid for downstream span-based code. - window_overlap defaults to 48 word-tokens, comfortably above the model's max_width (longest representable entity, 12 words), or an entity sitting on a window boundary would be lost by both neighbours. _merge collapses duplicate spans to the highest-scoring copy and drops any span strictly contained in a longer span of the same label - that is the clipped remains of an entity the neighbouring window saw whole.
…ntities Previously the model was queried at `threshold`, so anything less confident was discarded inside GLiNER and never reached the SDK. Setting candidate_threshold below threshold instead keeps those entities and labels them "Unknown". The rest of the pipeline already supports this: ontology filtering explicitly whitelists "Unknown" so low-confidence nodes survive pruning, and entity resolution prefers any specific type over "Unknown" when merging duplicates. Off by default, and measured that way deliberately. Lowering threshold outright raised entity recall 32% but dropped entity F1 0.568 -> 0.474 and triple F1 0.236 -> 0.211: low-confidence predictions are mostly noise, and should be marked rather than trusted. Keeping them behind an explicit opt-in means the default path is unchanged. A candidate_threshold above threshold is rejected at construction rather than silently discarding the entities it is meant to preserve.
…old as 'Unknown' by default candidate_threshold now defaults to threshold * (1 - CANDIDATE_BAND) with CANDIDATE_BAND = 0.25, so the band follows whichever model threshold is in effect (0.5625-0.75 for gliner_medium-v2.1, 0.375-0.5 for the bi-encoders) instead of assuming one score scale. Spans in the band reach the step-2 LLM labelled 'Unknown' and are re-typed or dropped against the text; spans below the band are still discarded inside GLiNER. candidate_threshold=None turns the band off; an explicit number sets the floor. Measured end to end on the 11-document benchmark, same code with the band off vs on: NER spans 1828 -> 2468, 0 'Unknown' nodes reached the graph (the LLM typed every kept one), entity recall 0.535 -> 0.571 at precision 0.605 -> 0.572 (F1 0.568 -> 0.572), triple relaxed F1 0.228 -> 0.236, triple exact F1 0.125 -> 0.128, QA 30.4 % -> 32.4 %, ingest cost +8 %. A fixed low floor of 0.30 measured earlier was net negative (F1 -0.027), which is why the band is relative and narrow.
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Team Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
The new windowing code introduces a few correctness/clarity issues (docstring contradiction, missing overlap validation, and concurrency/performance concerns) that should be addressed before approval.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR improves the default GLiNER-based entity extraction path in the ingestion pipeline by (1) ensuring long chunks are fully processed (no silent truncation), and (2) making confidence handling model-aware while preserving borderline spans for downstream re-typing.
Changes:
- Add windowed GLiNER inference over long chunks with overlap and a merge step for deduplication.
- Introduce per-model default thresholds plus a configurable “Unknown” candidate band below the main threshold.
- Update tests and docs to cover/explain the candidate band behavior.
File summaries
| File | Description |
|---|---|
| graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/entity_extractors.py | Implements windowed inference, per-model thresholds, and the candidate “Unknown” band logic. |
| graphrag_sdk/tests/test_entity_extractors.py | Adds unit tests validating default/explicit thresholds and candidate band behavior. |
| docs/extraction.mdx | Updates public documentation to describe the candidate band and configuration. |
Review details
- Files reviewed: 3/3 changed files
- Comments generated: 5
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+616
to
+625
| def _word_spans(self, model: Any, text: str) -> list[tuple[str, int, int]]: | ||
| """Split text the same way GLiNER does, keeping char offsets.""" | ||
| if self._splitter is None: | ||
| splitter = getattr(getattr(model, "data_processor", None), "words_splitter", None) | ||
| if splitter is None: | ||
| from gliner.data_processing import WordsSplitter | ||
|
|
||
| splitter = WordsSplitter() | ||
| self._splitter = splitter | ||
| return list(self._splitter(text)) |
Comment on lines
+643
to
+656
| kept: list[dict[str, Any]] = [] | ||
| for p in best.values(): | ||
| contained = any( | ||
| q is not p | ||
| and q["label"] == p["label"] | ||
| and q["start"] <= p["start"] | ||
| and p["end"] <= q["end"] | ||
| and (q["end"] - q["start"]) > (p["end"] - p["start"]) | ||
| for q in best.values() | ||
| ) | ||
| if not contained: | ||
| kept.append(p) | ||
| kept.sort(key=lambda p: (p["start"], p["end"])) | ||
| return kept |
| if len(words) <= window: | ||
| return model.predict_entities(text, labels, threshold=floor) | ||
|
|
||
| step = max(1, window - self._window_overlap) |
| beyond it **silently** — no exception, no warning, the tail simply never | ||
| reaches the model. Measured on ``gliner_medium-v2.1`` (``max_len`` 384): a | ||
| probe entity placed at word-token 388 is always returned, at 389 never. | ||
| The default model's limit is 2048, but real documents still exceed it — |
Comment on lines
+561
to
+563
| self._window_tokens = window_tokens | ||
| self._window_overlap = window_overlap | ||
| self._splitter: Any = None |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Extraction — GLiNER reads the whole chunk, per-model thresholds, "Unknown" band
Part of FalkorDB/research#88. Stacked on #319.
1. GLiNER reads the whole chunk
Problem: GLiNER truncates input at
config.max_lensilently — forgliner_medium-v2.1a probe entity at word 388 is always found, at word 389 never. Everything past that in a chunk was invisible; disabling the fix on the benchmark drops ceiling recall 0.805 → 0.621.Fix: text longer than the window is processed as overlapping windows (overlap 48 word-tokens, above the model's 12-word
max_width) and merged. Short text takes the old fast path.2. Per-model confidence thresholds
Problem: one hard-coded 0.75 for every model. Scores are not comparable across models — the bi-encoders return 2 entities for a whole corpus at 0.75.
Fix:
DEFAULT_THRESHOLDS(medium / large / medium-v2.5: 0.75; bi-small / bi-base: 0.50);threshold=Nonelooks the model up, unknown models fall back to 0.5 with a warning. Default model staysurchade/gliner_medium-v2.1(bi-small measured better in August but returns 0 entities under gliner 0.2.27 / transformers 5.6).3. "Unknown" band 25 % below the threshold
Problem: everything under the cutoff was discarded inside GLiNER; borderline real names were lost.
Fix:
candidate_thresholddefaults tothreshold × 0.75(0.5625–0.75 for the default model). Spans in that band reach the step-2 LLM labelledUnknown, which re-types or drops them against the text; below the band is still discarded.candidate_threshold=Noneturns it off.Measured (same code, band off vs on): NER spans 1,828 → 2,468; 0
Unknownnodes reached the graph (the LLM typed every kept one); entity recall 0.535 → 0.571 at precision 0.605 → 0.572 (F1 0.568 → 0.572); triple relaxed F1 0.228 → 0.236; answer accuracy 30.4 → 32.4 %; +8 % ingest cost. A wide fixed floor of 0.30 tested earlier was net negative (F1 −0.027), which is why the band is relative and narrow.Verification
Full suite: 1172 passed, 41 skipped; ruff clean.