custom-generate/hallucination_logits_processor.py

import torch
from transformers import LogitsProcessor

from .detectors.base_detector import BaseHallucinationDetector


class HallucinationLogitsProcessor(LogitsProcessor):
    """
    Generalized logits processor that works with any BaseHallucinationDetector.

    Optimizations:
      - debug_print: when False, ALL per-token prints are suppressed
        (counters still update). Significant runtime impact.
      - decoded-prefix cache: avoid O(N^2) re-decoding of the growing
        sequence on every step. With beam search, beams reorder per step,
        so the cache validates against the actual current_ids and falls
        back to a full decode on mismatch.
    """

    def __init__(
        self,
        hallucination_detector: BaseHallucinationDetector,
        last_k_tokens_to_consider: int = 4,
        top_k_logits: int = 10,
        penalty_value: float = float("-inf"),
        use_all_tokens: bool = False,
        skip_threshold: float = 1.0,
        debug_print: bool = False,
    ):
        self.hallucination_detector    = hallucination_detector
        self.last_k_tokens_to_consider = last_k_tokens_to_consider
        self.penalty_value             = penalty_value
        self.top_k_logits              = top_k_logits
        self.use_all_tokens            = use_all_tokens
        self.skip_threshold            = skip_threshold
        self.debug_print               = debug_print

        # Counters
        self.modifications_count = 0
        self.skip_count          = 0
        self.check_count         = 0

        # Decoded-prefix cache, per beam batch index. Each entry is a tuple
        # (last_id_int, decoded_text_so_far). On mismatch, we fall back to
        # full re-decode.
        self._decoded_cache: dict = {}

        if self.debug_print:
            tokens_context = (
                "all tokens" if use_all_tokens
                else f"last {last_k_tokens_to_consider} tokens"
            )
            print(
                f"Initialized LogitsProcessor with "
                f"{type(hallucination_detector).__name__}"
            )
            print(f"Context window: {tokens_context}")
            print(f"Skip threshold: {skip_threshold}")

    def _get_current_text(self, batch_idx: int, current_ids: torch.LongTensor) -> str:
        """
        Return the decoded text for current_ids, using a per-beam cache to
        avoid O(N^2) work. Cache invalidates on:
          - first call for this batch_idx (cold start)
          - the trailing token id changed without growing (beam reorder)
          - the sequence shrank (shouldn't happen, defensive)
        """
        cur_len = len(current_ids)
        cached = self._decoded_cache.get(batch_idx)
        tokenizer = self.hallucination_detector.tokenizer

        # Hot path: pure single-token append from a known previous state.
        if cached is not None:
            prev_len, prev_text = cached
            if cur_len == prev_len + 1:
                # Decode-the-difference: get just the appended token's text.
                # Decode from prev_len-1 onwards so leading-space markers
                # resolve correctly (BPE quirk on first token of a chunk).
                if prev_len > 0:
                    boundary_text = tokenizer.decode(
                        current_ids[prev_len - 1: prev_len + 1].tolist(),
                        skip_special_tokens=True,
                    )
                    last_token_only = tokenizer.decode(
                        current_ids[prev_len - 1: prev_len].tolist(),
                        skip_special_tokens=True,
                    )
                    appended_text = boundary_text[len(last_token_only):]
                else:
                    appended_text = tokenizer.decode(
                        current_ids[: cur_len].tolist(),
                        skip_special_tokens=True,
                    )
                new_text = prev_text + appended_text
                self._decoded_cache[batch_idx] = (cur_len, new_text)
                return new_text

        # Cold path / cache invalidation: full decode.
        new_text = tokenizer.decode(current_ids, skip_special_tokens=True)
        self._decoded_cache[batch_idx] = (cur_len, new_text)
        return new_text

    def __call__(
        self, input_ids: torch.LongTensor, scores: torch.FloatTensor,
    ) -> torch.FloatTensor:
        batch_size = input_ids.shape[0]

        for batch_idx in range(batch_size):
            current_ids = input_ids[batch_idx]
            current_text = self._get_current_text(batch_idx, current_ids)

            if self.use_all_tokens:
                context_window = len(current_ids)
            else:
                context_window = min(
                    self.last_k_tokens_to_consider, len(current_ids),
                )

            top_k_scores, top_k_indices = torch.topk(
                scores[batch_idx],
                k=min(self.top_k_logits, scores.shape[-1]),
            )
            top_k_probs = torch.softmax(top_k_scores, dim=-1)

            if self.debug_print:
                top_k_tokens = [
                    self.hallucination_detector.tokenizer.decode([idx.item()])
                    for idx in top_k_indices
                ]
                print("Top K tokens:")
                for rank, (token_str, prob) in enumerate(
                    zip(top_k_tokens, top_k_probs.tolist())
                ):
                    print(f"  {rank}: {repr(token_str)} ({prob:.4f})")

            self.check_count += 1

            # Skip HDM check when generator confidence STRICTLY exceeds the
            # threshold. With skip_threshold = 1.0, this branch is never taken.
            if top_k_probs[0].item() > self.skip_threshold:
                self.skip_count += 1
                if self.debug_print:
                    top_token_str = self.hallucination_detector.tokenizer.decode(
                        [top_k_indices[0].item()]
                    )
                    print(
                        f"Skipping check: top token {repr(top_token_str)} "
                        f"has prob {top_k_probs[0].item():.4f} > "
                        f"{self.skip_threshold}"
                    )
                continue

            # Walk top-k candidates, penalize the first hallucinatory one.
            for token_id in top_k_indices:
                token_id_item = token_id.item()

                if self.hallucination_detector.check_hallucination(
                    current_text, token_id_item, context_window,
                ):
                    scores[batch_idx][token_id_item] = self.penalty_value
                    self.modifications_count += 1
                    if self.debug_print:
                        token_str = self.hallucination_detector.tokenizer.decode(
                            [token_id_item]
                        )
                        print(
                            f"----------Modified logit for token: "
                            f"{token_str}----------"
                        )
                    continue

                if self.debug_print:
                    chosen_token_str = self.hallucination_detector.tokenizer.decode(
                        [token_id_item]
                    )
                    chosen_prob = torch.softmax(
                        scores[batch_idx], dim=-1
                    )[token_id_item].item()
                    input_text = self.hallucination_detector.input_text
                    if input_text in current_text:
                        generated_text = current_text[
                            current_text.find(input_text) + len(input_text):
                        ]
                    else:
                        generated_text = current_text
                    print(
                        f"Chosen token: {repr(chosen_token_str)} "
                        f"({chosen_prob:.4f}) | Generated so far: "
                        f"{repr(generated_text)}"
                    )

                # Greedy break after first non-hallucinatory candidate.
                break

        return scores