【AI大模型开发者指南】llama3 web部署实战与原理解析（logits_warper、logits_processor）

【AI大模型开发者指南】llama3 web部署实战与原理解析

streamlit脚本

import copy
import warnings
from dataclasses import asdict, dataclass
from typing import Callable, List, Optional

import streamlit as st
import torch
from torch import nn
from transformers.generation.utils import (LogitsProcessorList,
                                           StoppingCriteriaList)
from transformers.utils import logging

from transformers import AutoTokenizer, AutoModelForCausalLM  # isort: skip

logger = logging.get_logger(__name__)


@dataclass
class GenerationConfig:
    # this config is used for chat to provide more diversity
    max_length: int = 32768
    top_p: float = 0.8
    temperature: float = 0.8
    do_sample: bool = True
    repetition_penalty: float = 1.005


@torch.inference_mode()
def generate_interactive(
    model,
    tokenizer,
    prompt,
    generation_config: Optional[GenerationConfig] = None,
    logits_processor: Optional[LogitsProcessorList] = None,
    stopping_criteria: Optional[StoppingCriteriaList] = None,
    prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor],
                                                List[int]]] = None,
    additional_eos_token_id: Optional[int] = None,
    **kwargs,
):
    inputs = tokenizer([prompt], return_tensors='pt')
    input_length = len(inputs['input_ids'][0])
    for k, v in inputs.items():
        inputs[k] = v.cuda()
    input_ids = inputs['input_ids']
    _, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1]
    if generation_config is None:
        generation_config = model.generation_config
    generation_config = copy.deepcopy(generation_config)
    model_kwargs = generation_config.update(**kwargs)
    bos_token_id, eos_token_id = (  # noqa: F841  # pylint: disable=W0612
        generation_config.bos_token_id,
        generation_config.eos_token_id,
    )
    if isinstance(eos_token_id, int):
        eos_token_id = [eos_token_id]
    if additional_eos_token_id is not None:
        eos_token_id.append(additional_eos_token_id)
    has_default_max_length = kwargs.get(
        'max_length') is None and generation_config.max_length is not None
    if has_default_max_length and generation_config.max_new_tokens is None:
        warnings.warn(
            f"Using 'max_length''s default ({repr(generation_config.max_length)}) \
                to control the generation length. "
            'This behaviour is deprecated and will be removed from the \
                config in v5 of Transformers -- we'
            ' recommend using `max_new_tokens` to control the maximum \
                length of the generation.',
            UserWarning,
        )
    elif generation_config.max_new_tokens is not None:
        generation_config.max_length = generation_config.max_new_tokens + \
            input_ids_seq_length
        if not has_default_max_length:
            logger.warn(  # pylint: disable=W4902
                f"Both 'max_new_tokens' (={generation_config.max_new_tokens}) "
                f"and 'max_length'(={generation_config.max_length}) seem to "
                "have been set. 'max_new_tokens' will take precedence. "
                'Please refer to the documentation for more information. '
                '(https://huggingface.co/docs/transformers/main/'
                'en/main_classes/text_generation)',
                UserWarning,
            )

    if input_ids_seq_length >= generation_config.max_length:
        input_ids_string = 'input_ids'
        logger.warning(
            f"Input length of {input_ids_string} is {input_ids_seq_length}, "
            f"but 'max_length' is set to {generation_config.max_length}. "
            'This can lead to unexpected behavior. You should consider'
            " increasing 'max_new_tokens'.")

    # 2. Set generation parameters if not already defined
    logits_processor = logits_processor if logits_processor is not None \
        else LogitsProcessorList()
    stopping_criteria = stopping_criteria if stopping_criteria is not None \
        else StoppingCriteriaList()

    logits_processor = model._get_logits_processor(
        generation_config=generation_config,
        input_ids_seq_length=input_ids_seq_length,
        encoder_input_ids=input_ids,
        prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
        logits_processor=logits_processor,
    )

    stopping_criteria = model._get_stopping_criteria(
        generation_config=generation_config,
        stopping_criteria=stopping_criteria)
    logits_warper = model._get_logits_warper(generation_config)

    unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)
    scores = None
    while True:
        model_inputs = model.prepare_inputs_for_generation(
            input_ids, **model_kwargs)
        # forward pass to get next token
        outputs = model(
            **model_inputs,
            return_dict=True,
            output_attentions=False,
            output_hidden_states=False,
        )

        next_token_logits = outputs.logits[:, -1, :]

        # pre-process distribution
        next_token_scores = logits_processor(input_ids, next_token_logits)
        next_token_scores = logits_warper(input_ids, next_token_scores)

        # sample
        probs = nn.functional.softmax(next_token_scores, dim=-1)
        if generation_config.do_sample:
            next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
        else:
            next_tokens = torch.argmax(probs, dim=-1)

        # update generated ids, model inputs, and length for next step
        input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
        model_kwargs = model._update_model_kwargs_for_generation(
            outputs, model_kwargs, is_encoder_decoder=False)
        unfinished_sequences = unfinished_sequences.mul(
            (min(next_tokens != i for i in eos_token_id)).long())

        output_token_ids = input_ids[0].cpu().tolist()
        output_token_ids = output_token_ids[input_length:]
        for each_eos_token_id in eos_token_id:
            if output_token_ids[-1] == each_eos_token_id:
                output_token_ids = output_token_ids[:-1]
        response = tokenizer.decode(output_token_ids)

        yield response
        # stop when each sentence is finished
        # or if we exceed the maximum length
        if unfinished_sequences.max() == 0 or stopping_criteria(
                input_ids, scores):
            break


def on_btn_click():
    del st.session_state.messages


@st.cache_resource
def load_model(arg1):
    # model = AutoModelForCausalLM.from_pretrained(args.m).cuda()
    # tokenizer = AutoTokenizer.from_pretrained(args.m, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(arg1, torch_dtype=torch.float16).cuda()
    tokenizer = AutoTokenizer.from_pretrained(arg1, trust_remote_code=True)

  
    return model, tokenizer


def prepare_generation_config():
    with st.sidebar:
        max_length = st.slider('Max Length',
                               min_value=8,
                               max_value=8192,
                               value=8192)
        top_p = st.slider('Top P', 0.0, 1.0, 0.8, step=0.01)
        temperature = st.slider('Temperature', 0.0, 1.0, 0.7, step=0.01)
        st.button('Clear Chat History', on_click=on_btn_click)

    generation_config = GenerationConfig(max_length=max_length,
                                         top_p=top_p,
                                         temperature=temperature)

    return generation_config


user_prompt = '<|start_header_id|>user<|end_header_id|>\n\n{user}<|eot_id|>'
robot_prompt = '<|start_header_id|>assistant<|end_header_id|>\n\n{robot}<|eot_id|>'
cur_query_prompt = '<|start_header_id|>user<|end_header_id|>\n\n{user}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n'


def combine_history(prompt):
    messages = st.session_state.messages
    total_prompt = ''
    for message in messages:
        cur_content = message['content']
        if message['role'] == 'user':
            cur_prompt = user_prompt.format(user=cur_content)
        elif message['role'] == 'robot':
            cur_prompt = robot_prompt.format(robot=cur_content)
        else:
            raise RuntimeError
        total_prompt += cur_prompt
    total_prompt = total_prompt + cur_query_prompt.format(user=prompt)
    return total_prompt


def main(arg1):
    # torch.cuda.empty_cache()
    print('load model begin.')
    model, tokenizer = load_model(arg1)
    print('load model end.')

    st.title('Llama3-Instruct')

    generation_config = prepare_generation_config()

    # Initialize chat history
    if 'messages' not in st.session_state:
        st.session_state.messages = []

    # Display chat messages from history on app rerun
    for message in st.session_state.messages:
        with st.chat_message(message['role']):
            st.markdown(message['content'])

    # Accept user input
    if prompt := st.chat_input('Hello!'):
        # Display user message in chat message container
        with st.chat_message('user'):
            st.markdown(prompt)
        real_prompt = combine_history(prompt)
        # Add user message to chat history
        st.session_state.messages.append({
            'role': 'user',
            'content': prompt,
        })

        with st.chat_message('robot'):
            message_placeholder = st.empty()
            for cur_response in generate_interactive(
                    model=model,
                    tokenizer=tokenizer,
                    prompt=real_prompt,
                    additional_eos_token_id=128009,  # <|eot_id|>
                    **asdict(generation_config),
            ):
                # Display robot response in chat message container
                message_placeholder.markdown(cur_response + '▌')
            message_placeholder.markdown(cur_response)
        # Add robot response to chat history
        st.session_state.messages.append({
            'role': 'robot',
            'content': cur_response,  # pylint: disable=undefined-loop-variable
        })
        torch.cuda.empty_cache()


if __name__ == '__main__':

    import sys
    arg1 = sys.argv[1]
    main(arg1)

基于streamlit的web部署的核心便是实现streamlit app，以上是一个llama3 web demo的streamlit脚本，输入参数只有一个——llama3的模型目录

让我们来分析该脚本

主函数

main 函数作为一个基于 Streamlit 的交互式聊天应用程序的主入口，该程序使用机器学习模型进行自然语言处理。以下是对 main 函数关键组成部分和步骤的解析：

1. 模型和分词器加载：

• 函数开始时先清理 CUDA 缓存以释放 GPU 内存，尽管这一行被注释掉了（# torch.cuda.empty_cache()）。
• 然后，输出加载模型过程开始和结束的消息。
• 调用 load_model 函数并传入参数 arg1（可能是模型标识符），将预训练的因果语言模型和分词器加载到 CUDA 内存中，并将张量转换为 float16 以提高计算效率。

2. Streamlit 用户界面设置：

• 函数使用 prepare_generation_config 函数设置 Streamlit 界面，开始显示标题“Llama3-Instruct”。
• 在侧边栏使用滑块让用户可以调整生成参数，如 max_length、top_p 和 temperature。

3. 聊天历史管理：

• 初始化或维护聊天历史（st.session_state.messages）以跟踪对话。
• 通过 st.chat_message 在应用程序界面上显示过去的消息。

4. 用户输入处理：

• 应用程序通过 st.chat_input 接受用户输入，并使用 st.chat_message 显示用户消息。
• 每个用户消息都被追加到会话状态的聊天历史中。

5. 消息生成：

• 将用户输入与现有聊天历史结合，形成模型的提示。
• 调用 generate_interactive 函数从模型生成响应。该函数迭代处理生成过程，允许在聊天界面动态显示中间输出。
• 模型的最终响应被显示并添加到聊天历史中。

6. 资源管理：

• 在处理用户输入和模型响应后，函数显式清除 CUDA 缓存，以有效管理 GPU 内存。

7. 错误处理和循环结构：

• 函数包含错误处理，用于处理聊天历史中角色分配错误。
• 聊天应用程序设计为持续接受和处理用户输入，直到会话结束，由 Streamlit 的响应式编程模型支持。

8. 程序执行控制：

• if __name__ == '__main__': 代码块检查脚本是否作为主模块运行，并获取命令行参数传递给 main 函数。

总体而言，main 函数负责协调用户与 AI 模型之间的互动，高效地处理用户界面和计算后端。然而，可能需要更好的错误处理来应对运行时可能出现的问题，如模型加载失败或意外的用户输入。此外，评论为何禁用某些功能（如 CUDA 缓存清理）并确保在实时应用中清晰管理 GPU 资源，可以提高可靠性和性能。

重点关注for cur_response in generate_interactive

        with st.chat_message('robot'):
            message_placeholder = st.empty()
            for cur_response in generate_interactive(
                    model=model,
                    tokenizer=tokenizer,
                    prompt=real_prompt,
                    additional_eos_token_id=128009,  # <|eot_id|>
                    **asdict(generation_config),
            ):
                # Display robot response in chat message container
                message_placeholder.markdown(cur_response + '▌')
            message_placeholder.markdown(cur_response)
        # Add robot response to chat history
        st.session_state.messages.append({
            'role': 'robot',
            'content': cur_response,  # pylint: disable=undefined-loop-variable
        })

这一段代码是模型输出的关键

		output_token_ids = input_ids[0].cpu().tolist()
        output_token_ids = output_token_ids[input_length:]
        for each_eos_token_id in eos_token_id:
            if output_token_ids[-1] == each_eos_token_id:
                output_token_ids = output_token_ids[:-1]
        response = tokenizer.decode(output_token_ids)

yield response内容如上，接下来让我们观察原始token_ids输出过程

以上是output_token_ids的变化，可以看到输出过程

原理就是

input到模型，模型输出logits，logits首先经过logits_processor，然后经过logits_warper，最后经过softmax变成概率，根据概率选出tokens，拼到input_ids中（追加到input_ids中）

然后就到了yield response部分的token变化，每次根据input length获取到output tokens，当判断到最后一个token是eos_token时，截断

每次都解码整个output_tokens_ids作为自然文本输出

logits_processor 和 logits_warper

两者都是对logits的处理，可以参考
https://huggingface.co/docs/transformers/en/internal/generation_utils
https://huggingface.co/transformers/v4.1.1/_modules/transformers/generation_utils.html

可以看到，warper处理top_k、top_p、temperature这些参数，processor处理penalty、bad_words、min_length这些