Transformers Agent用のカスタムツールの作成

2023年5月28日 23:34

個人的に割とイケてるという感想のHuggingfaceのTransformers Agentのカスタムツールの作成を試してみます。
pineconeが公開してくれている例を参考にさせてもらいました。

Huggingface の Transformers Agentについては以下をご参照ください。
このAgentは、あまり注目されていない印象がありますが、実行時にエージェントがpythonの簡易コードを書いてから、具体的な処理に移ることや、外部ツール、プラグインとの連携が容易なことから、用途によってはlangchain 標準のzeroshot系のエージェントよりも有効な場面も多くありそうです。

０．前準備

!pip install langchain > /dev/null
!pip install openai > /dev/null
!pip install transformers > /dev/null

 from transformers.tools import OpenAiAgent
 pswd = 'OpenAI API key'
 agent = OpenAiAgent(model="gpt-3.5-turbo", api_key=pswd)

１．Transformers Agentが参照するツール

transformers agentが参照するツールは toolboxに設定されています。
初期状態で読み込まれるツールを表示します。

from transformers.tools.agents import PreTool

for i, (name, tool) in enumerate(agent.toolbox.items()):
    if type(tool) is PreTool:
        print(f"{i+1}: {name} (PreTool)")
    else:
        print(f"{i+1}: {name} (Tool)")

1: document_qa (PreTool)
2: image_captioner (PreTool)
3: image_qa (PreTool)
4: image_segmenter (PreTool)
5: transcriber (PreTool)
6: summarizer (PreTool)
7: text_classifier (PreTool)
8: text_qa (PreTool)
9: text_reader (PreTool)
10: translator (PreTool)
11: image_transformer (PreTool)
12: text_downloader (PreTool)
13: image_generator (PreTool)
14: video_generator (PreTool)

これらのツールは、agent呼び出し時に、各々どのようなツールの説明文(description)などがプロンプト追加されるので、入力トークンが増大したり、ツールの数が多すぎるとagentがどのツールを使うか迷ったり、誤動作してしまう可能性もありそうです。
このため、目的に合わせて適宜、整理するのがよさそうです。

２．toolboxを空にする

今回は簡単のためにtoolboxのPreToolを全部空にします。

# 削除リスト
del_list = []

for name, tool in agent.toolbox.items():
    if type(tool) is PreTool:
        del_list.append(name)

for name in del_list:
    del agent.toolbox[name]

すっきり空になりました。

agent.toolbox

{}

３．カスタムツールの定義

from transformers import Tool

class MeaningOfLifeTool(Tool):
    name = 'meaning_of_life'
    description = (
    """このツールは、万物の本質の答えを見つけるために使うものです、
    このツールは、人生、宇宙、すべての答えを見つけるために使用されます。
    ユーザーの入力'query'を受けて、彼らが探している究極の答えを返します。"""
    )
    inputs = ['text']
    outputs = ['text']

    def __call__(self, query: str):
        return "42"

meaning_of_life_tool = MeaningOfLifeTool()

agent = OpenAiAgent(
    model='gpt-3.5-turbo',
    api_key=pswd,
    additional_tools=[meaning_of_life_tool]
)

４．実行！

agent.run("人生、宇宙、すべての答えは？")

==Explanation from the agent==
I will use the following tool: `meaning_of_life` to find the answer to the ultimate question.

==Code generated by the agent==
answer = meaning_of_life(query="What is the answer to the ultimate question of life, the universe, and everything?")
print(f"The answer is {answer}.") # The answer is 42.

==Result==
The answer is 42.
'42'

agent.run("明日、月曜日に会社に行く意味は？")

==Result==
The ultimate answer is 42.
'42'

orz

元ネタを読んだことある人ってどの位いるんだろう？

おしまい。お読みいただき、ありがとうございました。