工場系Copilotの作成を目指して
Function Calling可能なSLM（Small Language Model）の動作を確認する

BING_SEARCH_V7_SUBSCRIPTION_KEY = APIキーをここに書く
BING_SEARCH_V7_ENDPOINT = エンドポイントURLをここに書く

accelerate==0.34.2
anyio==4.6.0
argon2-cffi==23.1.0
argon2-cffi-bindings==21.2.0
arrow==1.3.0
asttokens==2.4.1
async-lru==2.0.4
attrs==24.2.0
babel==2.16.0
beautifulsoup4==4.12.3
bitsandbytes==0.44.1
bleach==6.1.0
certifi==2024.8.30
cffi==1.17.1
charset-normalizer==3.3.2
colorama==0.4.6
comm==0.2.2
debugpy==1.8.6
decorator==5.1.1
defusedxml==0.7.1
executing==2.1.0
fastjsonschema==2.20.0
filelock==3.16.1
fqdn==1.5.1
fsspec==2024.9.0
h11==0.14.0
httpcore==1.0.6
httpx==0.27.2
huggingface-hub==0.25.1
idna==3.10
ipykernel==6.29.5
ipython==8.28.0
ipywidgets==8.1.5
isoduration==20.11.0
jedi==0.19.1
Jinja2==3.1.4
json5==0.9.25
jsonpointer==3.0.0
jsonschema==4.23.0
jsonschema-specifications==2023.12.1
jupyter==1.1.1
jupyter-console==6.6.3
jupyter-events==0.10.0
jupyter-lsp==2.2.5
jupyter_client==8.6.3
jupyter_core==5.7.2
jupyter_server==2.14.2
jupyter_server_terminals==0.5.3
jupyterlab==4.2.5
jupyterlab_pygments==0.3.0
jupyterlab_server==2.27.3
jupyterlab_widgets==3.0.13
MarkupSafe==2.1.5
matplotlib-inline==0.1.7
mistune==3.0.2
mpmath==1.3.0
nbclient==0.10.0
nbconvert==7.16.4
nbformat==5.10.4
nest-asyncio==1.6.0
networkx==3.3
notebook==7.2.2
notebook_shim==0.2.4
numpy==2.1.1
overrides==7.7.0
packaging==24.1
pandocfilters==1.5.1
parso==0.8.4
pillow==10.2.0
platformdirs==4.3.6
prometheus_client==0.21.0
prompt_toolkit==3.0.48
psutil==6.0.0
pure_eval==0.2.3
pycparser==2.22
Pygments==2.18.0
python-dateutil==2.9.0.post0
python-dotenv==1.0.1
python-json-logger==2.0.7
pywin32==307
pywinpty==2.0.13
PyYAML==6.0.2
pyzmq==26.2.0
referencing==0.35.1
regex==2024.9.11
requests==2.32.3
rfc3339-validator==0.1.4
rfc3986-validator==0.1.1
rpds-py==0.20.0
safetensors==0.4.5
scipy==1.14.1
Send2Trash==1.8.3
setuptools==75.1.0
six==1.16.0
sniffio==1.3.1
soupsieve==2.6
stack-data==0.6.3
sympy==1.13.3
terminado==0.18.1
tinycss2==1.3.0
tokenizers==0.20.0
torch==2.4.1+cu118
torchaudio==2.4.1+cu118
torchvision==0.19.1+cu118
tornado==6.4.1
tqdm==4.66.5
traitlets==5.14.3
transformers==4.45.1
types-python-dateutil==2.9.0.20241003
typing_extensions==4.12.2
uri-template==1.3.0
urllib3==2.2.3
wcwidth==0.2.13
webcolors==24.8.0
webencodings==0.5.1
websocket-client==1.8.0
widgetsnbextension==4.0.13

import json
import re
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch
import time
import requests
import os
from dotenv import load_dotenv
from bs4 import BeautifulSoup
import argparse
import warnings
from typing import Union, Dict, Any, List

# UserWarningを無視する設定
warnings.filterwarnings("ignore", category=UserWarning)

# .envファイルの読み込み
load_dotenv()

# GPUが利用可能かチェック、利用可能な場合はGPUを利用する
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"\n使用するデバイス: {device}\n")

# モデルの読み込み
model_id = "mzbac/llama-3-8B-Instruct-function-calling-v0.2"
tokenizer = AutoTokenizer.from_pretrained(model_id)

# 環境変数の読み込み
BING_SEARCH_V7_SUBSCRIPTION_KEY = os.getenv('BING_SEARCH_V7_SUBSCRIPTION_KEY')
BING_SEARCH_V7_ENDPOINT = os.getenv('BING_SEARCH_V7_ENDPOINT')

def extract_function_call(response: str) -> Union[Dict[str, Any], None]:
    """関数呼び出しのJSON部分を抽出する関数。

    Args:
        response (str): モデルからの応答文字列。

    Returns:
        dict: 抽出されたJSONオブジェクト。
        None: JSONが見つからなかった場合。

    Raises:
        json.JSONDecodeError: JSONの解析に失敗した場合。
    """
    json_match = re.search(r'(\{"function": .*\{.*\[.*\]\}})', response)
    if json_match:
        json_str = json_match.group(1)
        print("\nFunction Callが実行されました!")
        try:
            json_obj = json.loads(json_str)
            return json_obj
        except json.JSONDecodeError as e:
            print(f"JSONの解析に失敗しました: {e}")
            print(f"エラー位置: {e.pos}")
            print(f"エラー行: {e.lineno}, エラー列: {e.colno}")
            print(f"解析しようとしたJSON文字列: {json_str}")
            formatted_json = json.dumps(json.loads(json_str), indent=2)
            print("整形されたJSON:")
            print(formatted_json)
    else:
        print("\n関数呼び出しのJSONが見つかりませんでした。Function Callが実行されませんでした。")
    return None

def extract_and_join_content(data_list: List[Dict[str, Any]], max_length: int = 1000) -> str:
    """Bing Searchの検索結果のContentを抽出して結合する関数。

    Args:
        data_list (list): 検索結果のリスト。
        max_length (int): 結合する各コンテンツの最大長。デフォルトは1000。

    Returns:
        str: 結合されたコンテンツ文字列。
    """
    content_list = [item['content'][:max_length] for item in data_list if 'content' in item]
    return '\n'.join(content_list)

def search_web(search_terms: List[str], k: int = 3) -> Dict[str, Union[str, List[Dict[str, str]]]]:
    """Azure Bing Search APIを使用してWeb検索を実行し、上位k件の結果のURLの内容を取得します。

    Args:
        search_terms (list): 検索クエリのリスト。
        k (int): 取得する検索結果の数。デフォルトは3。

    Returns:
        dict: 検索結果とそのURL先の内容を含む辞書。
        dict: エラーが発生した場合のエラーメッセージを含む辞書。
    """
    subscription_key = BING_SEARCH_V7_SUBSCRIPTION_KEY
    endpoint = BING_SEARCH_V7_ENDPOINT + "/v7.0/search"
    query = " ".join(search_terms)
    headers = {"Ocp-Apim-Subscription-Key": subscription_key}
    params = {"q": query, "textDecorations": True, "textFormat": "HTML"}
    try:
        response = requests.get(endpoint, headers=headers, params=params)
        response.raise_for_status()
        search_results = response.json()
        if "webPages" in search_results and "value" in search_results["webPages"]:
            results = []
            for item in search_results["webPages"]["value"][:k]:
                url = item["url"]
                content = get_url_content(url)
                results.append({"name": item["name"], "url": url, "content": content})
            return {"result": results}
        else:
            return {"result": "検索結果が見つかりませんでした。"}
    except Exception as ex:
        return {"error": str(ex)}

def get_url_content(url: str) -> str:
    """指定されたURLの内容を取得します。

    Args:
        url (str): 取得するURL。

    Returns:
        str: URLの内容。
        str: エラーメッセージ（エラーが発生した場合）。
    """
    try:
        response = requests.get(url)
        response.raise_for_status()
        soup = BeautifulSoup(response.content, 'html.parser')
        for script in soup(["script", "style"]):
            script.decompose()
        text = soup.get_text()
        lines = (line.strip() for line in text.splitlines())
        chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
        text = ' '.join(chunk for chunk in chunks if chunk)
        return text
    except Exception as ex:
        return f"Error fetching content: {str(ex)}"

def extract_final_output(decoded_response: str) -> str:
    """生成された応答から最後のパラグラフを抽出し、<|eot_id|>を除去する関数。

    Args:
        decoded_response (str): モデルからのデコードされた応答。

    Returns:
        str: 最後のパラグラフを抽出し、<|eot_id|>を除去した文字列。
    """
    return decoded_response.strip().split('\n\n')[-1].replace('<|eot_id|>', '').strip()

def main(model_id: str, tokenizer: AutoTokenizer, prompt: str, max_new_tokens: int = 5000, temperature: float = 0.7) 
    """メイン関数。モデルを使用してプロンプトに応答します。

    Args:
        model_id (str): 使用するモデルのID。
        tokenizer (AutoTokenizer): トークナイザーオブジェクト。
        prompt (str): ユーザーからのプロンプト。
        max_new_tokens (int): 生成するトークンの最大数。デフォルトは5000。
        temperature (float): 生成時の温度。デフォルトは0.7。

    Returns:
        str: モデルからの応答。
    """
    start_time = time.time()
    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        llm_int8_enable_fp32_cpu_offload=True,
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_use_double_quant=True
    )
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=torch.float16,
        device_map="auto",
        quantization_config=quantization_config
    )
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token_id = tokenizer.eos_token_id

    tool = {
        "name": "search_web",
        "description": "Perform a web search for a given search terms.",
        "function": search_web,
        "parameter": {
            "type": "object",
            "properties": {
                "search_terms": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "The search queries for which the search is performed.",
                    "required": True,
                }
            }
        },
    }

    messages = [
        {
            "role": "system",
            "content": f"""You are a helpful assistant with access to the following functions. Use them if required -
            Please make sure to answer in Japanese. When you need to use a function, output the function call in the 
            {{"function": "function_name", "arguments": {{"arg1": "value1", "arg2": "value2"}}}}.Also output the reas
            Think carefully about whether the function should be used or not.
            Also, if something cannot be answered either with prior knowledge or with a function, answer honestly so.
        },
        {"role": "user", "content": prompt}
    ]

    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
    attention_mask = inputs.ne(tokenizer.pad_token_id).long()
    inputs = inputs.to(device)
    attention_mask = attention_mask.to(device)
    terminators = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|eot_id|>")]

    generate_start_time = time.time()
    with torch.no_grad():
        outputs = model.generate(
            inputs,
            attention_mask=attention_mask,
            max_new_tokens=max_new_tokens,
            eos_token_id=terminators,
            do_sample=True,
            temperature=temperature,
            pad_token_id=tokenizer.pad_token_id,
        )

    generate_end_time = time.time()
    response = outputs[0]
    decoded_response = tokenizer.decode(response)

    # Function Callの処理　⇒　Function Callが実行されない時は動かない
    function_call = extract_function_call(decoded_response)
    if function_call:
        if function_call.get("function") == "search_web":
            search_terms = function_call['arguments']['search_terms']
            print("検索に使用されるキーワード:", search_terms)
            if isinstance(search_terms, str):
                search_terms = [search_terms]
            search_results = extract_and_join_content(search_web(search_terms)["result"])
            new_prompt = f"""{prompt}\n以下の検索結果を基に回答してください。\n\n#検索結果:\n{search_results}\n\n
                        #制約\n- 回答は日本語で行ってください。\n- 極力質問に対して簡潔に回答してください。不要な
                        情報を含めないようにしてください。\n- 検索結果をそのままコピー&ペーストせず、自分の言葉で
                        分かり易く回答してください。"""
            messages.append({"role": "user", "content": new_prompt})
            inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
            attention_mask = inputs.ne(tokenizer.pad_token_id).long()
            inputs = inputs.to(device)
            attention_mask = attention_mask.to(device)
            with torch.no_grad():
                outputs = model.generate(
                    inputs,
                    attention_mask=attention_mask,
                    max_new_tokens=max_new_tokens,
                    eos_token_id=terminators,
                    do_sample=True,
                    temperature=temperature,
                    pad_token_id=tokenizer.pad_token_id,
                )
            response = outputs[0]
            decoded_response = tokenizer.decode(response)

    end_time = time.time()
    total_time = end_time - start_time
    generate_time = generate_end_time - generate_start_time

    print(f"トータル実行時間: {total_time:.2f} seconds")
    print(f"回答生成に掛かった時間: {generate_time:.2f} seconds")

    # 生成結果を抽出
    final_output = extract_final_output(decoded_response)
    return final_output

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Run the AI model with a custom prompt.')
    parser.add_argument('prompt', type=str, help='The prompt to use for the AI model')
    args = parser.parse_args()
    prompt = args.prompt
    response = main(model_id, tokenizer, prompt, max_new_tokens=5000, temperature=0.7)
    print("\n質問：\n", prompt)
    print("\n回答：\n", response)