Witllm/qwen/research_silu.py

124 lines
3.5 KiB
Python
Raw Normal View History

2024-02-04 23:48:51 +08:00
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")