智能编程助手代码解读
# 导入所需的库
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import streamlit as st
这段代码是Python语言编写的,用于导入一些在机器学习和深度学习中常用的库。下面是对每行代码的解释:
-
from transformers import AutoTokenizer, AutoModelForCausalLM
:- 这行代码从
transformers
库中导入了两个类:AutoTokenizer
和AutoModelForCausalLM
。 AutoTokenizer
是用于自动处理文本数据的分词器,它可以根据不同的预训练模型自动选择相应的分词方法。AutoModelForCausalLM
是用于因果语言模型(Causal Language Modeling,即生成文本的任务)的模型类。
- 这行代码从
-
import torch
:- 这行代码导入了PyTorch库,PyTorch是一个开源的机器学习库,广泛用于深度学习的研究和开发。它提供了强大的GPU加速的张量计算能力,以及构建深度学习模型的动态计算图。
-
import streamlit as st
:- 这行代码导入了Streamlit库,Streamlit是一个用于快速构建数据应用的开源库,它可以将Python脚本转换为交互式Web应用。
as st
是给库指定了一个别名st
,这样在代码中就可以通过st
来访问Streamlit库的所有功能。
- 这行代码导入了Streamlit库,Streamlit是一个用于快速构建数据应用的开源库,它可以将Python脚本转换为交互式Web应用。
# 创建一个标题和一个副标题
st.title("💬 Yuan2.0 智能编程助手")
在Streamlit应用中,st.title函数通常用来设置应用的主标题,它在页面的顶部显示,并且通常是用户加载页面后看到的第一个元素。这个标题通过一个表情符号和文字结合的方式,提供了一个友好而直观的欢迎信息,同时快速传达了应用的功能和用途。
# 源大模型下载
# 源大模型下载
from modelscope import snapshot_download
model_dir = snapshot_download('IEITYuan/Yuan2-2B-Mars-hf', cache_dir='./')
# 定义模型路径
path = './IEITYuan/Yuan2-2B-Mars-hf'
# 定义模型数据类型
torch_dtype = torch.bfloat16 # A10
# torch_dtype = torch.float16 # P100
这行代码将torch_dtype变量设置为torch.bfloat16,表示使用16位的二进制浮点数(bfloat16)。这种数据类型提供了与32位浮点数相似的动态范围,但只使用了一半的位数,可以在某些硬件上提供更好的性能和内存效率。注释# A10可能是指这种数据类型特别适合于NVIDIA A100系列GPU
# 定义一个函数,用于获取模型和tokenizer
@st.cache_resource
def get_model():
print("Creat tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(path, add_eos_token=False, add_bos_token=False, eos_token='<eod>')
tokenizer.add_tokens(['<sep>', '<pad>', '<mask>', '<predict>', '<FIM_SUFFIX>', '<FIM_PREFIX>', '<FIM_MIDDLE>','<commit_before>','<commit_msg>','<commit_after>','<jupyter_start>','<jupyter_text>','<jupyter_code>','<jupyter_output>','<empty_output>'], special_tokens=True)
print("Creat model...")
model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch_dtype, trust_remote_code=True).cuda()
return tokenizer, model
# 加载model和tokenizer
tokenizer, model = get_model()
这段代码定义了一个用于获取预训练模型和分词器(tokenizer)的函数,并使用Streamlit的st.cache_resource
装饰器来缓存函数的结果。下面是对代码的详细解释:
-
@st.cache_resource
:- 这是一个Streamlit提供的装饰器,用于缓存函数的执行结果。当函数再次被调用时,如果输入参数没有变化,Streamlit会使用缓存的结果而不是重新执行函数。这可以显著提高应用的性能,尤其是在加载大型模型或执行耗时操作时。
-
函数定义
def get_model()
:- 定义了一个名为
get_model
的函数,用于加载和返回分词器和模型。
- 定义了一个名为
-
tokenizer = AutoTokenizer.from_pretrained(path, add_eos_token=False, add_bos_token=False, eos_token='<eod>')
:- 使用
AutoTokenizer.from_pretrained
方法从指定路径path
加载预训练的分词器。 add_eos_token=False
和add_bos_token=False
表示在加载分词器时不自动添加结束符(EOS)和开始符(BOS)。eos_token='<eod>'
设置了分词器的结束符为'<eod>'
。
- 使用
-
tokenizer.add_tokens(['<sep>', '<pad>', ..., '<empty_output>'], special_tokens=True)
:- 向分词器中添加一系列特殊的标记,这些标记通常用于处理特定的NLP任务,如序列到序列的模型。
special_tokens=True
表示这些标记将作为特殊的分词器标记处理,它们将被加入到分词器的词汇表中。
-
model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch_dtype, trust_remote_code=True).cuda()
:- 使用
AutoModelForCausalLM.from_pretrained
方法从指定路径path
加载预训练的因果语言模型。 torch_dtype=torch_dtype
设置了模型使用的PyTorch数据类型,这将影响模型的精度和性能。trust_remote_code=True
表示信任远程加载的代码,这在加载自定义模型或第三方模型时可能需要。.cuda()
将模型移动到GPU上进行加速。
- 使用
-
return tokenizer, model
:- 函数返回加载的分词器和模型。
-
tokenizer, model = get_model()
:- 调用
get_model
函数并将其返回的分词器和模型分别赋值给tokenizer
和model
变量。
- 调用
整体来看,这段代码的作用是定义了一个函数来加载预训练的NLP模型和分词器,并通过Streamlit的缓存机制来优化性能。这样,当Web应用需要使用模型进行文本生成或其他NLP任务时,可以快速地获取已经加载和缓存的模型和分词器。
# 初次运行时,session_state中没有"messages",需要创建一个空列表
if "messages" not in st.session_state:
st.session_state["messages"] = []
# 每次对话时,都需要遍历session_state中的所有消息,并显示在聊天界面上
for msg in st.session_state.messages:
st.chat_message(msg["role"]).write(msg["content"])
这段代码使用了Streamlit库中的session_state
来管理Web应用会话状态,并在聊天界面上显示消息。下面是对代码的逐行解释:
-
if "messages" not in st.session_state:
:- 这行代码检查Streamlit的
session_state
字典中是否存在一个键名为"messages"
的项。session_state
是Streamlit提供的一个全局状态存储,用于跨多个页面重载保持状态。
- 这行代码检查Streamlit的
-
st.session_state["messages"] = []
:- 如果
"messages"
不在session_state
中,这行代码会在session_state
中创建一个新的空列表,并将其与"messages"
键关联。这个列表用于存储聊天界面的消息。
- 如果
-
for msg in st.session_state.messages:
:- 这行代码开始一个循环,遍历
session_state
中"messages"
键对应的消息列表。
- 这行代码开始一个循环,遍历
-
st.chat_message(msg["role"]).write(msg["content"])
:- 对于循环中的每条消息
msg
,这行代码使用st.chat_message
函数来显示消息。msg["role"]
指定了消息的角色(比如"user"或"assistant"),而msg["content"]
包含了消息的文本内容。 st.chat_message(...)
创建了一个聊天气泡,其角色由msg["role"]
决定,.write(...)
则将消息内容写入到这个聊天气泡中。
- 对于循环中的每条消息
整体来看,这段代码的作用是初始化会话状态中的聊天消息列表(如果它尚不存在),然后在聊天界面上显示所有存储在会话状态中的消息。这使得用户可以在Streamlit应用中体验到一个连续的对话流程,每次刷新或重新加载页面时,之前的聊天历史都能被保留并显示出来。
# 如果用户在聊天输入框中输入了内容,则执行以下操作
if prompt := st.chat_input():
# 将用户的输入添加到session_state中的messages列表中
st.session_state.messages.append({"role": "user", "content": prompt})
# 在聊天界面上显示用户的输入
st.chat_message("user").write(prompt)
# 调用模型
prompt = "<n>".join(msg["content"] for msg in st.session_state.messages) + "<sep>" # 拼接对话历史
inputs = tokenizer(prompt, return_tensors="pt")["input_ids"].cuda()
outputs = model.generate(inputs, do_sample=False, max_length=1024) # 设置解码方式和最大生成长度
output = tokenizer.decode(outputs[0])
response = output.split("<sep>")[-1].replace("<eod>", '')
# 将模型的输出添加到session_state中的messages列表中
st.session_state.messages.append({"role": "assistant", "content": response})
# 在聊天界面上显示模型的输出
st.chat_message("assistant").write(response)
这段代码是Streamlit应用程序中的一个聊天交互功能实现,用于处理用户输入、与预训练模型交互,并将模型的输出显示给用户。下面是对代码的逐行解释:
-
if prompt := st.chat_input():
:- 这行代码使用了Streamlit的
st.chat_input()
函数来创建一个聊天输入框,允许用户输入文本。 - Python的
:=
运算符(也称为海象运算符)在这里用于在条件语句中直接赋值。如果用户在输入框中输入了内容并提交,prompt
变量将存储这个输入的值。
- 这行代码使用了Streamlit的
-
st.session_state.messages.append({"role": "user", "content": prompt})
:- 如果用户提交了输入,这行代码将用户输入的文本添加到
session_state
中的messages
列表里,其中"role": "user"
表示消息的角色是用户。
- 如果用户提交了输入,这行代码将用户输入的文本添加到
-
st.chat_message("user").write(prompt)
:- 这行代码在聊天界面上显示用户的输入。
"user"
指定了消息的角色,write
函数用于将用户输入的文本写入到聊天界面中。
- 这行代码在聊天界面上显示用户的输入。
-
prompt = "<n>".join(msg["content"] for msg in st.session_state.messages) + "<sep>"
:- 这行代码将
session_state.messages
列表中的所有消息内容拼接成一个字符串,用于模型的输入。"<n>"
可能是对话中的一句话的分隔符,"<sep>"
作为整个输入序列的结束符。
- 这行代码将
-
inputs = tokenizer(prompt, return_tensors="pt")["input_ids"].cuda()
:- 使用之前加载的分词器对拼接好的
prompt
进行分词处理,转换成模型可以理解的输入ID。 return_tensors="pt"
指定了返回的张量类型为PyTorch张量。.cuda()
将输入张量移动到GPU上,以便使用GPU加速模型的生成过程。
- 使用之前加载的分词器对拼接好的
-
outputs = model.generate(inputs, do_sample=False, max_length=1024)
:- 调用模型生成文本。
do_sample=False
表示不进行采样,即生成确定性的输出;max_length=1024
设置了生成文本的最大长度。
- 调用模型生成文本。
-
output = tokenizer.decode(outputs[0])
:- 使用分词器将模型生成的ID序列解码回文本字符串。
-
response = output.split("<sep>")[-1].replace("<eod>", '')
:- 这行代码处理解码后的输出,通过分割
"<sep>"
来获取最后一段文本,并移除"<eod>"
标记,得到最终的模型响应。
- 这行代码处理解码后的输出,通过分割
-
st.session_state.messages.append({"role": "assistant", "content": response})
:- 将模型的输出添加到
session_state.messages
列表中,其中"role": "assistant"
表示消息的角色是助手。
- 将模型的输出添加到
-
st.chat_message("assistant").write(response)
:- 在聊天界面上显示模型的输出,
"assistant"
指定了消息的角色。
- 在聊天界面上显示模型的输出,
整体来看,这段代码实现了一个聊天机器人的基本功能:接收用户输入,使用预训练模型生成响应,并将用户输入和模型响应显示在聊天界面上。通过session_state
来维护聊天会话的状态,使得每次用户输入和模型响应都能够被记录下来并显示。