基于01-ai/Yi模型实现函数调用功能的实践指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00883/article/details/148442996

基于01-ai/Yi模型实现函数调用功能的实践指南

Yi 项目地址: https://gitcode.com/gh_mirrors/yi/Yi

引言

函数调用(Function Calling)是大语言模型(LLM)应用开发中的一项重要能力，它允许模型在理解用户请求后，决定是否需要调用外部函数来完成特定任务。本文将详细介绍如何使用01-ai/Yi-1.5-9B-Chat模型实现一个完整的函数调用流程，包括模型加载、函数定义、调用解析和执行等关键环节。

环境准备

在开始之前，我们需要确保开发环境满足以下要求：

Python 3.8或更高版本
足够的GPU显存（建议至少16GB以运行9B参数的Yi模型）
安装必要的Python库：

pip install transformers torch

核心实现步骤

1. 定义可调用函数

首先，我们需要明确模型可以调用哪些函数。在本示例中，我们定义了三个基础数学运算函数：

def multiply(a: int, b: int) -> int:
    return a * b

def plus(a: int, b: int) -> int:
    return a + b

def minus(a: int, b: int) -> int:
    return a - b

available_functions = {
    "multiply": multiply,
    "plus": plus,
    "minus": minus
}

扩展建议：在实际应用中，你可以根据需要添加更多功能函数，如数据库查询、API调用等。每个函数都应明确定义参数类型和返回值类型，这有助于模型正确理解和使用这些函数。

2. 加载Yi模型

使用Hugging Face的transformers库加载Yi-1.5-9B-Chat模型：

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_path = "01-ai/Yi-1.5-9B-Chat"
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
model = AutoModelForCausalLM.from_pretrained(
    model_path, 
    torch_dtype=torch.float16, 
    device_map="auto"
)

注意事项：

使用torch.float16可以减少显存占用
device_map="auto"会自动选择可用的设备（GPU或CPU）
如果显存不足，可以考虑使用量化版本或更小的模型

3. 实现核心功能函数

我们需要实现三个核心功能：

def generate_response(prompt):
    """生成模型响应"""
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, temperature=0.7, top_p=0.95)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response.split("Human:")[0].strip()

def parse_function_call(response):
    """解析模型响应中的函数调用信息"""
    try:
        start = response.index("{")
        end = response.rindex("}") + 1
        function_call_json = response[start:end]
        return json.loads(function_call_json)
    except (ValueError, json.JSONDecodeError):
        return None

def execute_function(function_name: str, arguments: dict):
    """执行指定的函数"""
    if function_name in available_functions:
        return available_functions[function_name](**arguments)
    raise ValueError(f"Function {function_name} not found")

4. 构建对话系统

完整的对话系统需要处理用户输入、生成响应、执行函数调用并维护对话历史：

def main():
    system_prompt = """You are an AI assistant capable of calling functions..."""  # 系统提示词
    
    conversation_history = [f"System: {system_prompt}"]
    
    while True:
        user_input = input("Human: ")
        if user_input.lower() == 'exit':
            break
            
        conversation_history.append(f"Human: {user_input}")
        full_prompt = "\n".join(conversation_history) + "\nAssistant:"
        
        response = generate_response(full_prompt)
        print(f"Model response: {response}")
        
        function_call = parse_function_call(response)
        if function_call:
            try:
                result = execute_function(function_call["function"], function_call["arguments"])
                print(f"Result: {result}")
                conversation_history.append(f"Assistant: The result is {result}")
            except Exception as e:
                print(f"Error: {str(e)}")
        else:
            conversation_history.append(f"Assistant: {response}")

系统提示词设计要点：

明确告知模型可用的函数及其功能
提供函数调用的格式示例
说明何时需要调用函数，何时直接回答

实际应用示例

下面是一个完整的交互示例：

Human: What is 5 plus 3?
Model response: Here's the function call to perform the addition:
{"function": "plus", "arguments": {"a": 5, "b": 3}}
Assistant: The result of plus({'a': 5, 'b': 3}) is 8