utils.py

import csv
import json
from pathlib import Path
from typing import List

import numpy as np
import torch
from transformers import AutoTokenizer

import tensorrt_llm
from tensorrt_llm.quantization import QuantMode
from tensorrt_llm.runtime import ModelConfig

from inference_types import Message

EOS_TOKEN = 2  # codellama
PAD_TOKEN = 2  # codellama

# EOS_TOKEN = 32021  # deepseek
# PAD_TOKEN = 32014  # deepseek


def get_engine_name(model, dtype, tp_size, pp_size, rank):
    if pp_size == 1:
        return "{}_{}_tp{}_rank{}.engine".format(model, dtype, tp_size, rank)
    return "{}_{}_tp{}_pp{}_rank{}.engine".format(model, dtype, tp_size, pp_size, rank)


def to_word_list_format(
    word_dict: List[List[str]], tokenizer=None, add_special_tokens=False
):
    """
    format of word_dict
        len(word_dict) should be same to batch_size
        word_dict[i] means the words for batch i
        len(word_dict[i]) must be 1, which means it only contains 1 string
        This string can contains several sentences and split by ",".
        For example, if word_dict[2] = " I am happy, I am sad", then this function will return
        the ids for two short sentences " I am happy" and " I am sad".
    """
    assert tokenizer != None, "need to set tokenizer"

    flat_ids = []
    offsets = []
    for word_dict_item in word_dict:
        item_flat_ids = []
        item_offsets = []

        if isinstance(word_dict_item[0], bytes):
            word_dict_item = [word_dict_item[0].decode()]

        words = list(csv.reader(word_dict_item))[0]
        for word in words:
            ids = tokenizer.encode(word, add_special_tokens=add_special_tokens)

            if len(ids) == 0:
                continue

            item_flat_ids += ids
            item_offsets.append(len(ids))

        flat_ids.append(np.array(item_flat_ids))
        offsets.append(np.cumsum(np.array(item_offsets)))

    pad_to = max(1, max(len(ids) for ids in flat_ids))

    for i, (ids, offs) in enumerate(zip(flat_ids, offsets)):
        flat_ids[i] = np.pad(ids, (0, pad_to - len(ids)), constant_values=0)
        offsets[i] = np.pad(offs, (0, pad_to - len(offs)), constant_values=-1)

    return np.array([flat_ids, offsets], dtype="int32").transpose((1, 0, 2))


def read_config(config_path: Path):
    with open(config_path, "r") as f:
        config = json.load(f)
    use_gpt_attention_plugin = config["plugin_config"]["gpt_attention_plugin"]
    remove_input_padding = config["plugin_config"]["remove_input_padding"]
    dtype = config["builder_config"]["precision"]
    tp_size = config["builder_config"]["tensor_parallel"]
    pp_size = config["builder_config"]["pipeline_parallel"]
    world_size = tp_size * pp_size
    assert (
        world_size == tensorrt_llm.mpi_world_size()
    ), f"Engine world size ({world_size}) != Runtime world size ({tensorrt_llm.mpi_world_size()})"
    num_heads = config["builder_config"]["num_heads"] // tp_size
    hidden_size = config["builder_config"]["hidden_size"] // tp_size
    vocab_size = config["builder_config"]["vocab_size"]
    num_layers = config["builder_config"]["num_layers"]
    num_kv_heads = config["builder_config"].get("num_kv_heads", num_heads)
    paged_kv_cache = config["plugin_config"]["paged_kv_cache"]
    tokens_per_block = config["plugin_config"]["tokens_per_block"]
    quant_mode = QuantMode(config["builder_config"]["quant_mode"])
    if config["builder_config"].get("multi_query_mode", False):
        tensorrt_llm.logger.warning(
            "`multi_query_mode` config is deprecated. Please rebuild the engine."
        )
        num_kv_heads = 1
    num_kv_heads = (num_kv_heads + tp_size - 1) // tp_size
    use_custom_all_reduce = config["plugin_config"].get("use_custom_all_reduce", False)

    # print("Engine config:")
    # print(f"  dtype: {dtype}")
    # print(f"  tensor_parallel: {tp_size}")
    # print(f"  pipeline_parallel: {pp_size}")
    # print(f"  num_heads: {num_heads}")
    # print(f"  hidden_size: {hidden_size}")
    # print(f"  vocab_size: {vocab_size}")
    # print(f"  num_layers: {num_layers}")
    # print(f"  num_kv_heads: {num_kv_heads}")
    # print(f"  paged_kv_cache: {paged_kv_cache}")
    # print(f"  tokens_per_block: {tokens_per_block}")
    # print(f"  quant_mode: {quant_mode}")
    # print(f"  use_custom_all_reduce: {use_custom_all_reduce}")

    model_config = ModelConfig(
        num_heads=num_heads,
        num_kv_heads=num_kv_heads,
        hidden_size=hidden_size,
        vocab_size=vocab_size,
        num_layers=num_layers,
        gpt_attention_plugin=use_gpt_attention_plugin,
        paged_kv_cache=paged_kv_cache,
        tokens_per_block=tokens_per_block,
        remove_input_padding=remove_input_padding,
        dtype=dtype,
        quant_mode=quant_mode,
        use_custom_all_reduce=use_custom_all_reduce,
    )

    return model_config, tp_size, pp_size


def pad_input_tokens(tokens, pad_token_ids, max_length):
    pad_tokens_length = max_length - len(tokens)
    pad_tokens = []
    while len(pad_tokens) < pad_tokens_length:
        for token_id in pad_token_ids:
            pad_tokens.append(token_id)
            if len(pad_tokens) >= pad_tokens_length:
                break
    return pad_tokens + tokens


def prepare_batch_inputs(
    batch_input_texts: List[str],
    tokenizer: AutoTokenizer,
):
    input_tokens = []
    for input_text in batch_input_texts:
        input_tokens.append(tokenizer.encode(input_text, add_special_tokens=False))

    max_input_length = max(len(tokens) for tokens in input_tokens)

    pad_token_ids = [836, 31857, 3078, 31908]  # codellama
    # pad_token_ids = [32014]  # deepseek-coder
    input_tokens = [
        pad_input_tokens(
            tokens, pad_token_ids=pad_token_ids, max_length=max_input_length
        )
        for tokens in input_tokens
    ]
    input_ids = np.array(input_tokens)
    input_ids = torch.tensor(input_ids, dtype=torch.int32, device="cuda")
    print(input_ids)

    input_lengths = torch.tensor(
        [max_input_length for x in input_tokens], dtype=torch.int32, device="cuda"
    )

    return input_ids, input_lengths


def parse_output(output_ids, input_lengths, max_output_length, tokenizer):
    def trim_output_tokens_after_eos_token(
        output_tokens: List[int], eos_token: int = 2
    ):
        if 2 in output_tokens:
            index_of_eos = output_tokens.index(eos_token)
            trimmed_output_tokens = output_tokens[:index_of_eos]
            return trimmed_output_tokens
        else:
            return output_tokens

    output_texts: List[str] = []
    output_tokens_lengths: List[int] = []
    for batch in range(input_lengths.size(0)):
        for beam in range(output_ids.size(1)):
            output_begin = input_lengths[batch]
            output_end = input_lengths[batch] + max_output_length
            outputs = output_ids[batch][beam][output_begin:output_end].tolist()
            outputs = trim_output_tokens_after_eos_token(outputs)
            output_texts.append(tokenizer.decode(outputs).strip())
            output_tokens_lengths.append(len(outputs))

    return output_texts, output_tokens_lengths


def format_llama_conversation(messages) -> str:
    if not isinstance(messages, list):
        raise ValueError("Invalid messages format")
    start_token = "<s>"
    end_token = "</s>"
    prompt = start_token
    right_after_system_prompt = False
    for message in messages:
        if isinstance(message, Message):
            message = message.__dict__
        role = message.get("role")
        content = message.get("content").strip()
        if role == "system":
            prompt += f"""[INST] <<SYS>>
{content}
<</SYS>>

"""
            right_after_system_prompt = True
            continue
        if role == "user":
            if right_after_system_prompt:
                prompt += f"""{content} [/INST]"""
                right_after_system_prompt = False
            else:
                prompt += f"""[INST] {content} [/INST]"""
        elif role == "assistant":
            prompt += f""" {content} {end_token}{start_token}"""
        elif role != "system":
            raise ValueError(f"Invalid message role {message['role']}")
    if prompt.endswith(start_token):
        prompt = prompt[: -len(start_token)]
    return prompt


def format_deepseek_conversation(messages) -> str:
    if not isinstance(messages, list):
        raise ValueError("Invalid messages format")
    prompt = "<｜begin▁of▁sentence｜>"
    for message in messages:
        if isinstance(message, Message):
            message = message.__dict__
        role = message.get("role")
        content = message.get("content").strip()
        if role == "system":
            prompt += content + "\n"
        else:
            if role == "user":
                prompt += "### Instruction:\n" + content + "\n"
            elif role == "assistant":
                prompt += "### Response:\n" + content + "\n<|EOT|>\n"
    last = messages[-1]
    if isinstance(last, Message):
        last = last.__dict__
    if last.get("role") != "assistant":
        prompt += "### Response:\n"
    return prompt


if __name__ == "__main__":
    messages = [
        {
            "role": "system",
            "content": "Act as a helpful, respectful, and honest programming assistant.\nYour answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content.\nNever mention that you're an AI.",
        },
        {"role": "user", "content": "Hello, how are you?"},
        {"role": "assistant", "content": "I'm good, how are you?"},
        {"role": "user", "content": "Tell me a story!"},
    ]
    llama_formatted_prompt = format_llama_conversation(messages)
    print(llama_formatted_prompt)
    deepseek_formatted_prompt = format_deepseek_conversation(messages)
    print(deepseek_formatted_prompt)