# GLiNER2 Integration — Cycle 5: Implementation Planning

**Status:** Historical implementation plan. Delivered in v1.2.0.
**Note:** Remote API tier (PIONEER_API_KEY) and fastcoref/spaCy coreference were removed. Extractor uses local GLiNER2 inference only. `--coref`/`--coref-device` flags were not shipped.
**Target:** extractor v1.2.0, schema v1.1.0
**Prerequisite:** extractor v1.1.0 (table reassembly, contextual prefix, async API) must ship first.

---

## 1. Version Assignment

| Component | Version |
|---|---|
| extractor library | v1.2.0 |
| Schema (element/chunk records) | v1.1.0 |
| SPEC | v1.1.0 (new §21) |

**Schema v1.1.0 is introduced in extractor v1.2.0.** The schema version decouples from the library version because schema consumers (vector stores, downstream apps) pin to `schema_version`, not `extractor_version`.

The `schema_version` field in all emitted records changes from `"1.0.0"` to `"1.1.0"` when running extractor v1.2.0. This is a MINOR bump — no breaking changes. All existing consumers that check `schema_version` and reject `MAJOR` mismatches will accept `"1.1.0"` payloads.

---

## 2. Phase Breakdown

### Phase 0: Prerequisites (v1.1.0 — not GLiNER2 work)

The following v1.1.0 items must be complete before Phase 1 begins:

- [ ] Table continuation reassembly (`table:continuation` merge before chunking)
- [ ] Contextual chunk prefix injection (deterministic, template-based)
- [ ] Async API (`aextract`, `aextract_batch`, `asyncio.Queue` + thread)

GLiNER2 integration must be designed to work cleanly with the async element stream from v1.1.0. `NERStep.annotate()` takes an `Iterable[BaseElement]`, which is compatible with both sync generators and async-to-sync adapters. Full async NER (`async def annotate()`) is deferred to v1.3.0.

---

### Phase 1: Schema Foundation (Week 1)

**Goal:** Define all Pydantic models and update schema_version. No inference code yet.

**Files modified:**
- `src/extractor/schema.py`
- `pyproject.toml` (version bump to 1.2.0-dev; `gliner>=2.0.0` and `torch>=2.0.0` are already in core dependencies)

**Tasks:**

1. Add `EntityAnnotation` Pydantic model to `schema.py`:
   - Fields: `entity_type`, `text`, `start`, `end`, `score`, `extraction_method`, `source_element_id`
   - `start`, `end`, `source_element_id` are `Optional` with `None` default
   - `model_config = ConfigDict(extra="forbid")`

2. Add `entities: list[EntityAnnotation] | None = None` to `BaseElement`:
   - Default `None` (absent in JSON when NER not run)
   - Field-level docstring: "Absent when NER not run. Empty list when NER ran, no entities found."

3. Add `entity_types: list[str] | None = None` to `ChunkMetadata`

4. Add `entities: list[EntityAnnotation] | None = None` to `ChunkRecord`

5. Update `schema_version` constant to `"1.1.0"` in `schema.py`

6. Add `[ner]` extras to `pyproject.toml`:
   ```toml
   [project.optional-dependencies]
   ner = [
       "gliner>=2.0.0",
       "torch>=2.0.0",
   ]
   ```

7. Update writer serialization: ensure `model.model_dump(exclude_none=True)` is used so `entities: None` is omitted from JSON output.

**Tests (Phase 1):**
- `test_schema_entity_annotation.py`: unit tests for `EntityAnnotation` model validation
- `test_schema_base_element_entities_absent.py`: verify `entities` field is absent (not null) when not set
- `test_schema_version_1_1_0.py`: verify emitted records carry `"schema_version": "1.1.0"`

---

### Phase 2: NER Module Core (Week 2)

**Goal:** Implement `src/extractor/ner.py` without CLI integration. Testable in isolation.

**Files created:**
- `src/extractor/ner.py`

**Files modified:**
- `src/extractor/__init__.py` (expose `preload_model` in public API)

**Tasks:**

1. Implement `NERStep` class per architecture spec (Cycle 4 §2):
   - `__init__()` with lazy `import gliner` guard
   - `annotate(elements)` as lazy generator (collects, batches, yields)
   - `_get_ner_text()`: element-type-aware text extraction (table cell text vs element.text)
   - `_truncate_text()`: char-count heuristic truncation at 256 cl100k equivalent
   - `_get_effective_threshold()`: per-type threshold adjustments
   - `_auto_batch_size()`: length-bucket-based batch sizing
   - `_run_inference_on_indices()`: batched GLiNER2 inference with failure isolation

2. Implement `NERStats` dataclass

3. Implement `preload_model()` module function

4. Implement module-level `_model_singleton` pattern

5. Implement `NER_SKIP_TYPES`, `NER_HIGH_THRESHOLD_TYPES`, `NER_LOW_THRESHOLD_TYPES` constants

6. Expose in `src/extractor/__init__.py`:
   ```python
   from extractor.ner import preload_model as preload_ner_model  # guarded import
   ```
   Guard with try/except ImportError — if `[ner]` not installed, `preload_ner_model` is not available but the package still imports cleanly.

**Tests (Phase 2):**

All NER tests must gate on `gliner` availability:
```python
gliner = pytest.importorskip("gliner")
```

- `test_ner_skip_types.py`: verify elements with NER_SKIP_TYPES pass through with `entities=None`
- `test_ner_threshold.py`: verify spans below threshold are dropped
- `test_ner_table_cell_text.py`: verify table elements use structured row text, not markdown
- `test_ner_truncation.py`: verify long elements are truncated and `stats.truncated_elements` incremented
- `test_ner_batch_failure_isolation.py`: mock GLiNER2 to raise exception; verify `entities=[]` on affected batch, extraction continues
- `test_ner_singleton.py`: verify second `NERStep` reuses cached model
- `test_ner_preload.py`: verify `preload_model()` initializes singleton

Integration test (requires gliner installed):
- `test_ner_end_to_end.py`: run `NERStep` on a small synthetic element stream; verify output shape

---

### Phase 3: Chunker Entity Rollup (Week 2, parallel with Phase 2)

**Goal:** When chunking, aggregate element-level entities into chunk-level rollup.

**Files modified:**
- `src/extractor/chunker.py`

**Tasks:**

1. Add `_build_chunk_entities(contributing_elements)` function (Cycle 4 §4)
2. Add `_build_chunk_entity_types(entities)` function
3. Call these functions in the chunk flush path, after the `ChunkRecord` is constructed
4. Set `chunk.entities` and `chunk.chunk_metadata.entity_types` accordingly

**Deduplication logic:**
- Key: `(entity_type, text.lower())`
- On collision: keep the `EntityAnnotation` with the highest `score`
- `source_element_id`: set to the element ID of the winning annotation

**Tests (Phase 3):**

- `test_chunker_entity_rollup_absent.py`: when NER not run, `chunk.entities` is `None` (absent)
- `test_chunker_entity_rollup_empty.py`: when NER ran but no entities on any element, `chunk.entities = []`
- `test_chunker_entity_rollup_dedup.py`: identical (type, text) from two elements → single entry with higher score
- `test_chunker_entity_types_summary.py`: `entity_types` is deduplicated sorted list of all entity_type values

---

### Phase 4: CLI Integration (Week 3)

**Goal:** Wire `--entities` and related flags into the `run` command.

**Files modified:**
- `src/extractor/cli.py`

**Tasks:**

1. Add four new parameters to `run()` command:
   - `--entities/--no-entities` (bool flag, **default True** — on by default)
   - `--entities-model` (str, default `"fastino/gliner2-multi-v1"`)
   - `--entities-types` (str | None, comma-separated)
   - `--entities-threshold` (float, default 0.50, range [0.0, 1.0])

2. Instantiate `NERStep` if `entities=True`:
   - Handle `ImportError` → **graceful fallback**: print one-line warning to stderr, emit `W_NER_MODEL_UNAVAILABLE` in manifest, continue without NER. **Never exit code 6 for missing gliner2.**
   - Call `preload_model()` before the document loop (eager load for batch use)

3. Wrap element stream with `ner_step.annotate(element_stream)` before passing to writer

4. Pass `ner_step.manifest_stats` to writer for manifest inclusion

5. Record `--entities`, `--entities-model`, `--entities-types`, `--entities-threshold` in `run_config` block of manifest

**Tests (Phase 4):**

- `test_cli_entities_default_on.py`: `extractor run doc.pdf` with no flags — verify NER attempted by default
- `test_cli_entities_no_entities.py`: `extractor run doc.pdf --no-entities` — verify entities field absent
- `test_cli_entities_fallback.py`: verify graceful fallback (W_NER_MODEL_UNAVAILABLE warning, no crash) when gliner not installed and `--entities` (default) is active
- `test_cli_entities_threshold_validation.py`: verify `--entities-threshold 1.5` raises error
- `test_cli_run_config_entities.py`: verify manifest `run_config` records entity flags

---

### Phase 5: Manifest and Writer Integration (Week 3)

**Goal:** Populate `stats.entities` and `parser_chain` NER step in manifest.

**Files modified:**
- `src/extractor/writer.py`

**Tasks:**

1. Accept optional `ner_stats: dict | None` in writer's manifest-building function
2. If `ner_stats` is not None, add `stats.entities = ner_stats` to manifest
3. Add NER step to `parser_chain` when NER was run:
   ```json
   {"step": N, "name": "ner_annotator", "impl": "gliner:MODEL_ID", "duration_ms": ...}
   ```
4. Add new warning codes to manifest warning emission:
   - `W_NER_MODEL_UNAVAILABLE`
   - `W_NER_INFERENCE_FAILED`
   - `W_NER_INPUT_TRUNCATED`
5. Add to `run_config`: `extract_entities`, `entities_model`, `entities_labels`, `entities_threshold`

**Tests (Phase 5):**

- `test_manifest_entities_stats.py`: verify `stats.entities` block structure when NER runs
- `test_manifest_entities_absent.py`: verify `stats.entities` is absent when `--entities` not used
- `test_manifest_parser_chain_ner.py`: verify NER step appears in `parser_chain`
- `test_manifest_ner_warnings.py`: verify `W_NER_INFERENCE_FAILED` appears when batch fails

---

### Phase 6: SPEC Update (Week 3–4)

**Goal:** Add §21 Entity Extraction to `docs/SPEC.md`.

**Files modified:**
- `docs/SPEC.md`

**Changes:**

1. Update header: `## Version 1.1.0`
2. Add `21` to Table of Contents: `21. [Entity Extraction](#21-entity-extraction)`
3. Update §2 Versioning Contract: note 1.0.0 → 1.1.0 bump for `entities` field addition
4. Update §5 Base Element Schema: add `entities` row to Optional Fields table
5. Update §8 Chunk Schema: add `entities` and `chunk_metadata.entity_types` fields, update example JSON
6. Update §11 Manifest Schema: add `run_config.extract_entities`, `run_config.entities_*`, `stats.entities` block, new warning codes
7. Update §13 Stability Contract: add `entities` to Experimental Fields list
14. Update §14 CLI Contract: add `--entities`, `--entities-model`, `--entities-types`, `--entities-threshold` to OPTIONS table
8. Add new §21 (full content below)

**New §21 Entity Extraction content outline:**

```
§21. Entity Extraction
  §21.1 Overview
  §21.2 Model: GLiNER2
  §21.3 Canonical Label Registry (Tier 1, Tier 2, Tier 3 tables)
  §21.4 EntityAnnotation Schema
    - Element-level annotation (with start/end)
    - Chunk-level rollup annotation (with source_element_id, no start/end)
  §21.5 Element Type Eligibility
    - NER-eligible types (table with explicit note on cell text)
    - NER-ineligible types (with reason column)
  §21.6 Confidence Threshold Policy
    - Global threshold
    - Per-type overrides (high threshold for structural:section_header; low for table:*)
    - Hard drop below 0.40
  §21.7 Absent vs Empty Semantics
    - entities absent ↔ NER not run
    - entities [] ↔ NER ran, no entities found above threshold
    - entities null ↔ NEVER (protocol violation)
  §21.8 BERT Token Limit Handling
    - 256 cl100k token truncation
    - W_NER_INPUT_TRUNCATED warning
  §21.9 Chunk-Level Entity Rollup
    - Deduplication algorithm
    - entity_types summary field
  §21.10 Manifest Integration
    - run_config fields
    - stats.entities block
    - parser_chain NER step
    - Warning codes
  §21.11 Installation
    - pip install extractor[ner]
    - Memory requirements
  §21.12 CLI Examples
```

---

### Phase 7: Documentation and Release (Week 4)

**Goal:** Complete all user-facing documentation, run full test suite, cut release.

**Tasks:**

1. Update `README.md`:
   - Add "Entity Extraction" section under Features
   - Add installation instructions for `[ner]` extras
   - Add usage example with `--entities`

2. Update `CHANGELOG.md`:
   - `v1.2.0` section with entity extraction feature
   - Schema version bump notice: 1.0.0 → 1.1.0

3. Run full test suite including NER tests (requires `pip install extractor[ner]` in CI):
   ```yaml
   - name: Test with NER extras
     run: |
       pip install -e ".[ner,dev]"
       pytest tests/ -m "not slow"
   ```

4. Verify CI passes without `[ner]` installed:
   ```yaml
   - name: Test without NER extras
     run: |
       pip install -e ".[dev]"
       pytest tests/ -k "not ner"
   ```

5. Bump version in `pyproject.toml`: `version = "1.2.0"`

6. Publish to PyPI via existing CI workflow

---

## 3. SPEC Sections to Add/Modify

| Section | Change Type | Description |
|---|---|---|
| Header | Modify | v1.0.0 → v1.1.0 |
| §2 Versioning Contract | Modify | Add row for entities field addition |
| §5 Base Element Schema | Modify | Add `entities` to Optional Fields table |
| §8 Chunk Schema | Modify | Add `entities`, `chunk_metadata.entity_types`; update example |
| §11 Manifest Schema | Modify | Add `run_config.extract_entities*`, `stats.entities`, 3 new warning codes |
| §13 Stability Contract | Modify | Add `entities`, `entity_types`, `EntityAnnotation.*` to Experimental list |
| §14 CLI Contract | Modify | Add 4 new flags to OPTIONS table; add examples |
| §21 Entity Extraction | **NEW** | Full entity extraction specification (12 subsections) |

---

## 4. Schema Version Changes

### What Changes in schema v1.1.0

All changes are additive optional fields. No existing field is renamed, removed, or made required.

**New optional fields on element records:**
```
entities: list[EntityAnnotation] | absent
```

**New optional fields on chunk records:**
```
entities: list[EntityAnnotation] | absent
chunk_metadata.entity_types: list[str] | absent
```

**New EntityAnnotation object schema:**
```
entity_type: str (required)
text: str (required)
start: int | absent (element-level only)
end: int | absent (element-level only)
score: float (required)
extraction_method: str (required)
source_element_id: str | absent (chunk-level only)
```

**New manifest fields:**
```
run_config.extract_entities: bool
run_config.entities_model: str
run_config.entities_labels: list[str]
run_config.entities_threshold: float
stats.entities: NERStats object
parser_chain[N]: NER annotator step
```

**Backward compatibility:** Consumers that do not check for `entities` continue to work unchanged. Consumers that check `schema_version` and process `1.x.x` payloads will accept `1.1.0`. No existing field is affected.

---

## 5. Testing Requirements

### Unit Tests (no gliner dependency)

| Test File | What It Tests |
|---|---|
| `tests/test_schema_entity_annotation.py` | `EntityAnnotation` model validation, forbidden extras |
| `tests/test_schema_entities_absent.py` | `entities: None` serializes as field omission, not null |
| `tests/test_schema_version.py` | All records carry `"schema_version": "1.1.0"` |
| `tests/test_chunker_entity_rollup.py` | Rollup: absent/empty/dedup cases |
| `tests/test_manifest_entities.py` | Manifest `stats.entities` structure |
| `tests/test_cli_entities_flags.py` | CLI flag parsing; exit code 6 on missing extras |

### Integration Tests (require `pip install extractor[ner]`)

Gated with `@pytest.mark.ner` marker + `pytest.importorskip("gliner")`:

| Test File | What It Tests |
|---|---|
| `tests/test_ner_step.py` | `NERStep.annotate()` end-to-end on synthetic elements |
| `tests/test_ner_skip_types.py` | All NER_SKIP_TYPES pass through with entities=None |
| `tests/test_ner_threshold.py` | Spans below threshold dropped; correct spans kept |
| `tests/test_ner_table_cells.py` | Table elements use `structured.rows` cell text |
| `tests/test_ner_truncation.py` | 1500-char element → truncated; `stats.truncated_elements += 1` |
| `tests/test_ner_batch_isolation.py` | Batch exception → entities=[] + stats.inference_errors += 1 |
| `tests/test_ner_singleton.py` | Two `NERStep` instances share the same model object |

### Invariant Tests

New testable invariants added to the existing invariant test suite:

1. If `run_config.extract_entities = false` → no element or chunk record in the JSONL has an `entities` field
2. If `run_config.extract_entities = true` → every eligible element (not in NER_SKIP_TYPES) has `entities` as a list (empty or non-empty, never absent)
3. `chunk.entities` is absent if and only if `all(e.entities is None for e in contributing_elements)`
4. Every entity in `chunk.entities` has `source_element_id` set and that ID appears in `chunk_metadata.element_ids`
5. `chunk_metadata.entity_types` is the sorted deduplicated list of `e.entity_type for e in chunk.entities` (or absent if entities absent)
6. No entity `score` is below the configured `run_config.entities_threshold` (except per-type threshold adjustments)

---

## 6. Optional Extras Packaging

### `pyproject.toml` additions

```toml
[project.optional-dependencies]
ner = [
    "gliner>=2.0.0,<3.0.0",
    "torch>=2.0.0",
]
audio = [
    "whisperx",
]
ocr = [
    "pytesseract",
    "easyocr",
]
all = [
    "extractor[ner,audio,ocr]",
]
```

### Installation Matrix

| Command | What's installed | NER available |
|---|---|---|
| `pip install extractor` | Core only | No |
| `pip install extractor[ner]` | Core + gliner + torch | Yes |
| `pip install extractor[all]` | Core + all extras | Yes |

### Dependency Notes

- `torch>=2.0.0` is a large dependency (~2.5 GB). Document this prominently in installation instructions.
- `gliner>=2.0.0` pins to major version 2 to avoid API breaks from major upgrades.
- Model weights (~550 MB at fp16) are downloaded on first use from HuggingFace Hub. Document this. For air-gapped environments, use `--entities-model /path/to/local/model`.

---

## 7. Documentation Requirements

### Required for Release

1. **SPEC §21** — complete entity extraction specification (see Phase 6)
2. **README entity extraction section** — installation, basic usage, output example
3. **`--entities` help text** — memory requirements, latency estimate, model info
4. **CHANGELOG v1.2.0 entry** — feature description, schema version note

### Recommended (not blocking release)

5. **Entity extraction cookbook** (separate doc): advanced patterns — vector store integration, GraphRAG routing, cross-document entity indexes
6. **Performance tuning guide**: fp16 vs fp32, batch size tuning, `preload_model()` for batch pipelines, `torch.compile` for dedicated batch servers

---

## 8. Release Checklist

- [ ] Phase 1: Schema models and [ner] extras group merged and green in CI
- [ ] Phase 2: `src/extractor/ner.py` complete, unit + integration tests passing
- [ ] Phase 3: Chunker entity rollup complete, tests passing
- [ ] Phase 4: CLI `--entities` flags wired and tested
- [ ] Phase 5: Manifest integration complete, manifest tests passing
- [ ] Phase 6: SPEC §21 written and reviewed
- [ ] CI passes WITHOUT `[ner]` installed (core tests only)
- [ ] CI passes WITH `[ner]` installed (all tests including NER integration)
- [ ] All invariant tests passing
- [ ] `pyproject.toml` version = `"1.2.0"`
- [ ] `schema_version` constant = `"1.1.0"` in schema.py
- [ ] CHANGELOG.md updated
- [ ] PyPI publish via CI workflow

---

## 8b. Phase 8: Analytics Module (v1.2.0, parallel with Phase 6)

**Goal:** Ship `extractor.analytics` — cross-document entity index, canonicalization, and citation layer.

**New package:** `src/extractor/analytics/`

```
src/extractor/analytics/
    __init__.py
    entity_index.py      # EntityIndex, EntityMention, CanonicalEntity — JSON-backed reverse index
    canonicalizer.py     # normalize_surface, merge_index_aliases, find_similar_entities
    corpus_analytics.py  # entity_frequency_table, entity_doc_cooccurrence, top_entities_by_section
    citations.py         # EntityCitation, cite_entity, explain_answer_entities
```

**Key design decisions:**
- `EntityIndex` is read-only after build, serializes to JSON (no SQLite)
- `Canonicalizer` is 3-layer: rule-based normalization → rapidfuzz clustering → optional char n-gram TF-IDF cosine similarity (no LLM, no network)
- `cite_entity()` returns `EntityCitation` objects with human-readable provenance strings using `source_element_id`, `source_filename`, `page`, and `section_path`
- `explain_answer_entities()` takes retrieved chunk dicts and surfaces citations for every entity — RAG explainability layer

**Dependencies added:** `rapidfuzz>=3.0.0` (already in `[ner]` extras)

**Optional dependency for corpus analytics:** `sentence-transformers>=2.2` in new `[canon-embeddings]` extras group.

**Tests:**
- `test_entity_index_build.py`: build from synthetic JSONL, verify mention counts and alias sets
- `test_entity_index_persistence.py`: save → load round-trip fidelity
- `test_canonicalizer_rules.py`: verify tier-1 normalization (suffixes, tickers, titles)
- `test_canonicalizer_fuzzy.py`: verify rapidfuzz clustering thresholds
- `test_citations_format.py`: verify citation string format with page/section/element data
- `test_explain_answer_entities.py`: end-to-end with mock retrieved chunks

---

## 8c. Phase 9: Coreference (v1.3.0 — deferred)

Intra-document coreference (fastcoref + LingMessCoref) and cross-document canonicalization via `EntityRegistry` are designed and documented in `docs/gliner2-coreference.md` but targeted for **v1.3.0**, not v1.2.0. Rationale:

- fastcoref + spaCy add ~400MB to the install footprint
- Requires full document in RAM before chunking (changes pipeline streaming model)
- Coreference quality is harder to test/validate than NER quality
- The v1.2.0 entity metadata (canonical_id-free) is already high-value for vector store filtering

**v1.3.0 additions:**
```toml
[project.optional-dependencies]
coref = [
    "fastcoref>=2.1.0",
    "spacy>=3.7",
]
canon-embeddings = [
    "sentence-transformers>=2.2",
    "numpy>=1.24",
]
```

New CLI flags in v1.3.0: `--coref`, `--canonicalize`, `--registry-path`, `--use-embeddings`, `--coref-device`

---

## 9. Risk Register

| Risk | Likelihood | Impact | Mitigation |
|---|---|---|---|
| GLiNER2 API changes between now and implementation | Low | Medium | Pin `gliner>=2.0.0,<3.0.0`; audit API at implementation start |
| torch version conflict with other user dependencies | Medium | Low | `torch>=2.0.0` is a wide range; document conflict resolution |
| BERT token limit causes silent entity misses | High (expected) | Low | 256 cl100k truncation + W_NER_INPUT_TRUNCATED warning; well-documented |
| Singleton causes model version mismatch in multi-model scenarios | Low | Medium | Singleton keyed on model_id; different model_ids create separate instances |
| Performance worse than projected on large docs | Low | Low | Feature is opt-in; latency documented; `preload_model()` available |
| False positive entities confuse RAG consumers | Medium | Low | Threshold documented; consumers should filter by `score >= 0.65` for index-ready entities |
