LMQL(Language Model Query Language)概観

ChatGPTのように対話的にAI（LLM）を扱う場合は、AIがズレた応答をしてきたときでも人間が補正／修正することで適切なアウトプットを求めることができます。しかしソフトウェアの中にLLMを組み込む場合は、自動化が目的なのもあって人間による補正を働かせることは困難なので、そのアウトプットには一定の安定性が求められます。

一方で、安定的な性能が出るプロンプトを作成するには特殊なノウハウが必要であり、そのためにプロンプトエンジニアリングという技術があります。ただこのプロンプトエンジニアリングにおいても、自然言語で書く以上は品質にばらつきが出ることは避けられず、またプロンプト変更時のインパクトを推し量ることも困難です。

LMQL（Language Model Query Language）はこのような問題を解決するために開発されている大規模言語モデル（LLM）向けのプログラミング言語です。プログラミング言語のような形式化による品質の安定化と、自然言語による処理記述の柔軟さを両手取りするような発想で面白いなと思ったので試してみました。

LMQL: Programming Large Language Models

LMQL Playgroundでクエリを試す

LMQLには動作を簡単に検証できるPlaygroundが用意されています。ローカルでPlaygroundを起動することもできます。

LMQL Playground

まずはGetting Startedで紹介されている以下のクエリを実行します。

argmax "Hello[WHO]" from "openai/text-ada-001" where len(WHO) < 10

「Run」ボタンをクリックするとOpenAIのAPI KEYを求められるので、入力します。

実行するとModel Responseの枠に結果が表示されます。

LMQLの基本構造

LMQLは記法的にはSQLと似ていて、以下のような構造を持っています。

デコーダ節（Decoder Clause）：
テキスト生成に使用するデコード・アルゴリズムを指定します。LMQLでは様々なデコード・アルゴリズムを選択することができ、今回の例ではargmaxデコードが選択されています。

プロンプト節（Prompt Clause）：
"Hello[WHO]"と記述されている部分がプロンプト節にあたります。[WHO]のように大カッコで囲まれた箇所はテンプレート変数となっており、モデルによって自動的に補完されます。

モデル節（Model Clause）：
from "openai/text-ada-001"と記述されている部分がモデル節にあたります。ここではテキスト生成に使用するモデルを指定します。現在のLMQLではGPT-3.5、GPT-4などのOpenAIモデルと、Transformersによるセルフホスティングモデルがサポートされています。

制約節（Constraint Clause）：
where len(WHO) < 10と記述されている部分が制約節にあたります。例えば今回の例では[WHO]にあたる箇所を10文字以内で文字生成するよう制約をつけています。

上記のLMQLは以下のようなコードにコンパイルされて実行されます。これもPlayground上から確認することができます。

import lmql
@lmql.compiled_query(output_variables=["WHO"])
async def query(WHO=None,argmax=None,context=None,len=None):
   yield lmql.runtime_support.context_call("set_model", 'openai/text-ada-001')
   yield lmql.runtime_support.context_call("set_decoder", 'argmax', )
   # where
   intm0 = lmql.ops.Lt([lmql.ops.LenOp([lmql.runtime_support.Var('WHO')]), 10])
   yield lmql.runtime_support.context_call("set_where_clause", intm0)
   # prompt
   (yield lmql.runtime_support.interrupt_call('query', f'Hello[WHO]'))
   WHO = (yield lmql.runtime_support.context_call('get_var', 'WHO'))
   yield ('result', (yield lmql.runtime_support.context_call("get_return_value", ())))

より複雑な例

より複雑な例として以下のLMQLを試してみます。

argmax
   """Review: We had a great stay. Hiking in the mountains was fabulous and the food is really good.
   Q: What is the underlying sentiment of this review and why?
   A:[ANALYSIS]
   Based on this, the overall sentiment of the message can be considered to be[CLASSIFICATION]"""
from
   "openai/text-davinci-003"
where
   not "\n" in ANALYSIS and CLASSIFICATION in [" positive", " neutral", " negative"]

実行すると以下のような出力になり、[ANALYSYS]、[CLASSIFICATION]の部分に制約通りの文字列が生成されていることが分かります。

プロンプト節は単に長くなっただけですが、制約節では[ANALYSIS]には改行が入らないよう、[CLASSIFICATION]にはpositive／neutral／negativeのいずれかの値が入るように制約しています。これを自然言語のみのプロンプトで書こうとすると表現と共に出力される値も揺れることがありますが、LMQLであれば機械的に記述することができます。

また、[CLASSIFICATION]の値によってプロンプト内で分岐させることも可能です。

argmax
   """Review: We had a great stay. Hiking in the mountains was fabulous and the food is really good.
   Q: What is the underlying sentiment of this review and why?
   A:[ANALYSIS]
   Based on this, the overall sentiment of the message can be considered to be[CLASSIFICATION]"""
   if CLASSIFICATION == " positive":
      "What is it that they liked about their stay? [FURTHER_ANALYSIS]"
   elif CLASSIFICATION == " neutral":
      "What is it that could have been improved? [FURTHER_ANALYSIS]"
   elif CLASSIFICATION == " negative":
      "What is it that they did not like about their stay? [FURTHER_ANALYSIS]"
from
    "openai/text-davinci-003"
where
    not "\n" in ANALYSIS and CLASSIFICATION in [" positive", " neutral", " negative"] and STOPS_AT(FURTHER_ANALYSIS, ".")

このようにif文を使ったり、for文を使ったりしてプロンプトを柔軟に記述する方法が用意されています。

Scripted Prompting

Constraints

Pythonとのインテグレーション

プログラムに組み込むことが目的なので、もちろんPython内からLMQLを呼び出すことができます。呼び出しのコードは以下のような形になります。

import lmql

@lmql.query
async def hello():
    '''lmql
    argmax
        "Hello[WHO]"
    from
        "openai/text-ada-001"
    where
        len(WHO) < 10
    '''

(await hello())[0].prompt

Python関数にlmql.queryデコレータをつけ、クエリコードを複数行の文字列として指定します。デコレータのついた関数は自動的にLMQLクエリにコンパイルされ、元のPython関数と同じ引数で呼び出すことができます。関数の戻り値はLMQLクエリの結果となります。

また、LMQLクエリはすべて非同期で実行されるため、クエリを実行するためには上記のコードにもある通りawaitキーワードを使用する必要があります。非同期処理はasyncioライブラリに依存しています。

LMQLの結果は以下のような LMQLResultオブジェクトのリストで返ります。

class LMQLResult:
    # full prompt with all variables substituted
    prompt: str
    # a dictionary of all assigned template variable values
    variables: Dict[str, str]

LangChainとのインテグレーション

LangChainとインテグレーションすることも可能です。LMQLクエリをシーケンシャルチェインの一部として実行する例を試してみます。

まず必要なライブラリを読み込みます。

from langchain import LLMChain, PromptTemplate
from langchain.chat_models import ChatOpenAI
from langchain.prompts.chat import (ChatPromptTemplate,HumanMessagePromptTemplate)
from langchain.llms import OpenAI

import lmql

次にLangChainのプロンプトを実行するようなチェインを用意します。

# setup the LM to be used by langchain
llm = OpenAI(temperature=0.9)

human_message_prompt = HumanMessagePromptTemplate(
        prompt=PromptTemplate(
            template="What is a good name for a company that makes {product}?",
            input_variables=["product"],
        )
    )
chat_prompt_template = ChatPromptTemplate.from_messages([human_message_prompt])
chat = ChatOpenAI(temperature=0.9)
chain = LLMChain(llm=chat, prompt=chat_prompt_template)

実行すると以下のような結果が返ります。

chain.run("colorful socks")

Rainbow Socks Inc.

次にLMQLクエリを用意します。

@lmql.query
async def write_catch_phrase(company_name: str):
    '''
    argmax "Write a catchphrase for the following company: {company_name}. [catchphrase]" from "chatgpt"
    '''

このクエリは会社名を引数にとって、その会社にふさわしいキャッチフレーズを考えます。単体で実行すると以下のようになります。

(await write_catch_phrase("Socks Inc"))[0].variables["catchphrase"]

"Step up your style with Socks Inc."

これをシーケンシャルチェインで繋げてみると、LangChainで定義したクエリとLMQLで定義したクエリを次のように連続で実行することができます。

from langchain.chains import SimpleSequentialChain
overall_chain = SimpleSequentialChain(chains=[chain, write_catch_phrase], verbose=True)

# Run the chain specifying only the input variable for the first chain.
catchphrase = overall_chain.run("colorful socks")
print(catchphrase)

> Entering new SimpleSequentialChain chain...
RainbowSocks Co.
"Step into a world of color with RainbowSocks Co.!"

> Finished chain.
"Step into a world of color with RainbowSocks Co.!"

所感

LMQLクエリをコンパイルした後のコードを見てみると、全てをLLMの実行結果に頼るのではなく、期待するアウトプットにあわせてコードレベルでもLLMからのアウトプットを補正するような処理が書かれています。

実際、LLMによる処理をコードに埋め込もうとすると、かなり揺れ幅のあるLLMからの出力結果と上手く付き合うような文字列処理の記述を行う必要があるため、その処理を自動化してくれるのと同時に、可読性の高い形でプロンプトを管理できそうなLMQLには期待が持てるなと感じました。

現場からは以上です。