[db:标题]
摘要:大语言模型-4.LLM基础能力实现 书生浦语大模型实战营学习笔记-2.LLM基础能力实现 本文包括第二期实战营的第2课内容。本来是想在笔记中给官方教程做做补充的,没想到官方教程的质量还是相当高的,跟着一步一步做没啥坑。所以这篇笔记主要学习一
大语言模型-4.LLM基础能力实现
书生浦语大模型实战营学习笔记-2.LLM基础能力实现
本文包括第二期实战营的第2课内容。本来是想在笔记中给官方教程做做补充的,没想到官方教程的质量还是相当高的,跟着一步一步做没啥坑。所以这篇笔记主要学习一下官方Demo中的一些代码等细节内容。
本文标题中大语言模型系列博客是笔者在学习大语言模型时做的博客;书生浦语大模型实战营学习笔记是笔者在参加书生浦语大模型第二期实战营做的学习笔记。
大语言模型的对话能力实现:以InternLM2-Chat-1.8B的官方Demo为例
我们来看看官方Demo的代码。首先它导入了一些必要的库,并创建变量存储了模型位置:
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
model_name_or_path = "/root/models/Shanghai_AI_Laboratory/internlm2-chat-1_8b"
然后使用了AutoTokenizer 和AutoModelForCausalLM。Hugging Face 的 AutoTokenizer 和 AutoModelForCausalLM 类熟悉大模型的不会陌生,用于自动加载预训练模型和相应的tokenizer。
加载模型时设置了相信远端代码以便从HuggingFace拉取确实模型权重,使用bf16量化节省内存,指定使用第一张显卡。同时使用model.eval()来取消梯度计算。
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True, device_map='cuda:0')
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map='cuda:0')
model = model.eval()
下面就是核心业务了设置system_prompt、接收input、调用model.stream_chat():
system_prompt = """You are an AI assistant whose name is InternLM (书生·浦语).
- InternLM (书生·浦语) is a conversational language model that is developed by Shanghai AI Laboratory (上海人工智能实验室). It is designed to be helpful, honest, and harmless.
- InternLM (书生·浦语) can understand and communicate fluently in the language chosen by the user such as English and 中文.
"""
messages = [(system_prompt, '')]
print("=============Welcome to InternLM chatbot, type 'exit' to exit.=============")
while True:
input_text = input("\nUser >>> ")
input_text = input_text.replace(' ', '') # 移除用户输入文本中的空格
if input_text == "exit": # 如果要退出,输入exit即可
break
length = 0
# 对模型的 stream_chat 方法进行迭代,该方法会生成一个对话的生成器。迭代过程中,每次生成一个回复消息 response 和一个占位符 _。
for response, _ in model.stream_chat(tokenizer, input_text, messages):
# 如果回复消息不为空,则打印回复消息中从上次打印位置 length 开始到结尾的部分,并刷新输出缓冲区。
