導入

LLM（特にtransformers等を使ったローカルLLM推論）において、構造化テキスト出力を支援するパッケージは様々なものが出ています。

代表的なものとしてOutlinesや、Guidance、lm-format-enforcerなどがあります。

そんな中で、比較的新しいパッケージとしてFormatronというものが出ているのを知りました。

READMEに記載の特徴を邦訳して抜粋すると以下のような感じ。

🔗 人気のライブラリ統合: transformers、exllamav2、vllm、RWKVをサポート。

🔌 プラグイン、ラッパーではない: サードパーティライブラリを大きくて扱いにくいクラスでラップするのではなく、Formatronは異なるライブラリのための便利でクリーンなプラグインを提供。

💡 ライブラリ、フレームワークではない: すべてを大きなフレームワークに統一するのではなく、Formatronはどこにでも埋め込むことができる柔軟なライブラリ。

✍️ 流暢なフォーマット: 自然言語を書くように簡単にフォーマットを記述。

📜 正規表現とCFGサポート: 正規表現と文脈自由文法（CFG）をフォーマットに簡単に組み込む。

⚙️ 効率的なJSON生成: Pydanticモデルやjsonスキーマに基づいた完全なJSON生成機能。

📤 バッチ推論: 1つのバッチで各シーケンスに異なるフォーマットを自由に指定可能！

🚀 最小のランタイムオーバーヘッド: Leo最適化、特殊な圧縮アルゴリズム、世代間でのCFGキャッシュを使用し、Rustで実装されたEarleyアルゴリズムは理論的にも実際的にも最速のアルゴリズム。

🔧 カスタマイズ可能: スキーマ生成、文法生成、関数呼び出しなどの生成後処理を含むすべてが設定可能。

特にlm-format-enforcerなど他のパッケージと比較して高速な処理を謳っているようです。

最近はLlama.cppやSGLangなどの推論パッケージには構造化出力機能をはじめから備えたものも多いのですが、こちらのパッケージもかなり使い勝手が良さそうだったので主にREADMEに記載のサンプルを試してみました。

検証はDatabricks on AWS上で行っています。
今回は推論パッケージとしてExLlamaV2を利用しました。
試してないのですが、transformersやvllmなどのよく使われるパッケージにも対応しているようです。

Step1. パッケージインストール

ExLlamaV2とFormatronに必要なパッケージをインストール。

%pip install -U flash-attn --no-build-isolation
%pip install https://github.com/turboderp/exllamav2/releases/download/v0.2.3/exllamav2-0.2.3+cu121.torch2.3.1-cp311-cp311-linux_x86_64.whl
%pip install -U "formatron==0.4.4" referencing

dbutils.library.restartPython()

Step2. モデルのロード

ExLlamaV2でモデルをロードします。
今回は小型モデルでもちゃんと動くこと確認したかったので、Google Gemma-2 2B JPN-itをEXL2形式に量子化したモデルを利用しました。


from exllamav2 import ExLlamaV2, ExLlamaV2Config, ExLlamaV2Cache_Q8, ExLlamaV2Tokenizer
from exllamav2.generator import ExLlamaV2DynamicGenerator

from formatron.schemas.pydantic import ClassSchema
from formatron.integrations.exllamav2 import create_formatter_filter
from formatron.formatter import FormatterBuilder
from typing import Literal, Optional

def format_prompt(p: str):
    """ チャット用のプロンプトへ変換 """
    return f"""<start_of_turn>user\n{p}<end_of_turn><start_of_turn>model\n"""


def generate(generator, prompts, filters):
    """ 推論の実行 """
    f_prompts = [format_prompt(p) for p in prompts]

    return generator.generate(
        prompt=f_prompts,
        filters=filters,
        max_new_tokens=512,
        add_bos=True,
        completion_only=True,
        stop_conditions=[tokenizer.eos_token_id],
    )

# 8bit量子化したGemma2のパス
model_path = "/Volumes/training/llm/model_snapshots/models--google--gemma-2-2b-jpn-it-exl2-bpw8_0/"
config = ExLlamaV2Config(model_path)
config.arch_compat_overrides()
model = ExLlamaV2(config)
cache = ExLlamaV2Cache_Q8(model, max_seq_len = 32768, lazy = True)

print("Loading Model...")
model.load_autosplit(cache, progress = True)

print("Loading tokenizer...")
tokenizer = ExLlamaV2Tokenizer(config)

generator = ExLlamaV2DynamicGenerator(
    model = model,
    cache = cache,
    tokenizer = tokenizer,
)

これで準備完了です。では、いろいろな構造化出力を試してみます。

Step3. 正規表現出力を試してみる

整数を意味する正規表現を指定して、その形で出力させてみます。

# Formatronのフォーマット設定構築インスタンス。これに処理を追加していく
f = FormatterBuilder()
digit = f.regex('([1-9][0-9]*)', capture_name='digit')
f.append_str(f"私の好きな数字は {digit} です。") # この形で出力される


prompts = ["なんか適当な数字を教えーて"]
# ExLlamaV2用のフィルタ作成
filters = [[create_formatter_filter(model, tokenizer, f)]]

generate(generator, prompts, filters)

出力

['私の好きな数字は 7 です。']

指定した形で出力されました。
数字だけが必要な場合はappend_strで出力する内容をdigitsだけにすればOK。

Step4. JSON出力を試してみる

次にJSON出力を試してみます。

まず、サンプル例を使ってJSON書式の結果を返すようにしてみます。

from formatron.schemas.dict_inference import infer_mapping
import json

f = FormatterBuilder()
schema = infer_mapping({"name": "foo", "height": "10m"})
f.append_line(f"{f.json(schema, capture_name='json')}")

prompts = ["東京には東京タワーという観光名所があります。高さは333mです。この文章から情報を抜き出してJSON形式で出力してください。"]
filters = [[create_formatter_filter(model, tokenizer, f)]]

output = generate(generator, prompts, filters)

# 結果文字列の出力
print(output)
# JSON変換してもエラー出ない
print(json.loads(output[0]))

出力

['{ \t"name": "東京タワー", \t"height": "333m" }\n']
{'name': '東京タワー', 'height': '333m'}

PydanticのClassSchemaを継承したクラスを使ってスキーマ指定をすることもできます。


# Pydanticでスキーマ情報作成
class Goods(ClassSchema):
    name: str
    price: float
    remaining: int

f = FormatterBuilder()
schema = Goods
f.append_line(f"{f.json(schema, capture_name='json')}")

prompts = ["リンゴが14個残っていて、それぞれの価格は300円です。この文から情報を抽出してJSONにしてください。"]
filters = [[create_formatter_filter(model, tokenizer, f)]]
output = generate(generator, prompts, filters)

# スキーマ情報
print("--- JSONスキーマ ---")
print(schema.schema())
# 結果
print("--- 推論結果 ---")
print(json.loads(output[0]))

出力

--- JSONスキーマ ---
{'properties': {'name': {'title': 'Name', 'type': 'string'}, 'price': {'title': 'Price', 'type': 'number'}, 'remaining': {'title': 'Remaining', 'type': 'integer'}}, 'required': ['name', 'price', 'remaining'], 'title': 'Goods', 'type': 'object'}
--- 推論結果 ---
{'name': 'リンゴ', 'price': 300, 'remaining': 14}

Step5. Function Callsを試してみる

特定のデコレータがついた関数を指定して、Tool/Function Call用の出力を得ることもできます。

from formatron import schemas

@schemas.pydantic.callable_schema
def add(a: int, b: int, /, *, c: int):
    return a + b + c

f = FormatterBuilder()
f.append_line(f"{f.json(add, capture_name='json')}")

prompts = ["8+12+5を計算したいです。適切な引数を含むjsonを生成してください。"]
filters = [[create_formatter_filter(model, tokenizer, f)]]

output = generate(generator, prompts, filters)

print("--- 推論結果 ---")
print(json.loads(output[0]))

出力

--- 推論結果 ---
{'a': 8, 'b': 12, 'c': 5}

Step6. JSONスキーマの直接指定を試してみる

直接JSONスキーマを指定することもできます。

from formatron.schemas import json_schema

schema = {
    "$id": "https://example.com/person.json",
    "$schema": "https://json-schema.org/draft/2020-12/schema",
    "type": "object",
    "properties": {
        "name": {
            "type": "string"
        },
        "age": {
            "type": "integer"
        }
    },
    "required": ["name", "age"]
}
schema = json_schema.create_schema(schema)

f = FormatterBuilder()
f.append_line(f"{f.json(schema, capture_name='json')}")

prompts = ["この文から情報をjsonに抽出してください: 私の名前はTomといいます。28歳の男性で趣味はサッカーです。"]
filters = [[create_formatter_filter(model, tokenizer, f)]]
output = generate(generator, prompts, filters)

print("--- 推論結果 ---")
print(json.loads(output[0]))

出力

--- 推論結果 ---
{'name': 'Tom', 'age': 28}

想定したフォーマット（構造）で結果を得ることができました。

まとめ

Formatronによる構造化出力を試してみました。
lm-format-enforcerに比べて早い・・・のかな？体感ではわからず。

書き方は割と直感的で使いやすい感じです。
Too/Function Call用に構造化出力させたい場合に使うのが良いかもしれません。
（Llama.cppを最初から使えばいいんじゃんって声も聞こえてきそうですが。。。）

軽く触っただけですが、プラガブルで使い勝手は良さそうに感じました。
機能的に少し足りないところがありますが、LLMの構造化出力制御に使っていこうと思います。