124 lines
3.5 KiB
Python
124 lines
3.5 KiB
Python
import torch
|
|
import sys
|
|
import math
|
|
from modelscope import snapshot_download
|
|
from transformers import AutoTokenizer
|
|
from transformers import AutoConfig
|
|
|
|
from modeling_qwen import QWenLMHeadModel
|
|
from modeling_qwen import QwenRunner
|
|
import numpy as np
|
|
|
|
import torch.nn.functional as F
|
|
|
|
from qwen_generation_utils import (
|
|
make_context,
|
|
decode_tokens,
|
|
)
|
|
|
|
sys.path.append("..")
|
|
from tools import show
|
|
|
|
seed = 4321
|
|
torch.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
|
|
model_dir = snapshot_download("qwen/Qwen-1_8B-Chat")
|
|
# model_dir = "/home/colin/.cache/modelscope/hub/qwen/Qwen-1_8B-Chat"
|
|
|
|
config, kwargs = AutoConfig.from_pretrained(
|
|
"./",
|
|
return_unused_kwargs=True,
|
|
trust_remote_code=True,
|
|
code_revision=None,
|
|
_commit_hash=None,
|
|
)
|
|
model = QWenLMHeadModel(config)
|
|
|
|
print(model)
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
|
|
|
|
model = model.from_pretrained(model_dir)
|
|
if torch.cuda.device_count() > 0:
|
|
model = model.cuda()
|
|
model = model.eval()
|
|
|
|
index = 0
|
|
|
|
|
|
class ResearchRunner(QwenRunner):
|
|
def __init__(self, model):
|
|
super().__init__(model)
|
|
|
|
def prepareInput(self, tokenizer, query, query_assistant, history, system):
|
|
start_to = [151644]
|
|
n_to = [198]
|
|
end_to = [151645]
|
|
system_str = "system\nYou are a helpful assistant."
|
|
user_str = "user\n" + query
|
|
aassistant_str = "assistant\n" + query_assistant
|
|
|
|
system_token = start_to + tokenizer.encode(system_str, allowed_special=set()) + end_to + n_to
|
|
user_token = start_to + tokenizer.encode(user_str, allowed_special=set()) + end_to + n_to
|
|
aassistant_token = start_to + tokenizer.encode(aassistant_str, allowed_special=set())
|
|
|
|
tokens = system_token + user_token + aassistant_token
|
|
tokens = user_token + aassistant_token
|
|
tokens = start_to + tokenizer.encode("user\nHi你好\nassistant\n我是", allowed_special=set())
|
|
|
|
return "", tokens
|
|
|
|
def forwardQWenBlock(
|
|
self,
|
|
block,
|
|
hidden_states,
|
|
rotary_pos_emb_list=None,
|
|
):
|
|
layernorm_output = block.ln_1(hidden_states)
|
|
|
|
attn_outputs = self.forwardAttention(block.attn, layernorm_output, rotary_pos_emb_list)
|
|
attn_output = attn_outputs[0]
|
|
layernorm_input = attn_output + hidden_states
|
|
|
|
layernorm_output = block.ln_2(layernorm_input)
|
|
a1 = block.mlp.w1(layernorm_output)
|
|
a2 = block.mlp.w2(layernorm_output)
|
|
|
|
activation = (F.relu(a2) > 0).to(float)
|
|
act_mean = torch.mean(activation, 2)
|
|
print("Layer:" + str(block.index))
|
|
print(act_mean.cpu())
|
|
global index
|
|
if index == 0:
|
|
activation = activation.reshape(activation.shape[1], 64, -1)
|
|
show.DumpTensorToImage(activation, "./temp/activation_layer_" + str(block.index) + ".png")
|
|
|
|
intermediate_parallel = a1 * F.silu(a2)
|
|
mlp_output = block.mlp.c_proj(intermediate_parallel)
|
|
|
|
hidden_states = layernorm_input + mlp_output
|
|
return hidden_states
|
|
|
|
def isFinish(self, next_tokens):
|
|
global index
|
|
index = index + 1
|
|
finish, next = super().isFinish(next_tokens)
|
|
return finish, next
|
|
|
|
|
|
para = list(model.parameters())
|
|
runner = ResearchRunner(model)
|
|
|
|
output_ids, history, decoded = runner.Chat(tokenizer, "你好!!", "")
|
|
print(decoded)
|
|
|
|
|
|
tokens = []
|
|
for i, token in enumerate(output_ids):
|
|
de = tokenizer.decode([token])
|
|
de = str(i + 1).zfill(3) + " : " + repr(de)
|
|
tokens.append(de)
|
|
|
|
show.DumpListToFile(tokens, "./temp/token_decode_list.txt")
|