# GLiNER2 Coreference Resolution — Design & Implementation

**Status:** SUPERSEDED — fastcoref/spaCy coreference was not implemented. Extractor uses GLiNER2-only entity extraction (local inference). Coreference cluster linking is not available.
**Approach (original design):** fastcoref (LingMessCoref) for intra-document coreference + 3-layer cross-document canonicalization

---

## 1. Problem Statement

The NER step produces raw entity spans per element. The same real-world entity appears under many surface forms across a document and across files:

| Variant type | Example |
|---|---|
| Full vs short name | "Tim Cook" → "Cook" |
| Title reference | "the CEO", "Apple's chief executive" |
| Pronouns | "he", "his", "him" |
| Abbreviation | "Apple Inc." → "Apple" → "AAPL" |
| Cross-file variant | "OpenAI" → "OpenAI Inc." → "OAI" |

Without resolution, entity-filtered retrieval (`WHERE person = "Tim Cook"`) misses mentions of "Cook" and "he" — defeating the purpose of NER.

---

## 2. Library Choice: fastcoref + LingMessCoref

**fastcoref** with the LingMessCoref model is the production choice for intra-document coreference:
- CPU-first (`device="cpu"` is a first-class parameter)
- 2x faster than AllenNLP alternatives
- SOTA accuracy on English documents
- spaCy 3.x compatible, but does not require spaCy NER (we supply entities via GLiNER2)
- Returns cluster spans directly — easy to map back to GLiNER2 entity offsets

GLiNER2 zero-shot labels (`"pronoun referring to a person"`) are used as a **supplementary signal** for pronouns that fastcoref misses, not as the primary coreference engine.

```bash
pip install fastcoref spacy
python -m spacy download en_core_web_sm
```

---

## 3. Pipeline Position

Coreference must run on the **full document element stream, before chunking**. Chunking breaks the span continuity that coreference models depend on.

```
parser.parse()
  → _normalize_stream()
  → _reassemble_table_continuations()
  → NERStep.annotate(elements)       # GLiNER2: adds entities[] with start/end offsets
  → CorefStep.resolve(elements)      # fastcoref: adds coref_cluster_id, canonical_text
  → CanonStep.canonicalize(elements) # 3-layer: assigns stable canonical_id
  → chunker                          # entities[] propagates to chunks
  → writer
```

CorefStep collects all text-bearing elements into a single document string, runs fastcoref on it, then maps cluster spans back to per-element entity annotations via an offset map.

---

## 4. EntityAnnotation Schema (Extended)

```python
from pydantic import BaseModel
from typing import Optional, Literal

class EntityAnnotation(BaseModel):
    entity_type: str
    text: str                              # surface form as extracted
    start: Optional[int] = None            # char offset within element.text
    end: Optional[int] = None
    score: float                           # GLiNER2 confidence [0,1]
    extraction_method: str = "gliner2"
    source_element_id: Optional[str] = None

    # Coreference fields (set by CorefStep)
    coref_cluster_id: Optional[str] = None    # intra-doc cluster: "coref_3_a1b2c3d4"
    is_pronoun: bool = False                  # True for "he", "she", "they", etc.
    is_alias: bool = False                    # True for "Cook", "the company", "AAPL"

    # Canonicalization fields (set by CanonStep)
    canonical_id: Optional[str] = None        # e.g. "person:tim-cook"
    canonical_text: Optional[str] = None      # e.g. "Tim Cook"
    merge_confidence: Optional[float] = None  # 1.0=exact, 0.92=embedding, 0.85=string
    merge_decision: Optional[Literal[
        "exact", "rule", "similarity", "embedding", "new", "review"
    ]] = None
```

---

## 5. CorefStep — fastcoref Integration

```python
# src/extractor/coref_step.py

import hashlib
import re
from typing import Any, Optional

PRONOUN_SET = frozenset({
    "he", "him", "his", "himself",
    "she", "her", "hers", "herself",
    "they", "them", "their", "theirs", "themselves",
    "it", "its", "itself",
})

DESCRIPTIVE_PATTERNS = re.compile(
    r"^(the|a|an)\s+(company|firm|ceo|chief executive|president|founder|"
    r"chairman|director|agency|bank|fund)\b",
    re.IGNORECASE,
)

# Element types eligible for coreference (same as NER-eligible types)
COREF_TEXT_TYPES = frozenset({
    "text:narrative", "text:paragraph", "text:caption", "text:footnote",
    "text:abstract", "text:note", "structural:section_header", "structural:title",
    "list:item", "list:item_ordered",
})


def _classify_mention(text: str) -> tuple[bool, bool]:
    """Returns (is_pronoun, is_alias)."""
    lower = text.strip().lower()
    if lower in PRONOUN_SET:
        return True, False
    if DESCRIPTIVE_PATTERNS.match(text):
        return False, True
    # Single capitalized token with no space (e.g. "Cook", "AAPL")
    if " " not in text and text[0].isupper() and len(text) > 1:
        return False, True
    return False, False


class CorefStep:
    """
    Intra-document coreference resolution using fastcoref LingMessCoref.

    Operates on the full list of BaseElement objects for one document.
    Mutates entities[] on each element to add:
      - coref_cluster_id, canonical_text, is_pronoun, is_alias
    """

    def __init__(self, device: str = "cpu", enabled: bool = True) -> None:
        self.enabled = enabled
        self._device = device
        self._model: Any = None

    def _load_model(self) -> None:
        if self._model is not None:
            return
        try:
            from fastcoref import LingMessCoref  # type: ignore[import-untyped]
            self._model = LingMessCoref(device=self._device)
        except ImportError as exc:
            raise ImportError(
                "fastcoref required for --coref. Install: pip install fastcoref spacy && "
                "python -m spacy download en_core_web_sm"
            ) from exc

    def resolve(self, elements: list[Any]) -> list[Any]:
        if not self.enabled:
            return elements

        # Build document string + offset map
        doc_parts: list[str] = []
        offset_map: list[tuple[int, int, int]] = []  # (doc_start, doc_end, elem_idx)
        cursor = 0

        for idx, el in enumerate(elements):
            if el.element_type not in COREF_TEXT_TYPES:
                continue
            text = getattr(el, "text", "") or ""
            if not text.strip():
                continue
            doc_parts.append(text)
            offset_map.append((cursor, cursor + len(text), idx))
            cursor += len(text) + 1
            doc_parts.append(" ")

        if not doc_parts:
            return elements

        doc_text = "".join(doc_parts)

        self._load_model()
        preds = self._model.predict(texts=[doc_text])
        clusters = preds[0].get_clusters(as_strings=False)

        # Determine canonical text per cluster (longest non-pronoun mention)
        cluster_canonical: dict[int, tuple[str, str]] = {}
        for cidx, cluster in enumerate(clusters):
            best_text = ""
            best_len = -1
            for (s, e) in cluster:
                span_text = doc_text[s:e].strip()
                is_pron, _ = _classify_mention(span_text)
                if not is_pron and len(span_text) > best_len:
                    best_text = span_text
                    best_len = len(span_text)
            if not best_text:
                # All pronouns — use longest
                for (s, e) in cluster:
                    span_text = doc_text[s:e].strip()
                    if len(span_text) > best_len:
                        best_text = span_text
                        best_len = len(span_text)
            intra_id = f"coref_{cidx}_{hashlib.sha256(best_text.encode()).hexdigest()[:8]}"
            cluster_canonical[cidx] = (best_text, intra_id)

        # Span-to-cluster lookup
        span_to_cluster: dict[tuple[int, int], int] = {
            span: cidx
            for cidx, cluster in enumerate(clusters)
            for span in cluster
        }

        # Map clusters back to per-element entity annotations
        for (doc_start, doc_end, elem_idx) in offset_map:
            el = elements[elem_idx]
            entities = getattr(el, "entities", None) or []
            for ent in entities:
                if ent.start is None or ent.end is None:
                    continue
                doc_ent_start = doc_start + ent.start
                doc_ent_end = doc_start + ent.end

                is_pron, is_alias = _classify_mention(ent.text)
                ent.is_pronoun = is_pron
                ent.is_alias = is_alias

                cidx = _find_cluster(doc_ent_start, doc_ent_end, span_to_cluster)
                if cidx is not None:
                    canonical_text, intra_id = cluster_canonical[cidx]
                    ent.coref_cluster_id = intra_id
                    ent.canonical_text = canonical_text
                    if ent.canonical_id is None:
                        ent.canonical_id = intra_id  # placeholder; CanonStep overwrites

        return elements


def _find_cluster(
    start: int, end: int,
    span_to_cluster: dict[tuple[int, int], int],
) -> Optional[int]:
    if (start, end) in span_to_cluster:
        return span_to_cluster[(start, end)]
    for (s, e), cidx in span_to_cluster.items():
        if s <= start < e or start <= s < end:
            return cidx
    return None
```

---

## 6. Pronoun Resolution Strategy

```
Pronoun encountered ("he", "the company", etc.)
   │
   ├─ In a fastcoref cluster?
   │    YES → Cluster contains non-pronoun mention?
   │           YES → canonical_text = best non-pronoun mention → resolved ✓
   │           NO  → canonical_text = None, coref_cluster_id set → emit as-is
   │    NO  → Same-element antecedent scan:
   │           find most recent compatible entity in same element
   │           (gender/animacy compatibility check)
   │           Found → canonical_text = antecedent.canonical_text → resolved ✓
   │           Not found → emit with is_pronoun=True, canonical_text=None
```

Unresolved pronouns are **never dropped** — they are emitted with `is_pronoun=True, canonical_text=None`. Downstream consumers filter on `is_pronoun=False OR canonical_text IS NOT NULL` for index-ready entities.

---

## 7. Three-Layer Cross-Document Canonicalization

### Layer 1: Rule-Based Normalization

```python
import re, unicodedata

ORG_SUFFIXES = re.compile(
    r"\b(Inc\.?|Ltd\.?|LLC\.?|Corp\.?|Co\.?|PLC\.?|GmbH\.?|"
    r"Incorporated|Limited|Corporation|Company|Holdings?|Group|"
    r"Partners?|Associates?|Ventures?|International|Global)\b\.?",
    re.IGNORECASE,
)
PERSON_TITLES = re.compile(
    r"^(Dr\.?|Mr\.?|Mrs\.?|Ms\.?|Prof\.?|Sir|Dame|"
    r"Gen\.?|Col\.?|Sen\.?|Rep\.?|Gov\.?)\s+",
    re.IGNORECASE,
)
TICKER_MAP: dict[str, str] = {
    "AAPL": "Apple", "MSFT": "Microsoft", "GOOGL": "Alphabet",
    "META": "Meta Platforms", "AMZN": "Amazon", "TSLA": "Tesla",
    "NVDA": "Nvidia", "JPM": "JPMorgan Chase", "BAC": "Bank of America",
}

def normalize_surface(text: str, entity_type: str) -> str:
    text = unicodedata.normalize("NFC", text)
    text = " ".join(text.split())
    if entity_type == "organization":
        if text.upper() in TICKER_MAP:
            return TICKER_MAP[text.upper()]
        text = ORG_SUFFIXES.sub("", text).strip()
    elif entity_type == "person":
        text = PERSON_TITLES.sub("", text).strip()
    return text
```

### Layer 2: rapidfuzz String Similarity Clustering

```python
from rapidfuzz import fuzz

AUTO_MERGE_THRESHOLD = 92   # token_set_ratio >= 92 → auto merge
REVIEW_THRESHOLD = 78       # 78-91 → flag for review
# < 78 → different entities

def similarity_score(a: str, b: str) -> float:
    return max(
        fuzz.token_set_ratio(a, b),
        fuzz.token_sort_ratio(a, b),
        fuzz.partial_ratio(a, b),
    )
```

### Layer 3: Embedding Similarity (Optional, No LLM)

For "review" decisions from Layer 2, `all-MiniLM-L6-v2` (22MB, CPU) resolves ambiguous cases:

```python
# Only loaded when use_embeddings=True
EMBEDDING_MERGE_THRESHOLD = 0.92

def embedding_similarity(text_a: str, text_b: str) -> float:
    from sentence_transformers import SentenceTransformer
    model = SentenceTransformer("all-MiniLM-L6-v2", device="cpu")
    embs = model.encode([text_a, text_b], normalize_embeddings=True)
    return float(embs[0] @ embs[1])
```

---

## 8. canonical_id Derivation

```python
import re, hashlib, unicodedata

def make_slug(text: str) -> str:
    text = unicodedata.normalize("NFC", text).lower()
    text = re.sub(r"[^\w\s-]", "", text)
    text = re.sub(r"[\s_]+", "-", text).strip("-")
    return text

def derive_canonical_id(entity_type: str, canonical_text: str) -> str:
    slug = make_slug(normalize_surface(canonical_text, entity_type))
    return f"{entity_type.lower()}:{slug}"

# Examples:
# "Tim Cook", person    → "person:tim-cook"
# "Apple Inc.", org     → "organization:apple"
# "AAPL", organization  → "organization:apple"  (via TICKER_MAP)
# "John Smith" (2nd)    → "person:john-smith-a3f2" (hash suffix via registry)
```

IDs are **deterministic** — same input always produces the same ID.

---

## 9. EntityRegistry — JSON-Backed Cross-Session Persistence

```python
import json, threading, hashlib
from dataclasses import dataclass, field, asdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

@dataclass
class RegistryEntry:
    canonical_id: str
    canonical_text: str
    entity_type: str
    aliases: list[str] = field(default_factory=list)
    first_seen_doc: str = ""
    doc_count: int = 1
    mention_count: int = 1
    chunk_ids: list[str] = field(default_factory=list)
    pending_review: list[dict] = field(default_factory=list)
    merge_log: list[dict] = field(default_factory=list)

class EntityRegistry:
    """
    JSON-backed canonical entity store.
    No SQLite. Thread-safe reads; write lock for updates.
    """

    def __init__(self) -> None:
        self._entries: dict[str, RegistryEntry] = {}
        self._alias_index: dict[str, str] = {}  # "entity_type:normalized" → canonical_id
        self._lock = threading.Lock()

    @classmethod
    def load(cls, path: Path) -> "EntityRegistry":
        reg = cls()
        if path.exists():
            raw = json.loads(path.read_text(encoding="utf-8"))
            for k, v in raw.items():
                if k.startswith("_"):
                    continue
                entry = RegistryEntry(**{
                    fk: v[fk] for fk in RegistryEntry.__dataclass_fields__ if fk in v
                })
                reg._entries[k] = entry
                reg._index_entry(entry)
        return reg

    def save(self, path: Path) -> None:
        with self._lock:
            out = {
                "_meta": {
                    "schema_version": "1.0.0",
                    "updated_at": datetime.now(timezone.utc).isoformat(),
                    "total_entities": len(self._entries),
                },
                **{k: asdict(v) for k, v in self._entries.items()},
            }
            tmp = path.with_suffix(".tmp")
            tmp.write_text(json.dumps(out, indent=2, ensure_ascii=False), encoding="utf-8")
            tmp.replace(path)

    def resolve(
        self,
        surface: str,
        entity_type: str,
        doc_id: str = "",
        chunk_id: Optional[str] = None,
        use_embeddings: bool = False,
    ) -> str:
        """
        Look up or create a canonical_id for the given surface form.
        Updates alias lists and mention counts. Returns canonical_id.
        """
        from .canon.normalizer import normalize_surface
        from .canon.clusterer import similarity_score, AUTO_MERGE_THRESHOLD, REVIEW_THRESHOLD

        with self._lock:
            normalized = normalize_surface(surface, entity_type)
            alias_key = f"{entity_type}:{normalized.lower()}"

            # Exact alias match
            if alias_key in self._alias_index:
                cid = self._alias_index[alias_key]
                self._record_mention(cid, surface, doc_id, chunk_id)
                return cid

            # String similarity search
            best_score = 0.0
            best_cid = None
            for cid, entry in self._entries.items():
                if entry.entity_type != entity_type:
                    continue
                score = max(
                    similarity_score(normalized, normalize_surface(a, entity_type))
                    for a in [entry.canonical_text] + entry.aliases
                )
                if score > best_score:
                    best_score = score
                    best_cid = cid

            if best_score >= AUTO_MERGE_THRESHOLD and best_cid:
                self._add_alias(best_cid, surface, normalized, entity_type, "similarity", best_score / 100)
                self._record_mention(best_cid, surface, doc_id, chunk_id)
                return best_cid

            if best_score >= REVIEW_THRESHOLD and best_cid:
                if use_embeddings:
                    from .canon.embedder import embedding_similarity, EMBEDDING_MERGE_THRESHOLD
                    sim = embedding_similarity(normalized, self._entries[best_cid].canonical_text)
                    if sim >= EMBEDDING_MERGE_THRESHOLD:
                        self._add_alias(best_cid, surface, normalized, entity_type, "embedding", sim)
                        self._record_mention(best_cid, surface, doc_id, chunk_id)
                        return best_cid

            # New entity
            return self._create_entry(surface, normalized, entity_type, doc_id, chunk_id)

    def _create_entry(self, surface: str, normalized: str, entity_type: str, doc_id: str, chunk_id: Optional[str]) -> str:
        from .canon.ids import derive_canonical_id, derive_canonical_id_with_disambiguation
        base_id = derive_canonical_id(entity_type, normalized)
        if base_id in self._entries:
            existing = self._entries[base_id].canonical_text
            if existing.lower() != normalized.lower():
                base_id = derive_canonical_id_with_disambiguation(entity_type, normalized, existing)

        entry = RegistryEntry(
            canonical_id=base_id,
            canonical_text=normalized,
            entity_type=entity_type,
            aliases=[surface] if surface != normalized else [],
            first_seen_doc=doc_id,
            chunk_ids=[chunk_id] if chunk_id else [],
        )
        self._entries[base_id] = entry
        self._index_entry(entry)
        return base_id

    def _add_alias(self, cid: str, surface: str, normalized: str, entity_type: str, decision: str, confidence: float) -> None:
        entry = self._entries[cid]
        if surface not in entry.aliases:
            entry.aliases.append(surface)
            entry.merge_log.append({"merged_from": surface, "decision": decision, "confidence": round(confidence, 4)})
        self._alias_index[f"{entity_type}:{normalized.lower()}"] = cid
        self._alias_index[f"{entity_type}:{surface.lower()}"] = cid

    def _record_mention(self, cid: str, surface: str, doc_id: str, chunk_id: Optional[str]) -> None:
        entry = self._entries[cid]
        entry.mention_count += 1
        if chunk_id and chunk_id not in entry.chunk_ids:
            entry.chunk_ids = (entry.chunk_ids + [chunk_id])[:1000]

    def _index_entry(self, entry: RegistryEntry) -> None:
        key = f"{entry.entity_type}:{entry.canonical_text.lower()}"
        self._alias_index[key] = entry.canonical_id
        for alias in entry.aliases:
            self._alias_index[f"{entry.entity_type}:{alias.lower()}"] = entry.canonical_id
```

---

## 10. CanonStep — Pipeline Integration

```python
# src/extractor/canon/canon_step.py

from pathlib import Path
from typing import Any, Optional

from .registry import EntityRegistry


class CanonStep:
    def __init__(self, registry_path: Path, enabled: bool = True, use_embeddings: bool = False) -> None:
        self.enabled = enabled
        self.registry_path = registry_path
        self.use_embeddings = use_embeddings
        self._registry: Optional[EntityRegistry] = None

    def _get_registry(self) -> EntityRegistry:
        if self._registry is None:
            self._registry = EntityRegistry.load(self.registry_path)
        return self._registry

    def canonicalize(self, elements: list[Any], doc_id: str, doc_sha256: str) -> list[Any]:
        if not self.enabled:
            return elements
        registry = self._get_registry()
        for el in elements:
            for ent in (getattr(el, "entities", None) or []):
                if ent.is_pronoun and not ent.canonical_text:
                    continue
                surface = ent.canonical_text or ent.text
                cid = registry.resolve(
                    surface=surface,
                    entity_type=ent.entity_type,
                    doc_id=doc_id,
                    chunk_id=getattr(el, "id", None),
                    use_embeddings=self.use_embeddings,
                )
                ent.canonical_id = cid
                if not ent.canonical_text:
                    ent.canonical_text = registry._entries[cid].canonical_text
        return elements

    def flush(self) -> None:
        if self._registry is not None:
            self._registry.save(self.registry_path)
```

---

## 11. CLI Flags

| Flag | Default | Description |
|---|---|---|
| `--entities / --no-entities` | on | Enable/disable NER |
| `--coref / --no-coref` | off | Enable intra-document coreference (requires `--entities`) |
| `--canonicalize / --no-canonicalize` | off | Cross-doc canonicalization (requires `--entities`) |
| `--registry-path PATH` | none | JSON file path for EntityRegistry (required with `--canonicalize`) |
| `--use-embeddings` | off | Use sentence-transformers for ambiguous entity disambiguation |
| `--coref-device TEXT` | `"cpu"` | Device for fastcoref: `"cpu"` or `"cuda"` |

---

## 12. Output Examples

**Element-level (resolved pronoun):**
```json
{
  "id": "el_7f3a",
  "element_type": "paragraph",
  "text": "Tim Cook joined Apple in 1998. He became CEO in 2011.",
  "entities": [
    {
      "entity_type": "person", "text": "Tim Cook", "score": 0.94,
      "start": 0, "end": 8,
      "canonical_id": "person:tim-cook", "canonical_text": "Tim Cook",
      "coref_cluster_id": "coref_0_a1b2c3d4",
      "is_pronoun": false, "is_alias": false,
      "merge_decision": "new"
    },
    {
      "entity_type": "pronoun referring to a person", "text": "He", "score": 0.89,
      "start": 31, "end": 33,
      "canonical_id": "person:tim-cook", "canonical_text": "Tim Cook",
      "coref_cluster_id": "coref_0_a1b2c3d4",
      "is_pronoun": true, "is_alias": false,
      "merge_decision": "rule"
    }
  ]
}
```

**Chunk-level rollup:**
- Pronouns excluded from chunk rollup (canonical entity already present)
- Deduplicated by `(canonical_id, entity_type)`
- `chunk_metadata.entity_types` = sorted flat list of unique types

---

## 13. Dependencies

```toml
[project.optional-dependencies]
ner = [
    "gliner>=2.0.0",
    "torch>=2.0.0",
    "rapidfuzz>=3.0.0",
]
coref = [
    "fastcoref>=2.1.0",
    "spacy>=3.7",
]
canon-embeddings = [
    "sentence-transformers>=2.2",
    "numpy>=1.24",
]
```

Coref and embeddings are separate opt-in extras — users who only need NER+canonicalization don't pay the fastcoref+spaCy install cost.

---

## 14. Limitations

1. **fastcoref requires full document in RAM.** Manageable for typical documents (PDFs < 500 pages). Very large documents may need chunked processing.

2. **Pronoun resolution accuracy ~85% on business/technical documents.** Unresolved pronouns always emitted with `canonical_text=None` — never silently dropped.

3. **Cross-doc clustering is O(N²) for new entities.** Use `EntityRegistry` + `--registry-path` to skip re-clustering of already-known entities on incremental runs.

4. **`sentence-transformers` is optional** — install only for the 5–15% of entity pairs that fall in the "review" zone between auto-merge and reject thresholds.
