基于TigerBot-13b训练其函数调用能力

写在前面

原生的tigerbot似乎并不支持函数调用，于是我来支持一下

 

数据集

我在huggingface上找了个英文的数据集

https://huggingface.co/datasets/sadmoseby/sample-function-call

这里面包含了1k组的函数调用，这个数据集的特点如下:

1. 包含有单个/多个/没有函数调用的情形

2. 描述函数的json_schema与OpenAI格式的一致（但多函数情况下，并没有用列表框起来）

3. 数据虽然是多轮对话的数据，但是每一个都是一整条的数据，且每个的开头与tigerbot的头不太一致

 

数据转换

我写了一个数据转换的代码，具体任务如下:

1. 将多个函数时没有用列表格式框选的情况给修复了

2. 切分为了多轮的对话，有多条训练数据

3. 修改了开头的情况

代码如下

import re

import json
import re

# system_prompt中可能有多个函数，多个函数的话要转为标准的[]格式
def get_function_json(input_string):
    # 使用正则表达式分割字符串，找出独立的 JSON 字符串
    json_strings = re.findall(r'\{[\s\S]+?\}\s*(?=\{|$)', input_string)

    # 解析每个 JSON 字符串并把它们加入到列表中
    json_objects = []
    for json_str in json_strings:
        input_string = input_string.replace(json_str, '')
        try:
            json_obj = json.loads(json_str)
            json_objects.append(json_obj)
        except json.JSONDecodeError as e:
            print(f"Error decoding JSON: {e}")
    # 打印结果或进行其他操作
    if json_objects:
        return input_string + json.dumps(json_objects, ensure_ascii=False, indent=4)
    else:
        return input_string

# 切分读入的数据
def split_string_with_keywords(s, keywords):
    # 将关键词列表转化为正则表达式，使用括号捕获分隔符
    # 比如 ['system', 'assistant'] 会被转换成 (system)|(assistant)
    regex_pattern = '({})'.format('|'.join(map(re.escape, keywords)))

    # 使用 re.split，它会返回包含分隔符的列表
    parts = re.split(regex_pattern, s)

    # 初始化结果列表
    result = []

    # 存储上一个匹配到的关键词，初始时没有关键词
    last_keyword = None

    # 遍历分割后的列表
    for part in parts:
        # 如果当前部分是关键词，记录下来并继续下一轮循环
        if part in keywords:
            last_keyword = part
            continue
        # 如果当前部分不是关键词，且上一部分是关键词，则将其作为结果加入
        if last_keyword:
            result.append((last_keyword, part.strip()))
            last_keyword = None  # 重置关键词

    return result

max_len = 0


def count_words_and_punctuation(s):
    # 使用正则表达式来匹配单词和标点符号
    # \w+ 匹配单词字符（字母、数字、下划线）出现一次或多次组成的单词
    # | 表示或，用来分隔不同的匹配规则
    # \s 表示空白字符
    # [^\w\s] 匹配任意不是单词字符和不是空白字符的字符，即标点符号
    matches = re.findall(r'\w+|[^\w\s]', s)

    # 计算匹配项的数量，即单词和标点符号的总数
    return len(matches)

def solve(input):
    global max_len
    max_len = max(max_len , count_words_and_punctuation(input))
    import json
    # 基础替换
    input = input.replace('<|endoftext|>', '')

    replace_map = {
        'SYSTEM:' : '\n\n### System:\n ',
        'ASSISTANT:': '\n\n### Response:\n ',
        'USER:': '\n\n### Instruction:\n ',
        'FUNCTION RESPONSE:': '\n\n### Function:\n '
    }

    data = split_string_with_keywords(input, list(replace_map.keys()))

    # 更换函数的格式
    if data[0][0] == 'SYSTEM:':
        data[0] = (data[0][0], get_function_json(data[0][1]))

    return_data = []
    train_str = ''
    for element in data:
        train_str += replace_map[element[0]]
        if element[0] == 'ASSISTANT:':
            return_data.append({
                "instruction": train_str,
                "input": "",
                "output": element[1]
            })
        train_str += element[1]

    return return_data

import pandas as pd

train_data = []

# 读取Parquet文件
df = pd.read_parquet('train-00000-of-00001.parquet')
column_name = df.columns[0]
for value in df[column_name]:
    train_data += solve(value)

with open('train_function_call.json', 'w', encoding='utf-8') as f:
    json.dump(train_data, f, ensure_ascii=False, indent=4)
print(max_len)

 改好格式的数据如下（以response来切分，response前为输入，response为需要模型生成的输出）：

### System:
 You are a helpful assistant with access to the following functions. Use them if required -
[
    {
        "name": "search_books",
        "description": "Search for books based on specified criteria",
        "parameters": {
            "type": "object",
            "properties": {
                "title": {
                    "type": "string",
                    "description": "The title of the book"
                },
                "author": {
                    "type": "string",
                    "description": "The author of the book"
                },
                "genre": {
                    "type": "string",
                    "description": "The genre of the book"
                }
            },
            "required": [
                "title"
            ]
        }
    }
]

### Instruction:
 I am looking for a book but I can't remember the full title. I know it has the word "Sun" in it and it's a science fiction novel.

### Response:
 <functioncall> {"name": "search_books", "arguments": '{"title": "Sun", "genre": "science fiction"}'}

### Function:
 {"results": [{"title": "The Sun Also Rises in Space", "author": "John Doe", "genre": "science fiction"}, {"title": "Sunset on Mars", "author": "Jane Doe", "genre": "science fiction"}]}

### Response:
 I found two science fiction books with "Sun" in the title. The first one is "The Sun Also Rises in Space" by John Doe and the second one is "Sunset on Mars" by Jane Doe.

 

 

启动训练

笔者依然在恒源云上，基于tigerbot-13b-chat-v5-4k进行训练。

考虑到vllm暂时不支持PEFT格式的adapter，此次依然采用了freeze训练。

为了尽可能地训练更多的层，笔者采用了单个A100-80G的显卡，这样可以在seq_len达到3072的情况下，训练10层的tranformer参数。

注意，此次的template和以前不太一样（因为有各种的function和自己添加的system），所以添加了一个新的模板

register_template(
    name="null",
    prefix=[
        ""
    ],
    prompt=[
        "{{query}}"
    ],
    system="",
    sep=[]
)

训练命令如下

python src/train_bash.py \
    --stage sft \
    --model_name_or_path /hy-tmp/tigerbot-13b-chat-v5-4k \
    --do_train True \
    --finetuning_type freeze \
    --num_layer_trainable 10 \
    --template null \
    --dataset_dir data \
    --dataset train_function_call \
    --cutoff_len 3072 \
    --learning_rate 1e-4 \
    --num_train_epochs 1.0 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 2 \
    --logging_steps 1 \
    --save_steps 10000 \
    --output_dir /hy-tmp/tigerbot-13b-function-call \
    --fp16 True \
    --plot_loss True \
    --overwrite_output_dir