InfLLM的笔记

news2024/11/18 15:41:48

文件中提供的代码是一个Python函数chat_loop，它是聊天系统的核心循环。以下是对这段代码逻辑的梳理：

函数定义与参数

chat_loop函数接收多个参数，用于配置聊天模型和聊天环境。
参数包括模型路径、设备类型、GPU数量、最大GPU内存、数据类型、是否加载8位模型、CPU卸载、对话模板、系统消息、温度参数、重复惩罚、最大新token数量、聊天IO对象等。

加载模型

使用load_model函数加载模型和分词器（tokenizer），这个函数根据提供的参数配置模型。
如果提供了inf_llm_config，则使用patch_hf函数对模型进行补丁处理。

设置模型类型和默认参数

根据模型类型（如T5、codet5p、xft等），可能需要设置特定的默认参数，例如T5模型的重复惩罚默认值设为1.2。

设置上下文长度

根据模型配置获取上下文长度，如果使用InfLLM补丁，则设置一个非常大的上下文长度。

新建和重新加载对话

new_chat函数用于创建一个新的对话实例，根据是否提供conv_template来选择对话模板。
reload_conv函数用于重新打印对话内容。

主聊天循环

使用while True创建一个无限循环，代表聊天系统的持续运行。
使用chatio.prompt_for_input函数提示用户输入，如果输入为空或者为退出命令（如"!!exit"），则退出循环。
支持对话控制命令，如"!!reset"重置对话，"!!remove"删除最后一条消息，"!!regen"重新生成最后一条消息，"!!save"保存对话，"!!load"加载对话。

生成输出

将用户输入添加到对话中，并生成提示（prompt）。
根据模型类型（如codet5p），可能需要特别处理提示。
设置生成文本的参数gen_params。
调用generate_stream_func函数生成输出流。
使用chatio.stream_output函数处理输出流并生成最终的文本输出。
如果设置了调试模式，将打印调试信息，包括对话模板、提示、输出和生成速度。

异常处理

使用try-except结构捕获KeyboardInterrupt异常，以便在用户尝试中断生成时处理。

清理和缓存管理

在生成输出后，根据需要清理缓存或更新对话状态。

整体而言，chat_loop函数是聊天系统的主控函数，负责管理聊天会话的流程，包括加载模型、处理用户输入、生成和输出文本、以及异常处理。

——————————————————————

该文件是一个Python脚本，它包含了一个基于FastChat模型的聊天系统，FastChat模型最初由LMSYS团队开发。这个脚本在原有代码的基础上进行了修改，增加了对InfLLM补丁的支持。以下是代码逻辑的梳理：

导入依赖：脚本开始部分导入了所需的所有库和模块，包括torch、json、argparse等。
Inference for FastChat models：这部分代码提供了FastChat模型的推理功能，定义了generate_stream函数，该函数用于生成聊天的输出流。
参数读取：在generate_stream函数内部，首先读取了一系列参数，包括prompt、temperature、repetition_penalty、top_p、top_k、max_new_tokens等，这些参数控制生成文本的行为。
日志概率处理器：使用prepare_logits_processor函数准备一个日志概率处理器，用于处理生成文本时的逻辑。
编码输入：将prompt转换为模型可理解的编码格式。
生成文本流：在generate_stream函数中，通过迭代的方式生成文本。在每次迭代中，模型都会生成一个或多个token，并根据设置的条件（如stream_interval、stop_token_ids等）决定是否输出这些token。
聊天循环：定义了chat_loop函数，该函数初始化模型和分词器，设置聊天环境，并进入一个循环，不断接收输入并生成输出，直到接收到退出命令。
模型加载：在chat_loop函数中，调用load_model函数来加载指定路径的模型。
聊天界面：根据命令行参数，初始化不同类型的聊天界面（SimpleChatIO、RichChatIO或ProgrammaticChatIO）。
命令行参数解析：脚本末尾部分定义了命令行参数解析逻辑，允许用户通过命令行指定模型路径、设备类型、温度参数、重复惩罚、最大新token数量等。
主函数：定义了main函数，它处理命令行参数，并启动聊天循环。
对话模板：定义了一个Llama3Conv类，用于生成和管理对话模板。
注册对话模板：通过register_conv_template函数注册了Llama3Conv类的一个实例，这个实例定义了对话的格式和角色。
入口点：脚本包含一个标准的Python入口点，即if __name__ == "__main__":部分，它解析命令行参数并调用main函数。