Attribution and Citation: Sourcing LLM Outputs › Code Implementation

# subprocess.run(["uv", "pip", "install", "transformers", "torch"], check=False, capture_output=True)
# Run once locally, then comment out

from transformers import pipeline

# Use a lightweight cross-encoder NLI model
nli_model = pipeline(
    "text-classification",
    model="cross-encoder/nli-deberta-v3-small",
    device=-1,  # CPU
)

def evaluate_citation(claim: str, passage: str, model) -> dict:
    """
    Evaluate whether a passage supports a claim using NLI.
    Returns entailment, neutral, and contradiction probabilities.
    """
    result = model(f"{passage} [SEP] {claim}", top_k=None)
    # result is a list of dicts: [{"label": ..., "score": ...}, ...]
    # When batched, it may be wrapped in an extra list; unwrap if needed
    if result and isinstance(result[0], list):
        result = result[0]

    scores = {item["label"].lower(): item["score"] for item in result}
    return {
        "entailment": scores.get("entailment", 0.0),
        "neutral": scores.get("neutral", 0.0),
        "contradiction": scores.get("contradiction", 0.0),
        "is_supported": scores.get("entailment", 0.0) > 0.5,
    }


def evaluate_attributed_response(
    claims_with_citations: list, passages: dict, model
) -> list:
    """
    Evaluate all claim-citation pairs in a response.

    claims_with_citations: list of {"claim": str, "citation_ids": list[int]}
    passages: dict mapping passage ID to passage text
    """
    results = []
    for item in claims_with_citations:
        claim = item["claim"]
        cited_passages = [
            passages[cid] for cid in item["citation_ids"] if cid in passages
        ]

        claim_results = []
        for idx, passage in zip(item["citation_ids"], cited_passages):
            score = evaluate_citation(claim, passage, model)
            claim_results.append(
                {
                    "passage_id": idx,
                    "passage_preview": passage[:80] + "..."
                    if len(passage) > 80
                    else passage,
                    **score,
                }
            )

        # A claim is considered supported if at least one cited passage entails it
        any_supported = any(r["is_supported"] for r in claim_results)
        results.append(
            {
                "claim": claim,
                "citation_results": claim_results,
                "claim_supported": any_supported,
            }
        )

    return results

# Simulated retrieved passages (as a RAG system would provide)
passages = {
    1: "The transformer architecture was introduced in the paper 'Attention is All You Need' by Vaswani et al. in 2017. It replaced recurrent neural networks with self-attention mechanisms.",
    2: "BERT (Bidirectional Encoder Representations from Transformers) was published by Devlin et al. in 2019. It uses masked language modeling for pretraining.",
    3: "The Amazon rainforest covers approximately 5.5 million square kilometers and is home to 10% of all species on Earth.",
}

# Simulated attributed response claims (what the model generated with citations)
claims_with_citations = [
    {
        "claim": "The transformer model was proposed by Vaswani et al. in 2017.",
        "citation_ids": [1],  # Correctly cited
    },
    {
        "claim": "BERT was first introduced in 2017.",  # Wrong date!
        "citation_ids": [
            2
        ],  # Passage 2 actually says 2019, but claim says 2017
    },
    {
        "claim": "Transformers use self-attention instead of recurrence.",
        "citation_ids": [
            3
        ],  # Completely wrong citation; passage 3 is about the Amazon
    },
]

# Evaluate attribution
evaluation_results = evaluate_attributed_response(
    claims_with_citations, passages, nli_model
)

def compute_citation_metrics(evaluation_results: list) -> dict:
    """
    Compute citation-level precision, recall, and F1.
    Precision: fraction of cited passages that support the claim.
    Recall: fraction of supported claims that have at least one supporting citation.
    """
    total_citations = 0
    supported_citations = 0
    total_claims = len(evaluation_results)
    supported_claims = 0

    for res in evaluation_results:
        for cr in res["citation_results"]:
            total_citations += 1
            if cr["is_supported"]:
                supported_citations += 1

        if res["claim_supported"]:
            supported_claims += 1

    citation_precision = (
        supported_citations / total_citations if total_citations > 0 else 0.0
    )
    claim_recall = supported_claims / total_claims if total_claims > 0 else 0.0

    f1 = (
        2
        * citation_precision
        * claim_recall
        / (citation_precision + claim_recall)
        if (citation_precision + claim_recall) > 0
        else 0.0
    )

    return {
        "citation_precision": citation_precision,
        "claim_recall": claim_recall,
        "f1": f1,
        "total_citations": total_citations,
        "supported_citations": supported_citations,
        "total_claims": total_claims,
        "supported_claims": supported_claims,
    }


metrics = compute_citation_metrics(evaluation_results)