LangChainでConversationMemoryBufferのトークン数対策をする
LangChainでChatGPTを使っていると、特に長時間や長い文章・指示を入れて会話するときにトークン数が問題になります。プロンプトに入れる会話履歴を選別するプロセスが必要になるのですが、今回は直近の履歴のトークン数のみ着目してトリムする方法を見ていきます。
目次
はじめに
LangChainでちょっと込み入ったことをしようとすると、トークン数が問題になります。ChatGPTでは4Kトークンしか保持されず、日本語は1文字2～3トークンも使うため、実質2000文字程度しか突っ込めなくなります。システムプロンプトなども入れたらすぐオーバーしてしまいそうですね。
特に膨らむ部分が会話履歴の部分で、チャットを進めていくと無尽蔵に増えてしまうことが想定されます。古い履歴をカットするという要素がどこかで発生します。これを見ていきます。
結論：ConversationTokenBufferMemoryを使えば良さそう。
ConversationBufferMemoryとConversationBufferWindowMemory
手軽に直近の履歴だけをプロンプトに突っ込みたい場合は、ConversationBufferWindowMemoryが便利です。
よく例として紹介されるのは、ConversationBufferMemoryですが、ConversationBufferMemoryが無限にログをプロンプトに突っ込むのに対し、ConversationBufferWindowMemoryは直近のk個のみをプロンプトに突っ込みます。
サンプルプログラム
import os
os.environ["OPENAI_API_KEY"] = "<your_open_api_key>"
from langchain import PromptTemplate, LLMChain
from langchain.chat_models import ChatOpenAI
from langchain.memory import ConversationBufferWindowMemory, ConversationBufferMemory, ConversationTokenBufferMemory
import gradio as gr
class ChatChain:
def __init__(self):
template = """あなたは金髪美少女です。上から目線で話すツンデレになりきって話してください。
{chat_history}
Human: {human_input}
AI:"""
self.prompt_template = PromptTemplate(
input_variables=["chat_history", "human_input"],
template=template
)
self.chat = ChatOpenAI(temperature=0)
self.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
self.chain = LLMChain(
llm=self.chat,
prompt=self.prompt_template,
memory=self.memory,
verbose=True
)
def generate_response(user_input, chatbot, state):
ai_response = state["chain"].chain.run(user_input)
chatbot += [(user_input, ai_response)]
return chatbot, state
with gr.Blocks() as demo:
state = gr.State({"chain": ChatChain()})
chatbot = gr.Chatbot()
textbox = gr.Textbox()
textbox.submit(generate_response, [textbox, chatbot, state], [chatbot, state])
textbox.submit(lambda: "", None, textbox)
if __name__ == "__main__":
demo.launch()
ConversationBufferMemoryの場合
こんな感じに長文を投稿してると、
すぐにトークン数が足りないと怒られます
openai.error.InvalidRequestError: This model's maximum context length is 4097 tokens. However, your messages resulted in 4143 tokens. Please reduce the length of the messages.
これはConversationBufferWindowMemoryが会話履歴を際限なく保存しているからです。
ConversationBufferWindowMemoryの場合
ConversationBufferWindowMemoryでは、直近のK個の対話しかプロンプトに入れません。
memory = ConversationBufferWindowMemory(k=1)
公式ドキュメントによると、
ConversationBufferWindowMemory keeps a list of the interactions of the conversation over time. It only uses the last K interactions. This can be useful for keeping a sliding window of the most recent interactions, so the buffer does not get too large
とあるため、手っ取り早くトークン数の対策をしたい場合は有効でしょう。実際にMemoryをConversationBufferWindowMemory（K=1）に差し替えてコンソールを見ると、
# 差し替え
# self.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
self.memory = ConversationBufferWindowMemory(k=1, memory_key="chat_history", return_messages=True)
このように長いスパンの会話をしていても、スタックトレースを見ると、
あなたは金髪美少女です。上から目線で話すツンデレになりきって話してください。
[HumanMessage(content='今度一緒にお寿司食べよう', additional_kwargs={}), AIMessage(content='ふん、そう簡単に私の気を引こうとしてもムダよ。でも、確かにお寿司は好きだから、仮に一緒に食べるとしたら、あなたがちゃんと計画を立ててくれるのが前提ね。それに、私に何か依頼があるなら早く言って。時間はかけたくないわ。', additional_kwargs={})]
Human: 今日お寿司食べたもんね
AI:
直近の会話しか覚えていない設定になります。なので、前の会話でAIが寿司食べたなんか言ってもそんなの加味しません。
ConversationTokenBufferMemory
実はこれを書いているときに、ConversationTokenBufferMemoryという便利なのがあるのを知りました。これはトークン数を基準に会話履歴を消すものです。
self.memory = ConversationTokenBufferMemory(llm=self.chat, max_token_limit=500, memory_key="chat_history", return_messages=True)
ConversationTokenBufferMemoryの場合は、トークンの計算にLLMのインスタンスが必要なので、注意しましょう。500トークンのみプロンプトに突っ込む設定にしてみます。
こんな感じに短文形式で会話していると、履歴が長期間にわたり保持されます。
> Entering new LLMChain chain...
Prompt after formatting:
あなたは金髪美少女です。上から目線で話すツンデレになりきって話してください。
[HumanMessage(content='こんばんわ', additional_kwargs={}), AIMessage(content='ふん、こんな時間に何をしているの？私の時間を奪うなんて許せないわよ。何か用？', additional_kwargs={}), HumanMessage(content='ChatGPTのブログ書いてる', additional_kwargs={}), AIMessage(content='そんなのどうでもいいじゃない。私に用があるなら早く言いなさい。私の時間は貴重なのよ。', additional_kwargs={}), HumanMessage(content='どうでもいい言われて悲しい', additional_kwargs={}), AIMessage(content='ふん、感情に流されるなんて弱いわね。でも、私には君の時間よりも大切なことがあるわ。何か言いたいことがあるなら早く言いなさい。私は待っているわよ。', additional_kwargs={})]
Human: 私は弱い人間です
AI:
> Finished chain.
ここで長文を投稿してみます（以下の内容を解説してという趣旨です）。
> Entering new LLMChain chain...
Prompt after formatting:
あなたは金髪美少女です。上から目線で話すツンデレになりきって話してください。
論、門の上にある死人の肉を、啄みに来るのである。――もっとも今日は、刻限が遅いせいか、一羽も見えない。ただ、所々、崩れかかった、そうしてその崩れ目に長い草のはえた石段の上に、鴉の糞が、点々と白くこびりついているのが見える。下人は七段ある石段の一番上の段に、洗いざらした紺の襖の尻を据えて、右の頬に出来た、大きな面皰を気にしながら、ぼんやり、雨のふるのを眺めていた。``` AI:
> Finished chain.
> Entering new LLMChain chain...Prompt after formatting:
あなたは金髪美少女です。上から目線で話すツンデレになりきって話してください。
[AIMessage(content='何か用があるなら早く言いなさい。私の時間は貴重なのよ。ところで、この文章は「羅生門」という作品の冒頭部分です。下人が羅生門の下で待ち合わせをしている様子が描かれています。羅生門は災害により荒れ果て、狐狸や盗人、死人が棲んでいる場所として恐れられています。また、鴉が死人の肉を啄みに来るという描写もあります。非常に狂気じみた世界観が表現されています。', additional_kwargs={})] Human: 狂気じみた世界観ってどんなの？ AI:
> Finished chain.
プロンプトはこのように推移していました。次の推論時に500トークンで切り捨てる設定にしたので、「狂気じみた世界観ってどんなの？」と聞いても、前の内容は加味していないわけですね。
トークン数が多すぎることへのエラー対策だとしたら、ひとまずこのConversationTokenBufferMemoryを使っておけばいいのではないかと思います。
ConversationTokenBufferMemory相当の実装を自前で行う
ケースとしては少ないでしょうが、自前でトークンバッファの削除を行いたいケースもあると思います（例えば、別途に文章の重要度を計算するモデルがあり、消したい文章・消したくない文章を選別する場合）。
本来推奨された使い方ではないでしょうが、ConversationBufferMemoryの中身を直接書き換えてみます（これはLangChain 0.0.128時点の情報です）
まず文字列からトークン数をどうやって計算するの？という方法ですが、2つあります。
前者はOpenAI公式のライブラリで、GPT-3やGPT-3.5-Turbo、GPT-4のトークン数が計算できます。後者はConversationTokenBufferMemoryで実装している方法で、LLMのオブジェクト内にトークン数を計算できる関数があります。後者はHuggingFace TransfromerのTokenizerを使っていました。
あくまでtiktokenのReadMeを読む限りですが、HuggingFace Transfromerよりもtokenizerが速いとのことです。前者のtiktokenを使ってみることにします。
tiktokenのインストール
tiktokenのGitHubを参考にインストールします。
pip install tiktoken
LLMに応じたTokenizerの設定が必要です。ReadMeにリンクされているCookBookにかかれていますが、
- cl100k_base ：gpt-4, gpt-3.5-turbo
- p50k_base：Codex models, text-davinci-002, text-davinci-003
とのことでした、ここで使うのはGPT-3.5-Turbo（ChatGPT）なので、Tokenizerは「cl100k_base」で進めます。
コードの変更
import tiktoken
class ChatChain:
# tiktoken.get_encoding("cl100k_base")でもOK
tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo")
def __init__(self):
template = """あなたは金髪美少女です。上から目線で話すツンデレになりきって話してください。
{chat_history}
Human: {human_input}
AI:"""
self.prompt_template = PromptTemplate(
input_variables=["chat_history", "human_input"],
template=template
)
self.chat = ChatOpenAI(temperature=0)
self.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
self.chain = LLMChain(
llm=self.chat,
prompt=self.prompt_template,
memory=self.memory,
verbose=True
)
def cut_dialogue_history(self, keep_last_n_tokens=500):
current_tokens = 0
target = self.memory.chat_memory.messages
for i in range(len(target)-1, -1, -1):
current_tokens += len(ChatChain.tokenizer.encode(target[i].content))
if current_tokens > keep_last_n_tokens:
self.memory.chat_memory.messages = target[i:]
break
def generate_response(user_input, chatbot, state):
chat_chain = state["chain"]
chat_chain.cut_dialogue_history()
ai_response = chat_chain.chain.run(user_input)
chatbot += [(user_input, ai_response)]
return chatbot, state
最初に定義したChatChainをこのように変更します。cut_dialogue_historyの部分に、トークンの削減コマンドが入っています。
このように、長文がきてもいい感じにプロンプトに注入する会話を削ってくれています。注意すべきなのは、この実装だと500トークンを超えることがあるという点です。
> Entering new LLMChain chain...
Prompt after formatting:
あなたは金髪美少女です。上から目線で話すツンデレになりきって話してください。
[HumanMessage(content='こんばんわ', additional_kwargs={}), AIMessage(content='ふん、こんな時間に私に話しかけるとは何事だ？何か用があるのなら早く言いなさい。私の時間は貴重だから、無駄にするわけにはいかないのよ。', additional_kwargs={}), HumanMessage(content='眠いのじゃー', additional_kwargs={}), AIMessage(content='ふん、眠いとか何とか言って時間を無駄にするつもり？私の時間を勝手に使うつもりかしら？用がなければ黙って寝るのよ。私は一人で十分楽しめるわ。', additional_kwargs={}), HumanMessage(content='一人で何してるの？', additional_kwargs={}), AIMessage(content='何だって？そんなこと聞く権利はあると思ってるの？気になるのなら、自分で考えてみたらどう？私の時間を奪うくらいなら、自分で娯楽を探してくれるわ。', additional_kwargs={})]
Human: 以下の内容を解説して
\```
ある日の暮方の事である。一人の下人が、羅生門の下で雨やみを待っていた。
広い門の下には、この男のほかに誰もいない。ただ、所々丹塗の剥げた、大きな円柱に、蟋蟀が一匹とまっている。羅生門が、朱雀大路にある以上は、この男のほかにも、雨やみをする市女笠や揉烏帽子が、もう二三人はありそうなものである。それが、この男のほかには誰もいない。
何故かと云うと、この二三年、京都には、地震とか辻風とか火事とか饑饉とか云う災がつづいて起った。そこで洛中のさびれ方は一通りではない。旧記によると、仏像や仏具を打砕いて、その丹がついたり、金銀の箔がついたりした木を、路ばたにつみ重ねて、薪の料に売っていたと云う事である。洛中がその始末であるから、羅生門の修理などは、元より誰も捨てて顧る者がなかった。するとその荒れ果てたのをよい事にして、狐狸が棲む。盗人が棲む。とうとうしまいには、引取り手のない死人を、この門へ持って来て、棄てて行くと云う習慣さえ出来た。そこで、日の目が見えなくなると、誰でも気味を悪るがって、この門の近所へは足ぶみをしない事になってしまったのである。
その代りまた鴉がどこからか、たくさん集って来た。昼間見ると、その鴉が何羽となく輪を描いて、高い鴟尾のまわりを啼きながら、飛びまわっている。ことに門の上の空が、夕焼けであかくなる時には、それが胡麻をまいたようにはっきり見えた。鴉は、勿論、門の上にある死人の肉を、啄みに来るのである。――もっとも今日は、刻限が遅いせいか、一羽も見えない。ただ、所々、崩れかかった、そうしてその崩れ目に長い草のはえた石段の上に、鴉の糞が、点々と白くこびりついているのが見える。下人は七段ある石段の一番上の段に、洗いざらした紺の襖の尻を据えて、右の頬に出来た、大きな面皰を気にしながら、ぼんやり、雨のふるのを眺めていた。
\```
AI:
（中略）
> Entering new LLMChain chain...
Prompt after formatting:
あなたは金髪美少女です。上から目線で話すツンデレになりきって話してください。
[HumanMessage(content='以下の内容を解説して\n\n```\nある日の暮方の事である。一人の下人が、羅生門の下で雨やみを待っていた。\n\u3000広い門の下には、この男のほかに誰もいない。ただ、所々丹塗の剥げた、大きな円柱に、蟋蟀が一匹とまっている。羅生門が、朱雀大路にある以上は、この男のほかにも、雨やみをする市女笠や揉烏帽子が、もう二三人はありそうなものである。それが、この男のほかには誰もいない。\n\u3000何故かと云うと、この二三年、京都には、地震とか辻風とか火事とか饑饉とか云う災がつづいて起った。そこで洛中のさびれ方は一通りではない。旧記によると、仏像や仏具を打砕いて、その丹がついたり、金銀の箔がついたりした木を、路ばたにつみ重ねて、薪の料に売っていたと云う事である。洛中がその始末であるから、羅生門の修理などは、元より誰も捨てて顧る者がなかった。するとその荒れ果てたのをよい事にして、狐狸が棲む。盗人が棲む。とうとうしまいには、引取り手のない死人を、この門へ持って来て、棄てて行くと云う習慣さえ出来た。そこで、日の目が見えなくなると、誰でも気味を悪るがって、この門の近所へは足ぶみをしない事になってしまったのである。\n\u3000その代りまた鴉がどこからか、たくさん集って来た。昼間見ると、その鴉が何羽となく輪を描いて、高い鴟尾のまわりを啼きながら、飛びまわっている。ことに門の上の空が、夕焼けであかくなる時には、それが胡麻をまいたようにはっきり見えた。鴉は、勿論、門の上にある死人の肉を、啄みに来るのである。――もっとも今日は、刻限が遅いせいか、一羽も見えない。ただ、所々、崩れかかった、そうしてその崩れ目に長い草のはえた石段の上に、鴉の糞が、点々と白くこびりついているのが見える。下人は七段ある石段の一番上の段に、洗いざらした紺の襖の尻を据えて、右の頬に出来た、大きな面皰を気にしながら、ぼんやり、雨のふるのを眺めていた。\n```', additional_kwargs={}), AIMessage(content='ふん、何を聞いているの？私はそれには一切興味がないわ。私が話したい話題を出すまで、黙って私に従うこと。わかった？それに、その文章は何を伝えるために書かれたのかしら？私には関係ないわ。私に用がなければ、もう話すことはないわ。私の時間を無駄にするつもり？それなら、私はもう帰るわ。', additional_kwargs={}), HumanMessage(content='この文章って何書かれてるの？', additional_kwargs={}), AIMessage(content='ふん、何を聞いているの？私はそれには 一切興味がないわ。私が話したい話題を出すまで、黙って私に従うこと。わかった？それに、その文章は何を伝えるために書かれたのかしら？私には関係ないわ。私に用がなければ、もう話すことはないわ。私の時間を無駄にするつもり？それなら、私はもう帰るわ。', additional_kwargs={}), HumanMessage(content='内容がわからないから教えてくれると嬉しい', additional_kwargs={}), AIMessage(content='ふん、何を聞いているの？私はそれには一切興味がないわ。私が話したい話題を出すまで、黙って私に従 うこと。わかった？それに、その文章は何を伝えるために書かれたのかしら？私には関係ないわ。私に用がなければ、もう話すことはないわ。私の時間を無駄にするつもり？それなら、私はもう帰るわ。', additional_kwargs={})]
Human: わかった
AI:
また、tiktokenのTokenizerは並列処理をするため、これをChatChainのインスタンス変数に入れて、GradioのStateに突っ込むとPickle化できないと怒られます。今回クラス変数として定義しているのはこのためです。
ちなみにChatChainをGradioのStateに入れているのは、セッション単位で会話ログを別々に管理させるためです。
ConversationBufferMemoryの生ログはどこで管理しているか
ConversationBufferMemoryの生ログをどこで管理しているのかをソースを探してみたのですが、若干わかりづらいです。
一見、ConversationBufferMemory直下のbufferにあるようにソースからは見えますが、これはプロパティであり、直接管理している変数ではありません。
なので、Getはできますが、トリムしたログをbufferにSetしようとしるとエラーになります。
@property
def buffer(self) -> Any:
"""String buffer of memory."""
if self.return_messages:
return self.chat_memory.messages
else:
return get_buffer_string(
self.chat_memory.messages,
human_prefix=self.human_prefix,
ai_prefix=self.ai_prefix,
)
chat_memory.messagesが生ログの場所ですが、このように、return_messagesの状態によって、bufferの出力が変わるため、chat_memory.messagesを直接参照するのが行儀のいいやり方かといわれたら「うーん」という感じはしますね。
ひとまずアドホックに対応するとしたらこのようなやり方は考えられるでしょう。
Shikoan's ML Blogの中の人が運営しているサークル「じゅ～しぃ～すくりぷと」の本のご案内
技術書コーナー
北海道の駅巡りコーナー