LLM for Coreference Resolution: Challenges and Opportunities

#engineering #oxlo #ai

Coreference resolution is still a bottleneck in pipelines that need to track entities across messy, real-world text. In this tutorial, I will show you how to ship a small resolver that clusters mentions into entities and maps them back to exact character spans. I run this on Oxlo.ai because cost does not scale with input length, so sending full paragraphs as context does not inflate the bill the way token-based providers do.

What you'll need

Python 3.10 or newer
An Oxlo.ai API key from https://portal.oxlo.ai
The OpenAI SDK: pip install openai

Step 1: Configure the client and data models

I start with a minimal dataclass schema so downstream code always works with native Python objects instead of raw JSON. The Oxlo.ai client uses the standard OpenAI SDK pattern with a custom base URL.

from openai import OpenAI
from dataclasses import dataclass
from typing import List
import json

client = OpenAI(base_url="https://api.oxlo.ai/v1", api_key="YOUR_OXLO_API_KEY")

@dataclass
class Mention:
    text: str
    start: int
    end: int

@dataclass
class Cluster:
    representative: str
    mentions: List[Mention]

Step 2: Write the system prompt

The system prompt is the entire specification. I force the model to emit exact mention strings in order of appearance so I can align them mechanically in Python instead of trusting character offsets from the LLM.

SYSTEM_PROMPT = """You are a coreference resolution engine. Analyze the text and identify all mentions that refer to the same real-world entity.

Rules:
- Output ONLY valid JSON. No markdown, no explanation.
- Top level key is "clusters".
- Each cluster has:
  - "representative": the most specific full name for the entity.
  - "mentions": list of exact substrings from the text, in the order they appear.
- Include proper names, pronouns, and descriptive noun phrases.
- Every mention string MUST appear verbatim in the input text."""

Step 3: Implement the resolver function

This function calls Oxlo.ai with JSON mode enabled and parses the response into our schema. I use llama-3.3-70b here because it follows structured instructions reliably for extraction tasks.

def resolve_coreference(text: str, model: str = "llama-3.3-70b") -> List[Cluster]:
    response = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": text},
        ],
        response_format={"type": "json_object"},
        temperature=0.1,
    )

    raw = response.choices[0].message.content
    data = json.loads(raw)

    clusters = []
    for item in data.get("clusters", []):
        clusters.append(Cluster(
            representative=item.get("representative", "unknown"),
            mentions=[]
        ))
        for mention_text in item.get("mentions", []):
            clusters[-1].mentions.append(Mention(text=mention_text, start=-1, end=-1))
    return clusters

Step 4: Align mentions to exact text spans

LLMs are bad at counting characters, so I recover exact spans by searching for each mention string sequentially through the original text. This guarantees that every offset is verifiable.

def align_mentions(text: str, clusters: List[Cluster]) -> List[Cluster]:
    aligned = []
    for cluster in clusters:
        mentions = []
        search_from = 0
        for mention in cluster.mentions:
            m_text = mention.text
            idx = text.find(m_text, search_from)
            if idx == -1:
                idx = text.find(m_text)
            if idx != -1:
                mentions.append(Mention(
                    text=m_text,
                    start=idx,
                    end=idx + len(m_text)
                ))
                search_from = idx + len(m_text)
        if mentions:
            aligned.append(Cluster(
                representative=cluster.representative,
                mentions=mentions
            ))
    return aligned

Step 5: Wire it into a single call

I bundle the resolution and alignment stages so the caller gets fully resolved clusters with character offsets in one line. This is the public API I expose to the rest of the pipeline.

def resolve(text: str, model: str = "llama-3.3-70b") -> List[Cluster]:
    raw_clusters = resolve_coreference(text, model)
    return align_mentions(text, raw_clusters)

Run it

Here is a realistic paragraph with ambiguous pronouns. I run it through the resolver and print the clusters.

if __name__ == "__main__":
    sample = ("Dr. Elena Varga presented her findings. "
              "She said her model was robust. "
              "Varga proved she was right.")

    clusters = resolve(sample, model="llama-3.3-70b")

    for c in clusters:
        print(f"\nCluster: {c.representative}")
        for m in c.mentions:
            print(f"  '{m.text}' [{m.start}:{m.end}]")

The output should look something like this:

Cluster: Dr. Elena Varga
  'Dr. Elena Varga' [0:15]
  'her' [26:29]
  'She' [40:43]
  'her' [49:52]
  'Varga' [71:76]
  'she' [84:87]

Wrap up and next steps

This resolver is already useful for preprocessing legal contracts or clinical notes before entity linking. Two concrete next steps: integrate spaCy mention detection to prune hallucinated clusters, and wire the character offsets into a Label Studio annotation project for human review.

DEV Community