Datawhale AI夏令营 第四期 大模型应用开发 Task1：案例：智能编程助手

智能编程助手代码解读

# 导入所需的库
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import streamlit as st

这段代码是Python语言编写的，用于导入一些在机器学习和深度学习中常用的库。下面是对每行代码的解释：

1. from transformers import AutoTokenizer, AutoModelForCausalLM:

   • 这行代码从transformers库中导入了两个类：AutoTokenizer和AutoModelForCausalLM。
   • AutoTokenizer是用于自动处理文本数据的分词器，它可以根据不同的预训练模型自动选择相应的分词方法。
   • AutoModelForCausalLM是用于因果语言模型（Causal Language Modeling，即生成文本的任务）的模型类。

2. import torch:

   • 这行代码导入了PyTorch库，PyTorch是一个开源的机器学习库，广泛用于深度学习的研究和开发。它提供了强大的GPU加速的张量计算能力，以及构建深度学习模型的动态计算图。

3. import streamlit as st:

   • 这行代码导入了Streamlit库，Streamlit是一个用于快速构建数据应用的开源库，它可以将Python脚本转换为交互式Web应用。as st是给库指定了一个别名st，这样在代码中就可以通过st来访问Streamlit库的所有功能。

# 创建一个标题和一个副标题
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装饰器来缓存函数的结果。下面是对代码的详细解释：

1. @st.cache_resource:

   • 这是一个Streamlit提供的装饰器，用于缓存函数的执行结果。当函数再次被调用时，如果输入参数没有变化，Streamlit会使用缓存的结果而不是重新执行函数。这可以显著提高应用的性能，尤其是在加载大型模型或执行耗时操作时。

2. 函数定义 def get_model():

   • 定义了一个名为get_model的函数，用于加载和返回分词器和模型。

3. 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>'。

4. tokenizer.add_tokens(['<sep>', '<pad>', ..., '<empty_output>'], special_tokens=True):

   • 向分词器中添加一系列特殊的标记，这些标记通常用于处理特定的NLP任务，如序列到序列的模型。
   • special_tokens=True表示这些标记将作为特殊的分词器标记处理，它们将被加入到分词器的词汇表中。

5. 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上进行加速。

6. return tokenizer, model:

   • 函数返回加载的分词器和模型。

7. 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应用会话状态，并在聊天界面上显示消息。下面是对代码的逐行解释：

1. if "messages" not in st.session_state::

   • 这行代码检查Streamlit的session_state字典中是否存在一个键名为"messages"的项。session_state是Streamlit提供的一个全局状态存储，用于跨多个页面重载保持状态。

2. st.session_state["messages"] = []:

   • 如果"messages"不在session_state中，这行代码会在session_state中创建一个新的空列表，并将其与"messages"键关联。这个列表用于存储聊天界面的消息。

3. for msg in st.session_state.messages::

   • 这行代码开始一个循环，遍历session_state中"messages"键对应的消息列表。

4. 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应用程序中的一个聊天交互功能实现，用于处理用户输入、与预训练模型交互，并将模型的输出显示给用户。下面是对代码的逐行解释：

1. if prompt := st.chat_input()::

   • 这行代码使用了Streamlit的st.chat_input()函数来创建一个聊天输入框，允许用户输入文本。
   • Python的:=运算符（也称为海象运算符）在这里用于在条件语句中直接赋值。如果用户在输入框中输入了内容并提交，prompt变量将存储这个输入的值。

2. st.session_state.messages.append({"role": "user", "content": prompt}):

   • 如果用户提交了输入，这行代码将用户输入的文本添加到session_state中的messages列表里，其中"role": "user"表示消息的角色是用户。

3. st.chat_message("user").write(prompt):

   • 这行代码在聊天界面上显示用户的输入。"user"指定了消息的角色，write函数用于将用户输入的文本写入到聊天界面中。

4. prompt = "<n>".join(msg["content"] for msg in st.session_state.messages) + "<sep>":

   • 这行代码将session_state.messages列表中的所有消息内容拼接成一个字符串，用于模型的输入。"<n>"可能是对话中的一句话的分隔符，"<sep>"作为整个输入序列的结束符。

5. inputs = tokenizer(prompt, return_tensors="pt")["input_ids"].cuda():

   • 使用之前加载的分词器对拼接好的prompt进行分词处理，转换成模型可以理解的输入ID。
   • return_tensors="pt"指定了返回的张量类型为PyTorch张量。
   • .cuda()将输入张量移动到GPU上，以便使用GPU加速模型的生成过程。

6. outputs = model.generate(inputs, do_sample=False, max_length=1024):

   • 调用模型生成文本。do_sample=False表示不进行采样，即生成确定性的输出；max_length=1024设置了生成文本的最大长度。

7. output = tokenizer.decode(outputs[0]):

   • 使用分词器将模型生成的ID序列解码回文本字符串。

8. response = output.split("<sep>")[-1].replace("<eod>", ''):

   • 这行代码处理解码后的输出，通过分割"<sep>"来获取最后一段文本，并移除"<eod>"标记，得到最终的模型响应。

9. st.session_state.messages.append({"role": "assistant", "content": response}):

   • 将模型的输出添加到session_state.messages列表中，其中"role": "assistant"表示消息的角色是助手。

10. st.chat_message("assistant").write(response):

    • 在聊天界面上显示模型的输出，"assistant"指定了消息的角色。

整体来看，这段代码实现了一个聊天机器人的基本功能：接收用户输入，使用预训练模型生成响应，并将用户输入和模型响应显示在聊天界面上。通过session_state来维护聊天会话的状态，使得每次用户输入和模型响应都能够被记录下来并显示。