構造化出力の制御（JSON・関数呼び出し） — LLMの出力をJSON等に制約するConstrained Decodingの理論と実装

あなたは情報抽出アシスタントです。
以下のテキストから人物情報を抽出し、次のJSON形式で返してください。

{
  "name": "氏名",
  "age": 年齢（数値）,
  "occupation": "職業"
}

JSONのみを出力し、余計な説明は付けないでください。

{
  "type": "object",
  "properties": {
    "name": {"type": "string"},
    "age": {"type": "integer"},
    "is_student": {"type": "boolean"}
  },
  "required": ["name", "age", "is_student"]
}

{
  "function": "get_weather",
  "arguments": {
    "city": "Tokyo",
    "unit": "celsius"
  }
}

{
  "type": "array",
  "items": {
    "type": "object",
    "properties": {
      "function": {"enum": ["get_weather", "search_web"]},
      "arguments": { ... }
    }
  }
}

import numpy as np
import re
from collections import defaultdict
from typing import Dict, List, Set, Tuple, Optional


class SimpleDFA:
    """正規表現パターンから簡易DFAを構築するクラス"""

    def __init__(self, states: Set[int], alphabet: Set[str],
                 transitions: Dict[Tuple[int, str], int],
                 start_state: int, accept_states: Set[int]):
        self.states = states
        self.alphabet = alphabet
        self.transitions = transitions
        self.start_state = start_state
        self.accept_states = accept_states
        # 各状態から受理状態に到達可能かを事前計算
        self.reachable = self._compute_reachability()

    def _compute_reachability(self) -> Set[int]:
        """受理状態から逆方向にBFSして到達可能な状態を列挙"""
        reachable = set(self.accept_states)
        # 逆遷移を構築
        reverse_trans = defaultdict(set)
        for (src, char), dst in self.transitions.items():
            reverse_trans[dst].add(src)
        # BFS
        queue = list(self.accept_states)
        while queue:
            state = queue.pop(0)
            for prev_state in reverse_trans[state]:
                if prev_state not in reachable:
                    reachable.add(prev_state)
                    queue.append(prev_state)
        return reachable

    def step(self, state: int, char: str) -> Optional[int]:
        """1文字遷移。遷移不可能なら None を返す"""
        return self.transitions.get((state, char), None)

    def step_token(self, state: int, token_str: str) -> Optional[int]:
        """トークン（複数文字）を順に遷移。途中で失敗したら None"""
        current = state
        for char in token_str:
            current = self.step(current, char)
            if current is None:
                return None
        return current

    def is_reachable_to_accept(self, state: int) -> bool:
        """この状態から受理状態に到達可能か"""
        return state in self.reachable


def build_integer_dfa() -> SimpleDFA:
    """整数 (0 | [1-9][0-9]*) にマッチするDFAを構築"""
    # 状態: 0=開始, 1=0を読んだ(受理), 2=1-9を読んだ(受理), 3=デッド
    states = {0, 1, 2, 3}
    digits = set("0123456789")
    alphabet = digits
    transitions = {}
    # 状態0: '0'→状態1, '1'-'9'→状態2
    transitions[(0, '0')] = 1
    for d in "123456789":
        transitions[(0, d)] = 2
    # 状態1: 0の後は何も来ない（JSONの整数: 0は単独）
    for d in digits:
        transitions[(1, d)] = 3  # デッド
    # 状態2: 後続の数字→状態2（ループ）
    for d in digits:
        transitions[(2, d)] = 2
    # 状態3: デッドステート
    for d in digits:
        transitions[(3, d)] = 3
    return SimpleDFA(states, alphabet, transitions,
                     start_state=0, accept_states={1, 2})

import matplotlib.pyplot as plt


class ConstrainedDecoder:
    """DFAベースのConstrained Decodingエンジン"""

    def __init__(self, dfa: SimpleDFA, vocabulary: List[str]):
        self.dfa = dfa
        self.vocabulary = vocabulary
        self.vocab_size = len(vocabulary)
        # トークンレベル遷移テーブルを事前計算
        self.token_transitions = self._precompute_token_transitions()

    def _precompute_token_transitions(self) -> Dict[Tuple[int, int], Optional[int]]:
        """全 (状態, トークン) ペアに対する遷移先を事前計算"""
        table = {}
        for state in self.dfa.states:
            for tok_idx, tok_str in enumerate(self.vocabulary):
                next_state = self.dfa.step_token(state, tok_str)
                if next_state is not None and self.dfa.is_reachable_to_accept(next_state):
                    table[(state, tok_idx)] = next_state
                else:
                    table[(state, tok_idx)] = None
        return table

    def get_allowed_tokens(self, state: int) -> List[int]:
        """現在の状態で許可されるトークンインデックスのリスト"""
        allowed = []
        for tok_idx in range(self.vocab_size):
            if self.token_transitions.get((state, tok_idx)) is not None:
                allowed.append(tok_idx)
        return allowed

    def apply_mask(self, logits: np.ndarray, state: int) -> np.ndarray:
        """ロジットにトークンマスクを適用"""
        allowed = self.get_allowed_tokens(state)
        masked_logits = np.full_like(logits, -np.inf)
        for idx in allowed:
            masked_logits[idx] = logits[idx]
        return masked_logits

    def sample(self, logits: np.ndarray, temperature: float = 1.0) -> int:
        """softmaxサンプリング"""
        if temperature <= 0:
            return int(np.argmax(logits))
        scaled = logits / temperature
        # 数値安定性のためにmaxを引く
        scaled -= np.max(scaled[scaled > -np.inf]) if np.any(scaled > -np.inf) else 0
        exp_logits = np.exp(scaled)
        probs = exp_logits / np.sum(exp_logits)
        return int(np.random.choice(len(probs), p=probs))

    def generate(self, logits_sequence: List[np.ndarray],
                 temperature: float = 1.0) -> Tuple[List[int], List[str]]:
        """制約付きデコーディングでトークン列を生成"""
        state = self.dfa.start_state
        token_ids = []
        token_strs = []
        for logits in logits_sequence:
            masked_logits = self.apply_mask(logits, state)
            if np.all(masked_logits == -np.inf):
                break  # 遷移可能なトークンがない → 生成終了
            tok_idx = self.sample(masked_logits, temperature)
            token_ids.append(tok_idx)
            token_strs.append(self.vocabulary[tok_idx])
            # 状態を更新
            state = self.token_transitions[(state, tok_idx)]
            # 受理状態に到達したら終了
            if state in self.dfa.accept_states:
                # 次のステップで遷移可能なトークンがあるか確認
                if not self.get_allowed_tokens(state):
                    break
        return token_ids, token_strs

np.random.seed(42)

# 簡易語彙: 数字の各桁 + 非数字トークン
vocabulary = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9",
              " ", ",", ".", "a", "b", "{", "}", '"']
vocab_size = len(vocabulary)

# 整数DFAを構築
dfa = build_integer_dfa()
decoder = ConstrainedDecoder(dfa, vocabulary)

# 擬似ロジットを生成（実際にはLLMが出力する）
# 数字トークンに高めのロジットを設定しつつ、非数字トークンにもある程度の確率を与える
def make_mock_logits(vocab_size: int, prefer_digits: bool = True) -> np.ndarray:
    logits = np.random.randn(vocab_size) * 0.5
    if prefer_digits:
        logits[:10] += 2.0  # 数字トークンを優遇
    else:
        logits[10:] += 1.5  # 非数字トークンを優遇
    return logits

n_trials = 1000
n_steps = 5

# 制約なしサンプリング
unconstrained_valid = 0
unconstrained_outputs = []
for _ in range(n_trials):
    output_tokens = []
    for step in range(n_steps):
        logits = make_mock_logits(vocab_size, prefer_digits=(step < 3))
        probs = np.exp(logits - np.max(logits))
        probs /= probs.sum()
        tok_idx = np.random.choice(vocab_size, p=probs)
        output_tokens.append(vocabulary[tok_idx])
    output_str = "".join(output_tokens)
    unconstrained_outputs.append(output_str)
    # 整数として有効か検証
    if re.match(r'^(0|[1-9][0-9]*)$', output_str):
        unconstrained_valid += 1

# 制約付きサンプリング
constrained_valid = 0
constrained_outputs = []
for _ in range(n_trials):
    logits_seq = [make_mock_logits(vocab_size, prefer_digits=(s < 3))
                  for s in range(n_steps)]
    token_ids, token_strs = decoder.generate(logits_seq, temperature=1.0)
    output_str = "".join(token_strs)
    constrained_outputs.append(output_str)
    if re.match(r'^(0|[1-9][0-9]*)$', output_str):
        constrained_valid += 1

print(f"制約なし: {unconstrained_valid}/{n_trials} 件が有効な整数 "
      f"({100*unconstrained_valid/n_trials:.1f}%)")
print(f"制約付き: {constrained_valid}/{n_trials} 件が有効な整数 "
      f"({100*constrained_valid/n_trials:.1f}%)")
print(f"\n制約なしの出力例: {unconstrained_outputs[:8]}")
print(f"制約付きの出力例: {constrained_outputs[:8]}")

fig, axes = plt.subplots(1, 3, figsize=(16, 5))

# 3つの状態でのマスキング効果を可視化
states_to_show = [
    (0, "State 0 (Start)"),
    (2, "State 2 (After digit 1-9)"),
    (1, "State 1 (After '0')")
]

for ax, (state, title) in zip(axes, states_to_show):
    logits = make_mock_logits(vocab_size, prefer_digits=True)

    # 元の確率分布
    probs_orig = np.exp(logits - np.max(logits))
    probs_orig /= probs_orig.sum()

    # 制約付き確率分布
    masked_logits = decoder.apply_mask(logits, state)
    valid_mask = masked_logits > -np.inf
    if np.any(valid_mask):
        probs_constrained = np.zeros(vocab_size)
        exp_vals = np.exp(masked_logits[valid_mask] - np.max(masked_logits[valid_mask]))
        probs_constrained[valid_mask] = exp_vals / exp_vals.sum()
    else:
        probs_constrained = np.zeros(vocab_size)

    x = np.arange(vocab_size)
    width = 0.35
    bars1 = ax.bar(x - width/2, probs_orig, width, label='Original',
                   color='#4a9eff', alpha=0.7)
    bars2 = ax.bar(x + width/2, probs_constrained, width, label='Constrained',
                   color='#ff6b6b', alpha=0.7)

    ax.set_xlabel('Token')
    ax.set_ylabel('Probability')
    ax.set_title(title)
    ax.set_xticks(x)
    ax.set_xticklabels(vocabulary, fontsize=8, rotation=45)
    ax.legend(fontsize=8)
    ax.grid(axis='y', alpha=0.3)

    # マスクされた確率質量を表示
    masked_mass = sum(probs_orig[i] for i in range(vocab_size) if not valid_mask[i])
    ax.text(0.02, 0.95, f'Masked mass: {masked_mass:.2%}',
            transform=ax.transAxes, fontsize=8,
            verticalalignment='top',
            bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

plt.tight_layout()
plt.savefig('constrained_decoding_distributions.png', dpi=150, bbox_inches='tight')
plt.show()

fig, ax = plt.subplots(1, 1, figsize=(10, 6))

# 非数字トークンの優遇度を変化させて、制約の影響度を測定
bias_levels = np.linspace(-2, 4, 20)
masked_masses = []
effective_vocab_ratios = []

for bias in bias_levels:
    logits = np.random.randn(vocab_size) * 0.5
    logits[:10] += 2.0       # 数字トークンの基本ロジット
    logits[10:] += bias      # 非数字トークンのバイアスを変化
    probs = np.exp(logits - np.max(logits))
    probs /= probs.sum()
    allowed = decoder.get_allowed_tokens(0)  # 開始状態での許可トークン
    masked_mass = sum(probs[i] for i in range(vocab_size) if i not in allowed)
    masked_masses.append(masked_mass)
    effective_vocab_ratios.append(len(allowed) / vocab_size)

ax.plot(bias_levels, masked_masses, 'o-', color='#ff6b6b',
        linewidth=2, markersize=6, label='Masked probability mass')
ax.axhline(y=0.5, color='gray', linestyle='--', alpha=0.5, label='50% threshold')
ax.fill_between(bias_levels, masked_masses, alpha=0.1, color='#ff6b6b')
ax.set_xlabel('Non-digit token bias', fontsize=12)
ax.set_ylabel('Masked probability mass', fontsize=12)
ax.set_title('Impact of Constraint Strength on Probability Mass', fontsize=14)
ax.legend(fontsize=11)
ax.grid(alpha=0.3)
ax.set_ylim(0, 1.0)

plt.tight_layout()
plt.savefig('constraint_impact.png', dpi=150, bbox_inches='tight')
plt.show()

class JsonObjectDFA:
    """
    固定スキーマのJSONオブジェクト生成用の簡易状態マシン
    スキーマ例: {"name": string, "age": integer}
    生成されるパターン: {"name": "...", "age": 数字}
    """

    def __init__(self, fields: List[Tuple[str, str]]):
        """
        fields: [(フィールド名, 型)] のリスト
        型は 'string' または 'integer' のみサポート
        """
        self.fields = fields
        self.states = {}  # 状態名 → 状態ID
        self.transitions = {}
        self.accept_states = set()
        self._build_states()

    def _build_states(self):
        """スキーマからDFA状態を構築"""
        state_id = 0
        # 状態を順に作成
        # '{' を期待する開始状態
        self.start_state = state_id
        self.states['start'] = state_id

        # 各フィールドについて状態を生成
        for i, (fname, ftype) in enumerate(self.fields):
            # '"' でフィールド名開始
            state_id += 1
            self.states[f'field_{i}_quote_open'] = state_id
            # フィールド名の各文字
            for j, char in enumerate(fname):
                state_id += 1
                self.states[f'field_{i}_name_{j}'] = state_id
            # '"' でフィールド名終了
            state_id += 1
            self.states[f'field_{i}_quote_close'] = state_id
            # ':' を期待
            state_id += 1
            self.states[f'field_{i}_colon'] = state_id
            # 値の開始
            state_id += 1
            self.states[f'field_{i}_value_start'] = state_id
            if ftype == 'string':
                # 文字列値の中身
                state_id += 1
                self.states[f'field_{i}_value_content'] = state_id
                # '"' で文字列値終了
                state_id += 1
                self.states[f'field_{i}_value_end'] = state_id
            elif ftype == 'integer':
                # 整数値（1桁以上の数字）
                state_id += 1
                self.states[f'field_{i}_value_digits'] = state_id

            # フィールド間の ',' または最後の '}'
            if i < len(self.fields) - 1:
                state_id += 1
                self.states[f'field_{i}_comma'] = state_id

        # 終了の '}'
        state_id += 1
        self.states['end_brace'] = state_id
        self.accept_states = {state_id}

        # 全状態集合
        self.all_states = set(range(state_id + 1))

    def get_allowed_chars(self, state_name: str) -> Set[str]:
        """状態名に基づいて許可される文字集合を返す"""
        if state_name == 'start':
            return {'{'}
        elif '_quote_open' in state_name:
            return {'"'}
        elif '_name_' in state_name:
            # フィールド名の次の文字を決定
            parts = state_name.split('_')
            field_idx = int(parts[1])
            char_idx = int(parts[3])
            fname = self.fields[field_idx][0]
            if char_idx < len(fname):
                return {fname[char_idx]}
            return {'"'}
        elif '_quote_close' in state_name:
            return {'"'}
        elif '_colon' in state_name:
            return {':'}
        elif '_value_start' in state_name:
            parts = state_name.split('_')
            field_idx = int(parts[1])
            ftype = self.fields[field_idx][1]
            if ftype == 'string':
                return {'"'}
            elif ftype == 'integer':
                return set('123456789')
        elif '_value_content' in state_name:
            # 文字列内容: 英数字 + 終了の '"'
            return set('abcdefghijklmnopqrstuvwxyz '
                      'ABCDEFGHIJKLMNOPQRSTUVWXYZ') | {'"'}
        elif '_value_digits' in state_name:
            parts = state_name.split('_')
            field_idx = int(parts[1])
            # 次がカンマか閉じ括弧かで許可文字が変わる
            return set('0123456789') | {',' if field_idx < len(self.fields) - 1 else '}'}
        elif '_value_end' in state_name:
            parts = state_name.split('_')
            field_idx = int(parts[1])
            if field_idx < len(self.fields) - 1:
                return {','}
            else:
                return {'}'}
        elif '_comma' in state_name:
            return {'"'}
        elif state_name == 'end_brace':
            return set()  # 受理状態、遷移なし
        return set()


# スキーマ定義
schema_fields = [("name", "string"), ("age", "integer")]
json_dfa = JsonObjectDFA(schema_fields)

# 各状態で許可される文字を表示
print("=== JSON Object DFA States ===")
for name, sid in sorted(json_dfa.states.items(), key=lambda x: x[1]):
    allowed = json_dfa.get_allowed_chars(name)
    is_accept = "  [ACCEPT]" if sid in json_dfa.accept_states else ""
    allowed_display = sorted(allowed) if allowed else "(none)"
    print(f"  State {sid:2d} ({name:30s}): allowed={allowed_display}{is_accept}")

def simulate_json_generation(json_dfa: JsonObjectDFA,
                             max_steps: int = 100) -> str:
    """JSONオブジェクトの制約付き生成をシミュレーション"""
    output = []
    state_names = {v: k for k, v in json_dfa.states.items()}
    current_state_id = json_dfa.states['start']

    for step in range(max_steps):
        current_name = state_names.get(current_state_id, 'unknown')
        allowed_chars = json_dfa.get_allowed_chars(current_name)
        if not allowed_chars:
            break

        # 擬似ロジット: 各許可文字にランダムなスコアを割り当て
        char_list = sorted(allowed_chars)
        logits = np.random.randn(len(char_list)) * 0.5

        # 文字列値では英字を、整数値では特定の数字を若干優遇
        for i, c in enumerate(char_list):
            if c.isalpha():
                logits[i] += 1.0
            elif c.isdigit():
                logits[i] += 0.5

        # softmaxサンプリング
        probs = np.exp(logits - np.max(logits))
        probs /= probs.sum()
        chosen_idx = np.random.choice(len(char_list), p=probs)
        chosen_char = char_list[chosen_idx]
        output.append(chosen_char)

        # 状態遷移（簡易版: 文字に応じて次状態を決定）
        current_state_id += 1
        if current_state_id in json_dfa.accept_states:
            break
        if current_state_id >= max(json_dfa.states.values()):
            break

    return "".join(output)

np.random.seed(123)
print("=== 制約付きJSON生成のシミュレーション ===\n")
for trial in range(5):
    result = simulate_json_generation(json_dfa)
    # 有効なJSONかチェック
    import json
    try:
        parsed = json.loads(result)
        valid = True
    except json.JSONDecodeError:
        valid = False
    print(f"Trial {trial+1}: {result}")
    print(f"  Valid JSON: {valid}")
    if valid:
        print(f"  Parsed: {parsed}")
    print()

fig, ax = plt.subplots(1, 1, figsize=(10, 6))

# 制約の強さと品質の関係を模式的に可視化
constraint_strength = np.linspace(0, 1, 100)

# 形式の正しさ: 制約なし→確率的、制約あり→100%
format_correctness_unconstrained = 0.85 - 0.3 * constraint_strength
format_correctness_constrained = np.ones_like(constraint_strength)

# 意味の品質: 制約が強すぎると低下
semantic_quality = 1.0 - 0.5 * constraint_strength**2

ax.plot(constraint_strength, format_correctness_unconstrained,
        '--', color='#4a9eff', linewidth=2, label='Format correctness (unconstrained)')
ax.plot(constraint_strength, format_correctness_constrained,
        '-', color='#4a9eff', linewidth=2, label='Format correctness (constrained)')
ax.plot(constraint_strength, semantic_quality,
        '-', color='#ff6b6b', linewidth=2, label='Semantic quality (constrained)')
ax.fill_between(constraint_strength,
                format_correctness_constrained, semantic_quality,
                alpha=0.1, color='green', label='Quality gap')

ax.set_xlabel('Constraint Strength (masked probability mass)', fontsize=12)
ax.set_ylabel('Quality Score', fontsize=12)
ax.set_title('Trade-off: Format Correctness vs. Semantic Quality', fontsize=14)
ax.legend(fontsize=10, loc='lower left')
ax.grid(alpha=0.3)
ax.set_ylim(0, 1.1)
ax.set_xlim(0, 1)

plt.tight_layout()
plt.savefig('quality_tradeoff.png', dpi=150, bbox_inches='tight')
plt.show()