Needle 使い方 入門｜26Mの超軽量モデルで爆速ツール呼び出しを実現する方法

# 仮想環境の作成（推奨）
python -m venv needle-env
source needle-env/bin/activate  # Windowsの場合は needle-env\Scripts\activate

# 必要なライブラリのインストール
pip install torch transformers accelerate

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# モデルIDの指定
model_id = "khulnasoft/needle-26m"

# トークナイザーとモデルのロード
# device_map="auto" を指定することで、GPUがあれば自動でGPUを使用します
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16, # 精度を保ちつつメモリを節約
    device_map="auto"
)

# ツール（関数）の定義をJSON形式で用意
# 26Mモデルなので、説明は簡潔にするのがコツです
tools = [
    {
        "name": "get_weather",
        "description": "指定された場所の天気を取得する",
        "parameters": {
            "type": "object",
            "properties": {
                "location": {"type": "string", "description": "都市名（例：東京）"}
            },
            "required": ["location"]
        }
    },
    {
        "name": "send_email",
        "description": "メールを送信する",
        "parameters": {
            "type": "object",
            "properties": {
                "to": {"type": "string", "description": "宛先のアドレス"},
                "body": {"type": "string", "description": "本文"}
            },
            "required": ["to", "body"]
        }
    }
]

def ask_needle(user_query):
    # Needle専用のプロンプト構築（システムプロンプト的な役割）
    prompt = f"Tools: {tools}\nUser: {user_query}\nAssistant: <call>"

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    # 推論実行。max_new_tokensは短くてOK（ツール名と引数だけなので）
    with torch.no_grad():
        output_tokens = model.generate(
            **inputs,
            max_new_tokens=128,
            temperature=0.1, # 決定論的な出力を得るために低めに設定
            do_sample=False
        )

    # 出力のデコード
    result = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
    return result

# テスト実行
query = "東京の天気を教えて"
print(f"Query: {query}")
print(f"Response: {ask_needle(query)}")

query = "ねぎさんに『お疲れ様です』とメールして"
print(f"\nQuery: {query}")
print(f"Response: {ask_needle(query)}")

Query: 東京の天気を教えて
Response: ... <call>get_weather(location='東京')

Query: ねぎさんに『お疲れ様です』とメールして
Response: ... <call>send_email(to='ねぎ', body='お疲れ様です')

import ast
import re

def parse_needle_output(output):
    # <call>以降の関数呼び出し部分を抽出する正規表現
    match = re.search(r"<call>(\w+)\((.*)\)", output)
    if not match:
        return None, None

    func_name = match.group(1)
    args_str = match.group(2)

    # 安全に引数を評価するためにast.literal_evalを使用
    # 注意: 複雑な引数の場合はパース処理を強化する必要があります
    try:
        # 引数部分を辞書形式に変換（簡易実装）
        # 例: location='東京' -> {"location": "東京"}
        args = {}
        for pair in args_str.split(','):
            key, value = pair.split('=')
            args[key.strip()] = ast.literal_eval(value.strip())
        return func_name, args
    except Exception as e:
        print(f"Parsing error: {e}")
        return func_name, None

# 実行例
raw_output = ask_needle("札幌の天気を調べて")
func, params = parse_needle_output(raw_output)

if func == "get_weather":
    print(f"【実行】{func}関数を、引数{params}で呼び出します。")
    # ここで実際のAPI（OpenWeatherMap等）を叩く処理を入れる