01-ai/Yi模型使用Transformers进行推理的完整指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00868/article/details/148443067

01-ai/Yi模型使用Transformers进行推理的完整指南

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

引言

在当今快速发展的人工智能领域，大型语言模型(LLM)已成为自然语言处理任务的核心工具。01-ai/Yi系列模型作为开源社区的重要贡献，提供了强大的语言理解和生成能力。本文将详细介绍如何使用Hugging Face Transformers库来加载和运行Yi-1.5-6B-Chat模型，帮助开发者快速上手这一先进技术。

环境准备

在开始之前，我们需要搭建适当的工作环境。以下是必需的软件包及其作用：

Transformers库：Hugging Face提供的核心库，支持加载和运行各种预训练模型
PyTorch：深度学习框架，提供模型运行的基础设施
Accelerate：优化模型加载和推理过程的工具
SentencePiece：处理模型分词的特殊组件

安装命令如下：

pip install transformers>=4.36.2
pip install gradio>=4.13.0
pip install torch>=2.0.1,<=2.3.0
pip install accelerate
pip install sentencepiece

模型加载详解

成功安装依赖后，我们可以开始加载Yi-1.5-6B-Chat模型。这个过程涉及两个关键组件：分词器(Tokenizer)和模型本身。

from transformers import AutoModelForCausalLM, AutoTokenizer

model_path = '01-ai/Yi-1.5-6B-Chat'

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)

# 加载模型
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map="auto",  # 自动选择可用设备(CPU/GPU)
    torch_dtype='auto'  # 自动选择合适的数据类型
).eval()  # 设置为评估模式

关键参数解析

device_map="auto"：自动检测并利用可用的计算资源，优先使用GPU
torch_dtype='auto'：根据硬件能力自动选择最佳数值精度
.eval()：将模型设置为推理模式，关闭训练特有的功能如Dropout

注意：首次加载模型时，系统会下载约12GB的模型文件，请确保有足够的磁盘空间和稳定的网络连接。

基础推理流程

让我们从一个简单的对话示例开始，了解模型的基本工作流程：

messages = [
    {"role": "user", "content": "你好！"}
]

# 将对话转换为模型可理解的格式
input_ids = tokenizer.apply_chat_template(
    conversation=messages,
    tokenize=True,
    add_generation_prompt=True,
    return_tensors='pt'
)

# 生成响应
output_ids = model.generate(input_ids.to('cuda'))

# 解码输出
response = tokenizer.decode(output_ids[0][input_ids.shape[1]:], skip_special_tokens=True)

print(f"用户: 你好！")
print(f"Yi: {response}")

流程分解

对话格式化：将自然语言对话转换为结构化表示
模板应用：使用apply_chat_template方法准备模型输入
生成响应：调用模型的generate方法产生文本
结果解码：将生成的token ID转换回可读文本

进阶：多轮对话实现

实际应用中，我们通常需要支持多轮对话。以下是一个完整的对话管理实现：

def chat_with_yi(user_input, history=[]):
    # 添加用户输入到历史记录
    history.append({"role": "user", "content": user_input})
    
    # 准备模型输入
    input_ids = tokenizer.apply_chat_template(
        conversation=history,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors='pt'
    )
    
    # 生成响应（限制最大token数）
    output_ids = model.generate(
        input_ids.to('cuda'),
        max_new_tokens=100
    )
    
    # 解码响应
    response = tokenizer.decode(
        output_ids[0][input_ids.shape[1]:],
        skip_special_tokens=True
    )
    
    # 更新对话历史
    history.append({"role": "assistant", "content": response})
    
    return response, history

# 测试对话
history = []
user_inputs = ["你好！", "你能告诉我一个笑话吗？", "谢谢，再见！"]

for user_input in user_inputs:
    print(f"用户: {user_input}")
    response, history = chat_with_yi(user_input, history)
    print(f"Yi: {response}\n")