infer.py

# import llm_utils
import dataclasses
import json
from transformers import AutoTokenizer, AutoConfig
import torch
from torchvision.transforms.functional import InterpolationMode
import numpy as np
from ml_dtypes import bfloat16
from axengine import InferenceSession
from tqdm import tqdm
import torchvision.transforms as T
from PIL import Image
import argparse


def post_process(data, topk=1, topp=0.9, temperature=0.6):
    def top_p(l: np.ndarray, p: float) -> np.ndarray:
        index = np.argsort(l)
        res = l.copy()
        sum_p = 0
        for i in index[::-1]:
            if sum_p >= p:
                res[i] = 0
            sum_p += res[i]
        return res / sum_p

    def softmax(l: np.ndarray) -> np.ndarray:
        l_max = l - l.max()
        l_exp = np.exp(l_max)
        res = l_exp / np.sum(l_exp)
        return res.astype(np.float64)

    r = data.astype(np.float32)
    r = r.flatten()
    # topk
    candidate_index = np.argpartition(r, -topk)[-topk:]
    candidate_value = r[candidate_index]
    # temperature
    candidate_value /= temperature
    # softmax
    candidate_soft = softmax(candidate_value)
    # topp
    candidate_soft = top_p(candidate_soft, topp)
    candidate_soft = candidate_soft.astype(np.float64) / candidate_soft.sum()
    pos = np.random.multinomial(1, candidate_soft).argmax()
    next_token = candidate_index[pos]
    return next_token, candidate_index, candidate_soft


def generate_slice_indices(token_len, prefill=128, expand=512):
    remaining = max(0, token_len - prefill)
    extra_blocks = (remaining + expand - 1) // expand
    return list(range(extra_blocks + 1))


if __name__ == "__main__":

    prompt = None
    parser = argparse.ArgumentParser(description="Model configuration parameters")
    parser.add_argument("--hf_model", type=str, default="../qwen2.5_tokenizer",
                        help="Path to HuggingFace model")
    parser.add_argument("--axmodel_path", type=str, default="../qwen2.5-1.5b-ctx-ax650",
                        help="Path to save compiled axmodel of llama model")
    parser.add_argument("-q", "--question", type=str, default="Please calculate the derivative of the function y=2x^2.",
                        help="Your question that you want to ask the model.")
    args = parser.parse_args()

    device = "cpu"
    hf_model_path = args.hf_model
    axmodel_path = args.axmodel_path

    cfg = AutoConfig.from_pretrained(hf_model_path, trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained(hf_model_path, trust_remote_code=True, use_fast=False)

    prompt = args.question

    print("注意: 这里预先设定了长 prompt 测试. ")
    # input_ids: 438 > 128
    prompt = '''你能将英文翻译成中文吗? 比如这句话:
    Once when I was six years old I saw a magnificent picture in a book, called True Stories from Nature, about the primeval forest. It was a picture of a boa constrictor in the act of swallowing an animal. Here is a copy of the drawing. 
    In the book it said: "Boa constrictors swallow their prey whole, without chewing it. After that they are not able to move, and they sleep through the six months that they need for digestion." 
    I pondered deeply, then, over the adventures of the jungle. And after some work with a colored pencil I succeeded in making my first drawing. My Drawing Number One. It looked like this: 
    I showed my masterpiece to the grown-ups, and asked them whether the drawing frightened them.
    But they answered: "Frighten? Why should any one be frightened by a hat?"
    My drawing was not a picture of a hat. It was a picture of a boa constrictor digesting an elephant. But since the grown-ups were not able to understand it, I made another drawing: I drew the inside of the boa constrictor, so that the grown-ups could see it clearly. They always need to have things explained. My Drawing Number Two looked like this: 
    The grown-ups‘ response, this time, was to advise me to lay aside my drawings of boa constrictors, whether from the inside or the outside, and devote myself instead to geography, history, arithmetic and grammar. That is why, at the age of six, I gave up what might have been a magnificent career as a painter. I had been disheartened by the failure of my Drawing Number One and my Drawing Number Two. Grown-ups never understand anything by themselves, and it is tiresome for children to be always and forever explaining things to them.
    So then I chose another profession, and learned to pilot airplanes.
    '''

    # prompt = "介绍一下你自己"
    # prompt = "今天是几号，天气怎么样"
    prompt = "你知道 `床前明月光,疑是地上霜`是谁写的吗?"

    messages = [
        {"role": "system", "content": "你的名字叫小智(allen), 你是一个人畜无害的 AI 助手. 深圳市今天(4月1日)阴天, 愚人节, 气温在 14°C 至 19°C 之间, 微风."},
        {"role": "user", "content": prompt}
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(device)
    token_ids = model_inputs.input_ids[0].cpu().numpy().tolist()

    embeds = np.load(f"{axmodel_path}/model.embed_tokens.weight.npy")
    prefill_data = np.take(embeds, token_ids, axis=0)
    token_len = len(token_ids)
    import pdb; pdb.set_trace()

    ##################
    lastN = 2559

    kv_dim = cfg.hidden_size // cfg.num_attention_heads * cfg.num_key_value_heads
    k_caches = [
        np.zeros((1, lastN, kv_dim), dtype=bfloat16)
        for _ in range(cfg.num_hidden_layers)
    ]
    v_caches = [
        np.zeros((1, lastN, kv_dim), dtype=bfloat16)
        for _ in range(cfg.num_hidden_layers)
    ]

    prefill_decoder_sessins = []

    for i in tqdm(range(cfg.num_hidden_layers), desc="Init InferenceSession"):
        session = InferenceSession(
            f"{axmodel_path}/qwen2_p128_l{i}_together.axmodel"
        )
        prefill_decoder_sessins.append(session)

    post_process_session = InferenceSession(
        f"{axmodel_path}/qwen2_post.axmodel"
    )
    print("model load done!")
    print("prefill token_len: ", token_len)

    """
    Model input shape:
        - kv_cache: g1:[1, 1, hidden_size] -> g2:[1, kv_mask_expand_lens, kv_dim] -> g3:[1, kv_mask_expand_lens * 2, kv_dim]
        - mask: g1:[1, input_prefill_len, input_prefill_len] -> g2:[1, input_prefill_len, input_prefill_len+kv_mask_expand_lens] -> g3:[1, input_prefill_len, input_prefill_len+kv_mask_expand_lens*2]
        - indices: g1:[1, input_prefill_len] -> g2:[1, input_prefill_len] -> g3:[1, input_prefill_len]
        - input: g1:[1, input_prefill_len, hidden_size] -> g2:[1, input_prefill_len, hidden_size] -> g3:[1, input_prefill_len, hidden_size]
    """

    input_prefill_len = 128
    kv_mask_expand_len = 128 # 512

    """
    Model output shape:
        - kv_cache: g1:[1, input_prefill_len, kv_dim] -> g2:[1, input_prefill_len, kv_dim] -> g3:[1, input_prefill_len, kv_dim]
        - output: g1:[1, input_prefill_len, hidden_size] -> g2:[1, input_prefill_len, hidden_size] -> g3:[1, input_prefill_len, hidden_size]
    """
    slice_indexs = generate_slice_indices(token_len, input_prefill_len, input_prefill_len)
    print(f"slice_indexs is {slice_indexs}")

    """
        prefill
    """
    if input_prefill_len > 0:
        for slice_index in slice_indexs:
            if slice_index == 0:
                current_slice_len = input_prefill_len
            else:
                current_slice_len = kv_mask_expand_len

            indices = np.array(
                list(
                    range(
                        slice_index * input_prefill_len,
                        (slice_index + 1) * input_prefill_len,
                    )
                ),
                np.uint32,
            ).reshape((1, input_prefill_len))

            mask = (
                np.zeros((1, input_prefill_len, current_slice_len * slice_index + input_prefill_len))
                - 65536
            )
            data = np.zeros((1, input_prefill_len, cfg.hidden_size)).astype(bfloat16)
            for i, t in enumerate(
                range(
                    slice_index * input_prefill_len,
                    (slice_index + 1) * input_prefill_len,
                )
            ):
                if t < len(token_ids):
                    mask[:, i, : slice_index * input_prefill_len + i + 1] = 0
                    data[:, i : i + 1, :] = (
                        prefill_data[t]
                        .reshape((1, 1, cfg.hidden_size))
                        .astype(bfloat16)
                    )

            if slice_index == slice_indexs[-1]:
                curlen_procd = token_len - slice_index * input_prefill_len # curlen_procd 是当前处理数据的长度
            else:
                curlen_procd = input_prefill_len

            mask = mask.astype(bfloat16)
            for i in range(cfg.num_hidden_layers):
                input_feed = {
                    "K_cache": (
                        k_caches[i][:, 0: current_slice_len * slice_index, :]
                        if slice_index
                        else np.zeros((1, 1, cfg.hidden_size), dtype=bfloat16)
                    ),
                    "V_cache": (
                        v_caches[i][:, 0: current_slice_len * slice_index, :]
                        if slice_index
                        else np.zeros((1, 1, cfg.hidden_size), dtype=bfloat16)
                    ),
                    "indices": indices,
                    "input": data,
                    "mask": mask,
                }
                outputs = prefill_decoder_sessins[i].run(None, input_feed, shape_group=slice_index + 1)

                k_caches[i][
                    :,
                    slice_index
                    * input_prefill_len : slice_index
                    * input_prefill_len + curlen_procd, # current_slice_len
                    :,
                ] = outputs[0][:, :curlen_procd, :]

                v_caches[i][
                    :,
                    slice_index
                    * input_prefill_len : slice_index
                    * input_prefill_len + curlen_procd, # current_slice_len
                    :,
                ] = outputs[1][:, :curlen_procd, :]

                data = outputs[2]

            print("slice prefill done", slice_index)

        post_out = post_process_session.run(
            None,
            {
                "input": data[
                    :, token_len - (len(slice_indexs) - 1) * input_prefill_len - 1, None, :
                ]
            }
        )[0]
        next_token, posssible_tokens, possible_soft = post_process(post_out)
        posibles = [tokenizer.decode([t]) for t in posssible_tokens]
        posible_soft = [str((t, s)) for t, s in zip(posibles, possible_soft)]
        token_ids.append(next_token)

    print("answer >>", tokenizer.decode(token_ids[token_len], skip_special_tokens=True), end='', flush=True)
    # print("answer >>", end='', flush=True)

    # set to decoder
    kv_cache_len = lastN
    mask = np.zeros((1, 1, kv_cache_len + 1), dtype=np.float32).astype(bfloat16)
    mask[:, :, :kv_cache_len] -= 65536
    if input_prefill_len > 0:
        mask[:, :, :token_len] = 0

    # for start_indice in tqdm(range(kv_cache_len), desc="Decode"):
    for start_indice in range(kv_cache_len):
        if input_prefill_len > 0 and start_indice < token_len:
            continue

        next_token = token_ids[start_indice]
        indices = np.array([start_indice], np.uint32).reshape((1, 1))
        data = embeds[next_token, :].reshape((1, 1, cfg.hidden_size)).astype(bfloat16)
        for i in range(cfg.num_hidden_layers):
            input_feed = {
                "K_cache": k_caches[i],
                "V_cache": v_caches[i],
                "indices": indices,
                "input": data,
                "mask": mask,
            }
            outputs = prefill_decoder_sessins[i].run(None, input_feed, shape_group=0)
            k_caches[i][:, start_indice, :] = outputs[0][:, :, :]
            v_caches[i][:, start_indice, :] = outputs[1][:, :, :]
            data = outputs[2]
        mask[..., start_indice] = 0
        if start_indice < token_len - 1:
            pass
        else:
            post_out = post_process_session.run(None, {"input": data})[0]
            next_token, posssible_tokens, possible_soft = post_process(post_out)
            token_ids.append(next_token)
            if next_token == tokenizer.eos_token_id and next_token > token_len:
                break

        print(tokenizer.decode(next_token, skip_special_tokens=True), end='', flush=True)