FunctionGemmaに検索ワードを生成させる

概要

Function Callingができる軽いモデルを動かしてみる。
GPUなしでも動くFunctionGemmaで検索ワード生成を試す。
FunctionGemmaは検索ワード生成としてはいまいちという結果。

FunctionGemmaのダウンロード

google版のモデルはサインインが必要なので、unsloth版をPCにダウンロードして使用する。

今回は、ブラウザで全ファイルダウンロードして、functiongemma270itとして保存する。

ソースコード

from transformers import AutoProcessor, AutoModelForCausalLM

model_id = "functiongemma270it" # ご自身の環境のモデルパス
processor = AutoProcessor.from_pretrained(
    model_id,
    device_map="cpu"
)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype="auto",
    device_map="cpu"
)

search_function_schema = {
    "type": "function",
    "function": {
        "name": "get_search_query",
        "description": "ユーザーの要望を分析し、検索エンジンでヒットしやすい最適なキーワードを生成して検索します。",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "検索クエリ。話し言葉は検索用キーワードに変換してください（例: 'いい感じの' -> 'おいしい' や 'おすすめ' など）。スペース区切りで指定します。",
                },
            },
            "required": ["query"],
        },
    }
}

message = [
    {
        "role": "developer",
        "content": "あなたは優秀な検索アシスタントです。ユーザーの要望に対して、get_search_query を使って検索を実行してください。クエリを作成する際は、「いい感じの」といった曖昧な話し言葉を、「おいしい」「おすすめ」「人気」といった検索にヒットしやすい具体的なキーワードに変換してください。"
    },
    {
        "role": "user", 
        "content": "いい感じのお寿司屋さんを探したい"
    }
]

inputs = processor.apply_chat_template(message, tools=[search_function_schema], add_generation_prompt=True, return_dict=True, return_tensors="pt")

out = model.generate(**inputs.to(model.device), pad_token_id=processor.eos_token_id, max_new_tokens=128)
output = processor.decode(out[0][len(inputs["input_ids"][0]):], skip_special_tokens=True)

print(output)

実行結果

"いい感じのお寿司屋さんを探したい"というチャット入力した。
理想は、"おいしい すし屋"という検索ワード。developerやsearch_function_schemaにもしつこく書いている。
結果は、"いい感じのお寿司屋さん"という検索ワードが生成された。使えない。

所感

FunctionGemmaはチャット入力の気の利いた言い換えとかできない感じ。
キーワードの抽出とかはできそうだけど。