Logits ProcessorでQwen（LLM）のハルシネーション対策

{
  "Title": "業務改善提案",
  "Content": "会議の効率化を図るため、..."
}

"Content": "会議の効率化を図るため、アジェンダを事前に共有し、アジェンダを事前に共有し、アジェンダを事前に共有し、アジェンダを事前に共有し..."

generation_config = {
    "repetition_penalty": 1.3,
}

generation_config = {
    "no_repeat_ngram_size": 3,
}

from __future__ import annotations

from datetime import datetime
from typing import Optional

import torch
from transformers.generation.logits_process import (
    LogitsProcessor,
    LogitsProcessorList,
    _calc_banned_ngram_tokens,
)


class ContentNoRepeatNGramLogitsProcessor(LogitsProcessor):
    """JSONのContentフィールド内でのみno-repeat-ngramを適用するLogitsProcessor"""

    def __init__(
        self,
        tokenizer,
        ngram_size: int,
        decode_max_tokens: int = 2048,
        debug: bool = False,
    ):
        if not isinstance(ngram_size, int) or ngram_size <= 0:
            raise ValueError(
                f"`ngram_size` has to be a strictly positive integer, but is {ngram_size}"
            )
        if not isinstance(decode_max_tokens, int) or decode_max_tokens <= 0:
            raise ValueError(
                "`decode_max_tokens` has to be a strictly positive integer, "
                f"but is {decode_max_tokens}"
            )
        self.tokenizer = tokenizer
        self.ngram_size = ngram_size
        self.decode_max_tokens = decode_max_tokens
        self.debug = debug
        self._debug_calls = 0
        self._debug_hits = 0

    @staticmethod
    def _extract_open_content_text(decoded_text: str) -> Optional[str]:
        """
        デコード済みテキストから、開いた状態のContentフィールドのテキストを抽出する。
        閉じていれば（"が見つかれば）Noneを返す。
        """
        marker = '"Content":"'
        marker_index = decoded_text.rfind(marker)
        if marker_index == -1:
            return None

        content_start = marker_index + len(marker)
        escaped = False
        for index in range(content_start, len(decoded_text)):
            char = decoded_text[index]
            if escaped:
                escaped = False
                continue
            if char == "\\":
                escaped = True
                continue
            if char == '"':
                # 閉じ引用符が見つかった→Content終了済み
                return None
        return decoded_text[content_start:]

    @staticmethod
    def _debug_timestamp() -> str:
        return datetime.now().strftime("%H:%M:%S")

    def __call__(
        self, input_ids: torch.LongTensor, scores: torch.FloatTensor
    ) -> torch.FloatTensor:
        self._debug_calls += 1
        num_batch_hypotheses = scores.shape[0]
        scores_processed = scores.clone()

        for i in range(num_batch_hypotheses):
            # デコードするトークン数を制限してパフォーマンスを確保
            decode_ids = input_ids[i, -self.decode_max_tokens :]
            decoded_text = self.tokenizer.decode(
                decode_ids,
                skip_special_tokens=False,
                clean_up_tokenization_spaces=False,
            )

            # Contentフィールドが開いているか確認
            content_text = self._extract_open_content_text(decoded_text)
            if content_text is None:
                if self.debug:
                    print(
                        f"[{self._debug_timestamp()}][ContentNoRepeat] "
                        f"content_text is None. {decoded_text[-80:]}"
                    )
                continue

            if self.debug:
                preview = content_text[-80:].replace("\n", "\\n")
                print(
                    f"[{self._debug_timestamp()}][ContentNoRepeat] "
                    f"step={self._debug_calls} batch={i} "
                    f"content_chars={len(content_text)} tail={preview!r}"
                )

            # Contentフィールドのテキストをトークン化
            content_ids = self.tokenizer.encode(
                content_text,
                add_special_tokens=False,
            )
            if len(content_ids) + 1 < self.ngram_size:
                # まだngram_sizeに達していない→制限不要
                continue

            content_input_ids = input_ids.new_tensor(content_ids).unsqueeze(0)
            banned_batch_tokens = _calc_banned_ngram_tokens(
                self.ngram_size,
                content_input_ids,
                1,
                content_input_ids.shape[-1],
            )

            banned_tokens = banned_batch_tokens[0]
            if banned_tokens:
                self._debug_hits += 1
                if self.debug:
                    print(
                        f"[{self._debug_timestamp()}][ContentNoRepeat] "
                        f"banned {len(banned_tokens)} token(s) "
                        f"for batch={i} at step={self._debug_calls}"
                    )
                # 禁止トークンのlogitを-infに設定→実質的に出力不可にする
                scores_processed[i, banned_tokens] = -float("inf")

        return scores_processed


class ContentNoRepeatGenerationMixin:
    """CLIやモデルクラスに組み込むためのMixin"""

    @staticmethod
    def add_content_no_repeat_cli_args(parser) -> None:
        parser.add_argument(
            "--content_no_repeat_ngram_size",
            type=int,
            default=0,
            help='JSONの"Content"フィールド内にのみno-repeat-ngramを適用するサイズ（0で無効）',
        )
        parser.add_argument(
            "--content_no_repeat_decode_max_tokens",
            type=int,
            default=1024,
            help='Contentフィールドの検出に使う最大トークン数',
        )
        parser.add_argument(
            "--content_no_repeat_debug",
            action="store_true",
            help='デバッグログを出力する',
        )

    @staticmethod
    def build_content_no_repeat_logits_processor(
        tokenizer,
        content_no_repeat_ngram_size: int = 0,
        content_no_repeat_decode_max_tokens: int = 2048,
        content_no_repeat_debug: bool = False,
    ) -> Optional[LogitsProcessorList]:
        processors = LogitsProcessorList()
        if content_no_repeat_ngram_size and content_no_repeat_ngram_size > 0:
            processors.append(
                ContentNoRepeatNGramLogitsProcessor(
                    tokenizer=tokenizer,
                    ngram_size=content_no_repeat_ngram_size,
                    decode_max_tokens=content_no_repeat_decode_max_tokens,
                    debug=content_no_repeat_debug,
                )
            )
        return processors or None

pip install transformers torch

from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers.generation.logits_process import LogitsProcessorList

# モデルとトークナイザーの読み込み
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-7B-Instruct")

# LogitsProcessorの構築
logits_processor = ContentNoRepeatGenerationMixin.build_content_no_repeat_logits_processor(
    tokenizer=tokenizer,
    content_no_repeat_ngram_size=4,       # 4-gramで繰り返しを禁止
    content_no_repeat_decode_max_tokens=1024,
    content_no_repeat_debug=True,          # 開発時はTrueで動作確認
)

# 推論
inputs = tokenizer("JSONを生成してください：", return_tensors="pt")
generation_config = {
    "max_new_tokens": 512,
    "temperature": 0.7,
    "do_sample": True,
}

output_ids = model.generate(
    **inputs,
    **generation_config,
    logits_processor=logits_processor,
)

print(tokenizer.decode(output_ids[0], skip_special_tokens=True))