Hallucination Detection: NLI, Self-Consistency & Learned Models › Code Implementation

# Install required packages
import subprocess

subprocess.run(["uv", "pip", "install", "transformers", "torch"], check=True)

from typing import Dict, List

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer


class EntailmentDetector:
    """Detects hallucinations using NLI-based entailment scoring."""

    def __init__(self, model_name: str = "cross-encoder/nli-deberta-v3-small"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(
            model_name
        )
        self.model.train(False)  # Set to inference mode
        # Labels order varies by model; check the id2label mapping
        self.label2id = {v: k for k, v in self.model.config.id2label.items()}
        self.entail_id = self.label2id.get(
            "ENTAILMENT", self.label2id.get("entailment", 2)
        )
        self.contradict_id = self.label2id.get(
            "CONTRADICTION", self.label2id.get("contradiction", 0)
        )

    def score_pair(self, premise: str, hypothesis: str) -> Dict[str, float]:
        """Score a single premise-hypothesis pair."""
        inputs = self.tokenizer(
            premise,
            hypothesis,
            return_tensors="pt",
            truncation=True,
            max_length=512,
            padding=True,
        )
        with torch.no_grad():
            logits = self.model(**inputs).logits
        probs = torch.softmax(logits, dim=-1).squeeze().numpy()
        return {
            "entailment": float(probs[self.entail_id]),
            "contradiction": float(probs[self.contradict_id]),
            "neutral": float(
                1.0 - probs[self.entail_id] - probs[self.contradict_id]
            ),
        }

    def score_document(
        self, source: str, generated_sentences: List[str]
    ) -> List[Dict[str, float]]:
        """Score each generated sentence against the source document."""
        return [self.score_pair(source, sent) for sent in generated_sentences]

source = """
Marie Curie was a Polish and naturalised-French physicist and chemist who conducted
pioneering research on radioactivity. She was the first woman to win a Nobel Prize,
the first person to win the Nobel Prize twice, and the only person to win the Nobel Prize
in two scientific fields. She won the Nobel Prize in Physics in 1903 and the Nobel Prize
in Chemistry in 1911.
"""

faithful_claims = [
    "Marie Curie won the Nobel Prize in Physics in 1903.",
    "She was the first woman to receive a Nobel Prize.",
    "Curie won Nobel Prizes in two different scientific fields.",
]

hallucinated_claims = [
    "Marie Curie was born in France.",
    "She won the Nobel Prize in Physics in 1905.",
    "Curie was the only woman ever to win the Nobel Prize.",
]

detector = EntailmentDetector()
faithful_scores = detector.score_document(source, faithful_claims)
hallucinated_scores = detector.score_document(source, hallucinated_claims)

import re
from typing import Tuple

import numpy as np


def split_into_sentences(text: str) -> List[str]:
    """Simple sentence splitter using regex."""
    sentences = re.split(r"(?<=[.!?])\s+", text.strip())
    return [s for s in sentences if len(s) > 10]


def summac_score(
    source: str, generated_text: str, detector: EntailmentDetector
) -> Tuple[float, List[float]]:
    """
    Compute SummaC-style score for generated text against source.
    Returns (document_score, per_sentence_scores).
    """
    source_sentences = split_into_sentences(source)
    gen_sentences = split_into_sentences(generated_text)

    row_scores = []
    for gen_sent in gen_sentences:
        # Score the generated sentence against every source sentence
        sentence_scores = []
        for src_sent in source_sentences:
            result = detector.score_pair(src_sent, gen_sent)
            sentence_scores.append(result["entailment"])
        # Take the maximum across all source sentences
        row_scores.append(max(sentence_scores))

    doc_score = float(np.mean(row_scores)) if row_scores else 0.0
    return doc_score, row_scores

faithful_text = (
    "Marie Curie won the Nobel Prize in Physics in 1903. "
    "She was the first woman to receive this honor. "
    "Curie's research focused on radioactivity."
)

hallucinated_text = (
    "Marie Curie conducted pioneering research on radioactivity. "
    "She won the Nobel Prize in Physics in 1905. "
    "She was the first woman to win a Nobel Prize."
)

faithful_doc_score, faithful_sent_scores = summac_score(
    source, faithful_text, detector
)
hallucinated_doc_score, hallucinated_sent_scores = summac_score(
    source, hallucinated_text, detector
)

def selfcheck_nli_score(
    claim: str, sampled_passages: List[str], detector: EntailmentDetector
) -> float:
    """
    Compute SelfCheckGPT NLI score for a claim against sampled passages.
    Returns the average contradiction probability across samples.
    Higher = more likely to be a hallucination.
    """
    contradiction_scores = []
    for passage in sampled_passages:
        scores = detector.score_pair(passage, claim)
        contradiction_scores.append(scores["contradiction"])
    return float(np.mean(contradiction_scores))


def selfcheck_document(
    generated_sentences: List[str],
    sampled_passages: List[str],
    detector: EntailmentDetector,
) -> List[float]:
    """Score each sentence in the generated text against sampled passages."""
    return [
        selfcheck_nli_score(sent, sampled_passages, detector)
        for sent in generated_sentences
    ]

# Simulate 3 independent samples from a model about Marie Curie.
# In a real system, these would be generated by the LLM with temperature > 0.
sampled_passages = [
    """Marie Curie was a physicist and chemist born in Warsaw, Poland.
    She conducted pioneering research on radioactivity and was the first woman
    to win a Nobel Prize, receiving the award in Physics in 1903.""",
    """Marie Curie, originally from Poland, became famous for her work on radioactivity.
    She received the Nobel Prize in Physics in 1903, becoming the first woman to win
    the award. She later won a second Nobel Prize in Chemistry in 1911.""",
    """The Polish-born scientist Marie Curie made groundbreaking discoveries in radioactivity.
    She was awarded the Nobel Prize in Physics in 1903 and the Nobel Prize in Chemistry
    in 1911, making her the first person to win Nobel Prizes in two sciences.""",
]

# Claims to evaluate: two true and two false
test_sentences = [
    "Marie Curie won the Nobel Prize in Physics in 1903.",  # True
    "Marie Curie was born in Paris, France.",  # False
    "Marie Curie was the first woman to win a Nobel Prize.",  # True
    "Marie Curie won three Nobel Prizes throughout her career.",  # False
]

selfcheck_scores = selfcheck_document(
    test_sentences, sampled_passages, detector
)