Structured Outputs の仕組み JSONスキーマのガイダンス

[Structured Outputs](https://platform.openai.com/docs/guides/structured-outputs) を用いることで LLM（Large Language Model）の出力を厳密に JSONスキーマに準拠させられますが、逐次的にテキストを生成する LLM で、どうすればそのようなことを実現できるのか。その仕組みを [Hugging Face の Guidance](https://huggingface.co/docs/text-generation-inference/conceptual/guidance) を参考に調査しました。


結論からいうと有限状態マシン（FSM: Finite State Machine、有限オートマトンともいう）を用いて LLM が予測する次トークンの候補から JSONスキーマに準拠するものだけを選抜します。

例えば、"ACC" または "ABC" のいずれかのテキストだけを生成したい場合、その有限状態マシンは次のような状態遷移図で表せます。

```mermaid
graph LR
    %% 初期状態
    start_node[ ] -- 開始 --> 0;
    style start_node fill:#fff,stroke:#fff

    %% 状態定義
    0(0);
    1(1);
    2(2);
    3(3);
    4(4);

    %% 遷移定義
    0 -- A --> 1;
    1 -- B --> 3;
    1 -- C --> 2;
    2 -- C --> 4;
    3 -- C --> 4;

    4 -- 終了 --> final_node[ ];
    style final_node fill:#fff,stroke:#fff
```

ここで、0 から 4 の数値の箱は状態を、状態から次の状態へ遷移する矢印はその状態で入力可能な文字を表します。例えば初期状態０は文字 A だけを許容し、A が入力されたら状態１へ遷移します。状態１ は B 又は C を許容し、C が入力されたら状態２へ、B が入力されたら状態３へ遷移します。

さて、文字 A をLLM へ入力すると次の文字（次トークン）を次表のように予測したとします。

|次トークン|確率|
|:-:|:-:|
|A|0.1|
|B|0.2|
|C|0.2|
|D|0.5|


このままだと A の次の文字は D の可能性が高いですが、有限状態マシンに従うと A の次は B または C しか許容されません。なのでそれ以外の文字 A, D を候補から除外し残りの候補から次の文字を選択します（次表）。

|次のトークン|確率|
|:-:|:-:|
|A|0|
|B|0.2|
|C|0.2|
|D|0|


この操作を繰り返すことで有限状態マシンに従うテキストを生成します。すなわち JSONスキーマに基づく有限状態マシンを用いれば、JSONスキーマに準拠したテキストを生成できるということです。

有限状態マシンは正規表現から作成できるので、JSONスキーマを正規表現へ変換できればその正規表現から JSONスキーマの有限状態マシンを作成することができます。

以降では、Python にてパッケージ [outlines-core](https://github.com/dottxt-ai/outlines-core
) と [interegular](https://github.com/MegaIng/interegular) を用いて JSONスキーマの有限状態マシンを作成します。outlines-core で JSONスキーマを正規表現へ変換し、interegular でその正規表現の有限状態マシンを作成します（次図）。



```mermaid
flowchart LR
    A[JSONスキーマ] -->|outlines-core| B(正規表現)
    B -->|interegular | C[有限状態マシン]
```



# 実行環境

<b>実行環境</b>： Google Colab

<b>バージョン</b>：
- Python 3.11.11
- outlines-core 0.1.26
- interegular 0.3.3

outlines-coreをインストールすると依存関係で interegular もインストールされます。
```
pip install -q outlines-core
```

# 正規表現の有限状態マシン作成
JSONスキーマへ取り掛かる前に正規表現 `A(B|C)C` の有限状態マシンを作成してみます。

```python
import outlines_core
import interegular

# 正規表現から有限状態マシンを作成
fsm = interegular.parse_pattern(r'A(B|C)C').to_fsm()

print(fsm)
```
実行結果
```
  name final? A B C anything_else 
----------------------------------
* 0    False  1                   
  1    False    2 3               
  2    False      4               
  3    False      4               
  4    True                       
```

ここで `name` は状態を表します。 `*` は初期状態を表し、`final?` は最終状態であるか否かを表します。`A B C` はその文字を受理した場合の遷移先です。

冒頭の有限状態マシンの状態遷移図はこれを図にしたものです。

```mermaid
graph LR
    %% 初期状態
    start_node[ ] -- 開始 --> 0;
    style start_node fill:#fff,stroke:#fff

    %% 状態定義
    0(0);
    1(1);
    2(2);
    3(3);
    4(4);

    %% 遷移定義
    0 -- A --> 1;
    1 -- B --> 3;
    1 -- C --> 2;
    2 -- C --> 4;
    3 -- C --> 4;

    4 -- 終了 --> final_node[ ];
    style final_node fill:#fff,stroke:#fff
```



# 有限状態マシンで逐次評価
次の関数 `accepts` は、有限状態マシンを用いて入力 `inputs` を1文字ずつ読み取り逐次的に評価するコード例です。

```python
def accepts(fsm, inputs):
    anything_else = interegular.fsm.anything_else

    state = fsm.initial
    for symbol in inputs:
        print(symbol, end='')
        if anything_else in fsm.alphabet and not symbol in fsm.alphabet:
            print(" is anything_else", end='')
            symbol = anything_else
        transition = fsm.alphabet[symbol]

        # Missing transition = transition to dead state
        if not (state in fsm.map and transition in fsm.map[state]):
            print()
            return False

        post_state = fsm.map[state][transition]
        print(f" {state} -> {post_state}", end='')
        state = post_state
        print()

    return state in fsm.finals


print(accepts(fsm, 'ACC'))
print()
print(accepts(fsm, 'ABC'))
print()
print(accepts(fsm, 'AAC'))
```
このコードは [interegular の accepts 関数](https://github.com/MegaIng/interegular/blob/master/interegular/fsm.py#L220) に少し手を加えたものです。

実行結果
```
A 0 -> 1
C 1 -> 3
C 3 -> 4
True

A 0 -> 1
B 1 -> 2
C 2 -> 4
True

A 0 -> 1
A
False
```
ここで有限状態マシンに合致する "ACC" と "ABC" は戻り値が `True` ですが、合致しない "AAC" は2文字目の "A" で `False` となっていることが分かります。 


雰囲気が分かったので次は、JSONスキーマの有限状態マシンを作成します。


# JSONスキーマを正規表現へ変換
パッケージ outlines-core で JSONスキーマを正規表現へ変換します。

```mermaid
flowchart LR
    A[JSONスキーマ] -->|outlines-core| B(正規表現)
```

まず、Pydantic のモデル `User` を定義しその JSONスキーマを生成します。
```python
import json
from pydantic import BaseModel

# Pydantic モデル
class User(BaseModel):
    id: int
    name: str

# JSONスキーマを生成
schema = json.dumps(User.model_json_schema(), indent=2)

print(schema)
```
実行結果：JSONスキーマ
```json
{
  "properties": {
    "id": {
      "title": "Id",
      "type": "integer"
    },
    "name": {
      "title": "Name",
      "type": "string"
    }
  },
  "required": [
    "id",
    "name"
  ],
  "title": "User",
  "type": "object"
}
```

outlines-core の `build_regex_from_schema` 関数を使って JSONスキーマを正規表現へ変換します。
```python
from outlines_core.fsm.json_schema import build_regex_from_schema

# JSONスキーマを正規表現へ変換
regex_str = build_regex_from_schema(schema, whitespace_pattern=None)

print(regex_str)
```
実行結果：JSONスキーマの正規表現
```
\{[ ]?"id"[ ]?:[ ]?(-)?(0|[1-9][0-9]*)[ ]?,[ ]?"name"[ ]?:[ ]?"([^"\\\x00-\x1F\x7F-\x9F]|\\["\\])*"[ ]?\}
```

# 正規表現から有限状態マシンを作成
パッケージ interegular で JSONスキーマの正規表現から有限状態マシンを作成します。

```mermaid
flowchart LR
    B(正規表現) -->|interegular | C[有限状態マシン]
```

JSONスキーマの正規表現を用いて有限状態マシンを作成
```python
# 正規表現から有限状態マシンを作成
schema_fsm = interegular.parse_pattern(regex_str).to_fsm()
```

状態遷移表を確認します。
```python
print(schema_fsm)
```

次は、状態遷移表からの抜粋です。
```
  name final?     "  ,  -  0  1  2  3  4  5  6  7  8  9  :  \\ a  d  e  i  m  n  {  } anything_else 
-----------------------------------------------------------------------------------------------------
* 0    False                                                                     1                               
  1    False   2  3                                                                                              
  2    False      3                                                                                              
  3    False                                                            4                                        
  4    False                                                      5                                              
  5    False      6                                                                                              
  6    False   7                                         8                                                       
  7    False                                             8                                                       
  8    False   10       12 11 9  9  9  9  9  9  9  9  9                                                          
  9    False   15    14    13 13 13 13 13 13 13 13 13 13                                                         
  10   False            12 11 9  9  9  9  9  9  9  9  9                                                          
  11   False   15    14                                                                                          
  12   False               11 9  9  9  9  9  9  9  9  9                                                          
  13   False   15    14    13 13 13 13 13 13 13 13 13 13                                                         
  14   False   16 17                                                                                             
  15   False         14                                                                                          
  16   False      17                                                                                             
  17   False                                                                  18                                 
  18   False                                                   19                                                
  19   False                                                               20                                    
  20   False                                                         21                                          
  21   False      22                                                                                             
  22   False   23                                        24                                                      
  23   False                                             24                                                      
  24   False   25 26                                                                                             
  25   False      26                                                                                             
  26   False   27 29 27 27 27 27 27 27 27 27 27 27 27 27 27 28 27 27 27 27 27 27 27 27 27            
  27   False   27 29 27 27 27 27 27 27 27 27 27 27 27 27 27 28 27 27 27 27 27 27 27 27 27            
  28   False      30                                        30                                                   
  29   False   31                                                                   32                           
  30   False   27 29 27 27 27 27 27 27 27 27 27 27 27 27 27 28 27 27 27 27 27 27 27 27 27            
  31   False                                                                        32                           
  32   True                                                                                                      
```


# 有限状態マシンでJSONを評価
JSON文字列が JSONスキーマに準拠するかどうかを有限状態マシンで評価します。
```python
# 有限状態マシンでJSON文字列を評価
accepts(schema_fsm, '{"id": 123, "name": "アリス"}')
```
実行結果
```
{ 0 -> 1
" 1 -> 3
i 3 -> 4
d 4 -> 5
" 5 -> 6
: 6 -> 8
  8 -> 10
1 10 -> 9
2 9 -> 13
3 13 -> 13
, 13 -> 14
  14 -> 16
" 16 -> 17
n 17 -> 18
a 18 -> 19
m 19 -> 20
e 20 -> 21
" 21 -> 22
: 22 -> 24
  24 -> 25
" 25 -> 26
ア is anything_else 26 -> 27
リ is anything_else 27 -> 27
ス is anything_else 27 -> 27
" 27 -> 29
} 29 -> 32
True
```
`True` なのでJSONスキーマに準拠しています。

JSONスキーマに準拠しない場合は、その時点で `False` が返ります。
```python
accepts(schema_fsm, '{"foo": "bar"}')
```
実行結果
```
{ 0 -> 1
" 1 -> 3
f is anything_else
False
```


以上です。
最後までお読みいただきありがとうございます。


# 参考文献
### ドキュメント Hugging Face Guidance
https://huggingface.co/docs/text-generation-inference/conceptual/guidance

### ドキュメント Outlines
https://dottxt-ai.github.io/outlines/latest/

### パッケージ outlines-core
https://github.com/dottxt-ai/outlines-core

### パッケージ interegular
https://github.com/MegaIng/interegular