[ModuleFlow] support Qwen2.5

sophgo · Feb 10, 2025 · 71354f5 · 71354f5
1 parent 6674e40
commit 71354f5
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 6 deletions.
diff --git a/models/language_model/python_demo/pipeline.py b/models/language_model/python_demo/pipeline.py
@@ -35,7 +35,7 @@ def map(self, model_type, tokenizer_path):
         """Abstract model-specific mapper into a dictionary."""
         if model_type == "qwen2":
             self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
-            self.EOS = self.tokenizer.eos_token_id
+            self.EOS = [self.tokenizer.eos_token_id]
             self.append_user = lambda history, input_str: history.append(
                 {"role": "user", "content": input_str}
             )
@@ -46,9 +46,24 @@ def map(self, model_type, tokenizer_path):
                 history, tokenize=False, add_generation_prompt=True
             )
             self.system_prompt = {"role": "system", "content": "You are a helpful assistant."}
+        elif model_type == "qwen2.5":
+            self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
+            ID_IM_END = self.tokenizer.convert_tokens_to_ids("<|im_end|>")
+            ID_END = self.tokenizer.convert_tokens_to_ids("<|end|>")
+            self.EOS = [self.tokenizer.eos_token_id, ID_IM_END, ID_END]
+            self.append_user = lambda history, input_str: history.append(
+                {"role": "user", "content": input_str}
+            )
+            self.append_assistant = lambda history, answer_str: history.append(
+                {"role": "assistant", "content": answer_str}
+            )
+            self.apply_chat_template = lambda history: self.tokenizer.apply_chat_template(
+                history, tokenize=False, add_generation_prompt=True
+            )
+            self.system_prompt = {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."}
         elif model_type == "qwen":
             self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
-            self.EOS = self.tokenizer.im_end_id
+            self.EOS = [self.tokenizer.im_end_id]
             self.append_user = lambda history, input_str: history.append(
                 "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n".format(input_str)
             )
@@ -188,7 +203,6 @@ def stream_answer(self, tokens):
         token = self.model.forward_first(tokens)
         first_end = time.time()
         # Following tokens
-        total_tokens = [tokens]
         full_word_tokens = []
         while token not in self.EOS and self.model.total_length < self.model.SEQLEN:
             full_word_tokens.append(token)

diff --git a/models/module_flow.py b/models/module_flow.py
@@ -103,7 +103,7 @@ def init_from_args(self, args):
             self.dtype = torch.float
         else:
             os.environ['CUDA_VISIBLE_DEVICES'] = "0"
-            if self.model_type in ["qwen2"]:
+            if self.model_type in ["qwen2", "qwen2.5"]:
                 self.dtype = torch.bfloat16
             else:
                 raise ValueError(f"{self.model_type} not support now")
@@ -214,7 +214,9 @@ def stream_answer(self, tokens):
         print(word, end="")
 
         # decoding
-        while int(token) != self.EOS and token_len < self.seq_length:
+        ID_IM_END = self.tokenizer.convert_tokens_to_ids("<|im_end|>")
+        ID_END = self.tokenizer.convert_tokens_to_ids("<|end|>")
+        while int(token) not in self.EOS and token_len < self.seq_length:
             token_len += 1
             input_id = torch.tensor([token]).to(self.device)
             hidden_states = self.embed(input_id).view(1, 1, self.HIDDEN_SIZE)
@@ -237,7 +239,7 @@ def stream_answer(self, tokens):
 
 
 if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='llm_exporter', formatter_class=argparse.RawTextHelpFormatter)
+    parser = argparse.ArgumentParser(description='module_flow_test', formatter_class=argparse.RawTextHelpFormatter)
     parser.add_argument('-t', '--torch_path', type=str, required=True, help='torch path, like ./Qwen2-VL-2B-Instruct')
     parser.add_argument('--seq_length', type=int, required=True, help="sequence length")
     parser.add_argument('--visual_length', type=int, default=1024, help="visual length for vision transformer")