qwen/modeling_qwen.py

@@ -198,7 +198,7 @@ class QwenRunner:
     ):
         qwen = self.qwen
         history = copy.deepcopy(history)
-        raw_text, context_tokens = self.prepareInput(tokenizer, query, query_assistant, history, system)
+        raw_text, context_tokens = make_context(tokenizer, query, query_assistant, history=history, system=system)
         input_ids = torch.tensor([context_tokens]).to(next(qwen.parameters()).device)
         eos_token_id_tensor = torch.tensor([qwen.config.eos_token_id]).to(input_ids.device)
         pad_token_id = qwen.config.pad_token_id
@@ -354,9 +354,6 @@ class QwenRunner:
 
         return lm_logits
 
-    def prepareInput(self, tokenizer, query, query_assistant, history, system):
-        return make_context(tokenizer, query, query_assistant, history=history, system=system)
-
     def repetition_penalty(self, input_ids, next_token_scores):
         penalty = self.qwen.config.repetition_penalty
         score = torch.gather(next_token_scores, 1, input_ids)

qwen/research_attention.py

@@ -7,7 +7,6 @@ from transformers import AutoConfig
 
 from modeling_qwen import QWenLMHeadModel
 from modeling_qwen import QwenRunner
-import numpy as np
 
 import torch.nn.functional as F
 
@@ -70,24 +69,18 @@ def Dump_lm_head_weight(model):
 
 # Dump_lm_head_weight(model)
 
-qk_sum = []
-qk_index = []
-
 
 def DumpQK(query, key, causal_mask, index):
     scale_factor = 1 / math.sqrt(query.size(-1))
     attn_weight = query @ key.transpose(-2, -1) * scale_factor
+    attn_weight = torch.softmax(attn_weight, dim=-1)
+    size = query.shape[2]
     attn_mask = torch.ones(causal_mask.shape, dtype=query.dtype, device=query.device)
     attn_mask.masked_fill_(causal_mask.logical_not(), float(0))
-    attn_weight = attn_weight * attn_mask
-    attn_weight = torch.softmax(attn_weight, dim=-1)
-    attn_weight = attn_weight * attn_mask
-    size = query.shape[2]
-    qk = attn_weight[0]
-    # prePath = "./temp/" + "q@k_seq_" + str(size) + "_layer_" + str(index) + ".png"
-    # show.DumpTensorToImage(qk, prePath, GridValue=255)
-    qk_sum.append(qk.sum(0))
-    qk_index.append(size)
+    qk = attn_weight * attn_mask
+    qk = qk[0]
+    prePath = "./temp/" + "q@k_seq_" + str(size) + "_layer_" + str(index) + ".png"
+    show.DumpTensorToImage(qk, prePath, GridValue=255)
 
 
 class ResearchRunner(QwenRunner):
@@ -106,48 +99,18 @@ class ResearchRunner(QwenRunner):
         attn_output = attention.c_proj(context_layer)
         return attn_output
 
-    def sample(self, next_token_scores):
-        qk_sum_cat = torch.stack(qk_sum, 0)
-        qk_sum.clear()
-        prePath = "./temp/" + "q@k_sum_seq_" + str(qk_index[-1]) + ".png"
-        show.DumpTensorToImage(qk_sum_cat, prePath, GridValue=255)
-
-        return super().sample(next_token_scores)
-
-    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\n你好\nassistant\n", allowed_special=set())
-
-        return "", tokens
-
-
 runner = ResearchRunner(model)
 
 # 第一轮对话
-# output_ids, history, decoded = runner.Chat(tokenizer, "东南亚国家日本的首都是什么市", "日本的首都是")
-# print(decoded)
-
-output_ids, history, decoded = runner.Chat(tokenizer, "你好!!", "")
+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)
+    de = str(i).zfill(3) + " : " + repr(de)
     tokens.append(de)
+    print(de)
 
 # <|im_start|>system
 # You are a helpful assistant.<|im_end|>
@@ -159,9 +122,5 @@ for i, token in enumerate(output_ids):
 
 show.DumpListToFile(tokens, "./temp/token_decode_list.txt")
 
-# if decoded.split("\n")[-2] != """日本的首都东京。<|im_end|>""":
-#     raise ()
-
-
-# normal   (x - mean) / (std + eps) =>  sum(y)==0
-# softmax  exp(x) / sum(exp(x))  =>  0 < y < 1   sum(y)==1
+if decoded.split("\n")[-2] != """日本的首都东京。<|im_end|>""":
+    raise ()